kinetic_encoder_ocsf/

lib.rs

#![deny(clippy::unwrap_used)]

pub mod generated {
    include!(concat!(env!("OUT_DIR"), "/generated_schema.rs"));
}

use arrow_array::builder::*;
use arrow_schema::{DataType, Schema};
use bytes::Bytes;
use kinetic_core::encode::{
    Decoder, DecoderConfig, Encoder, EncoderConfig, Error as EncodeError, Result as EncodeResult,
};
use kinetic_core::event::EventBatch;
use kinetic_core::metadata::ArcEventMetadata;
use kinetic_doc_derive::ComponentDoc;
use serde::{Deserialize, Serialize};
use snafu::Snafu;
use sonic_rs::{JsonContainerTrait, JsonValueMutTrait, JsonValueTrait};
use std::sync::Arc;

#[derive(Debug, Snafu)]
pub enum Error {
    #[snafu(display("Failed to parse OCSF JSON: {}", source))]
    SonicParse { source: sonic_rs::Error },

    #[snafu(display("Arrow IPC encoding error: {}", source))]
    ArrowEncode { source: arrow_schema::ArrowError },

    #[snafu(display("Kinetic core error: {}", source))]
    KineticCore { source: kinetic_core::Error },

    #[snafu(display("Unsupported OCSF version: {}", version))]
    UnsupportedVersion { version: String },

    #[snafu(display("Schema mismatch: {}", message))]
    SchemaMismatch { message: String },
}

impl From<Error> for EncodeError {
    fn from(err: Error) -> Self {
        match err {
            Error::SonicParse { source } => EncodeError::Decode {
                source: Box::new(source),
            },
            Error::ArrowEncode { source } => EncodeError::Encode {
                source: Box::new(source),
            },
            Error::KineticCore { source } => EncodeError::Decode {
                source: Box::new(source),
            },
            Error::UnsupportedVersion { version } => EncodeError::Decode {
                source: format!("Unsupported OCSF version: {}", version).into(),
            },
            Error::SchemaMismatch { message } => EncodeError::Decode {
                source: message.into(),
            },
        }
    }
}

pub type Result<T, E = Error> = std::result::Result<T, E>;

/// Configuration for the OCSF Decoder.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, ComponentDoc)]
#[component(type = "decoder", name = "ocsf")]
pub struct OcsfDecoderConfig {
    /// OCSF schema version to use (e.g., "1.0").
    pub version: Option<String>,
    /// Maximum size of a single message in bytes.
    pub max_size: Option<usize>,
}

impl Default for OcsfDecoderConfig {
    fn default() -> Self {
        Self {
            version: None,
            max_size: Some(10 * 1024 * 1024), // 10MB default
        }
    }
}

impl DecoderConfig for OcsfDecoderConfig {
    fn build(&self, _schema: Arc<Schema>) -> EncodeResult<Arc<dyn Decoder>> {
        // Note: For OCSF we use the internally generated schemas.
        // We can use the version provided in config or detect it per-payload.
        Ok(Arc::new(OcsfDecoder {
            config: self.clone(),
        }))
    }
}

/// The OCSF Decoder.
pub struct OcsfDecoder {
    pub config: OcsfDecoderConfig,
}

impl Decoder for OcsfDecoder {
    fn decode(&self, data: &[u8], metadata: ArcEventMetadata) -> EncodeResult<EventBatch> {
        self.decode(data, metadata)
    }
}

fn append_value(
    builder: &mut dyn arrow_array::builder::ArrayBuilder,
    data_type: &DataType,
    value: Option<&sonic_rs::Value>,
) -> Result<()> {
    match data_type {
        DataType::Utf8 => {
            let b = builder
                .as_any_mut()
                .downcast_mut::<StringBuilder>()
                .ok_or_else(|| Error::SchemaMismatch {
                    message: "Expected StringBuilder".to_string(),
                })?;
            match value {
                Some(v) => {
                    if let Some(s) = v.as_str() {
                        b.append_value(s);
                    } else if !v.is_null() {
                        if let Ok(s) = sonic_rs::to_string(v) {
                            b.append_value(s);
                        } else {
                            b.append_null();
                        }
                    } else {
                        b.append_null();
                    }
                }
                None => b.append_null(),
            }
        }
        DataType::Int64 => {
            let b = builder
                .as_any_mut()
                .downcast_mut::<Int64Builder>()
                .ok_or_else(|| Error::SchemaMismatch {
                    message: "Expected Int64Builder".to_string(),
                })?;
            match value {
                Some(v) => {
                    if let Some(i) = v.as_i64() {
                        b.append_value(i);
                    } else if let Some(u) = v.as_u64() {
                        b.append_value(u as i64);
                    } else {
                        b.append_null();
                    }
                }
                None => b.append_null(),
            }
        }
        DataType::Float64 => {
            let b = builder
                .as_any_mut()
                .downcast_mut::<Float64Builder>()
                .ok_or_else(|| Error::SchemaMismatch {
                    message: "Expected Float64Builder".to_string(),
                })?;
            match value {
                Some(v) => {
                    if let Some(f) = v.as_f64() {
                        b.append_value(f);
                    } else if let Some(i) = v.as_i64() {
                        b.append_value(i as f64);
                    } else if let Some(u) = v.as_u64() {
                        b.append_value(u as f64);
                    } else {
                        b.append_null();
                    }
                }
                None => b.append_null(),
            }
        }
        DataType::Boolean => {
            let b = builder
                .as_any_mut()
                .downcast_mut::<BooleanBuilder>()
                .ok_or_else(|| Error::SchemaMismatch {
                    message: "Expected BooleanBuilder".to_string(),
                })?;
            match value {
                Some(v) => {
                    if let Some(b_val) = v.as_bool() {
                        b.append_value(b_val);
                    } else {
                        b.append_null();
                    }
                }
                None => b.append_null(),
            }
        }
        DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, _) => {
            let b = builder
                .as_any_mut()
                .downcast_mut::<TimestampMillisecondBuilder>()
                .ok_or_else(|| Error::SchemaMismatch {
                    message: "Expected TimestampMillisecondBuilder".to_string(),
                })?;
            match value {
                Some(v) => {
                    if let Some(i) = v.as_i64() {
                        b.append_value(i);
                    } else if let Some(u) = v.as_u64() {
                        b.append_value(u as i64);
                    } else {
                        b.append_null();
                    }
                }
                None => b.append_null(),
            }
        }
        DataType::List(field) => {
            if let Some(b) = builder
                .as_any_mut()
                .downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
            {
                if let Some(v) = value {
                    if let Some(arr) = v.as_array() {
                        for item in arr {
                            append_value(b.values(), field.data_type(), Some(item))?;
                        }
                        b.append(true);
                    } else {
                        b.append(false);
                    }
                } else {
                    b.append(false);
                }
            }
        }
        DataType::Struct(fields) => {
            if let Some(b) = builder.as_any_mut().downcast_mut::<StructBuilder>() {
                if let Some(v) = value.and_then(|v| v.as_object()) {
                    for (i, field) in fields.iter().enumerate() {
                        let field_val = v.get(field.name());
                        let dt = field.data_type();

                        match dt {
                            DataType::Utf8 => append_value(
                                b.field_builder::<StringBuilder>(i).ok_or_else(|| {
                                    Error::SchemaMismatch {
                                        message: format!("Field {} builder missing", field.name()),
                                    }
                                })?,
                                dt,
                                field_val,
                            )?,
                            DataType::Int64 => append_value(
                                b.field_builder::<Int64Builder>(i).ok_or_else(|| {
                                    Error::SchemaMismatch {
                                        message: format!("Field {} builder missing", field.name()),
                                    }
                                })?,
                                dt,
                                field_val,
                            )?,
                            DataType::Float64 => append_value(
                                b.field_builder::<Float64Builder>(i).ok_or_else(|| {
                                    Error::SchemaMismatch {
                                        message: format!("Field {} builder missing", field.name()),
                                    }
                                })?,
                                dt,
                                field_val,
                            )?,
                            DataType::Boolean => append_value(
                                b.field_builder::<BooleanBuilder>(i).ok_or_else(|| {
                                    Error::SchemaMismatch {
                                        message: format!("Field {} builder missing", field.name()),
                                    }
                                })?,
                                dt,
                                field_val,
                            )?,
                            DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, _) => {
                                append_value(
                                    b.field_builder::<TimestampMillisecondBuilder>(i)
                                        .ok_or_else(|| Error::SchemaMismatch {
                                            message: format!(
                                                "Field {} builder missing",
                                                field.name()
                                            ),
                                        })?,
                                    dt,
                                    field_val,
                                )?
                            }
                            DataType::List(_) => append_value(
                                b.field_builder::<ListBuilder<Box<dyn ArrayBuilder>>>(i)
                                    .ok_or_else(|| Error::SchemaMismatch {
                                        message: format!("Field {} builder missing", field.name()),
                                    })?,
                                dt,
                                field_val,
                            )?,
                            DataType::Struct(inner_fields) => append_value(
                                b.field_builder::<StructBuilder>(i).ok_or_else(|| {
                                    Error::SchemaMismatch {
                                        message: format!("Field {} builder missing", field.name()),
                                    }
                                })?,
                                &DataType::Struct(inner_fields.clone()),
                                field_val,
                            )?,
                            _ => {}
                        }
                    }
                    b.append(true);
                } else {
                    // Recursive null append for all fields
                    for (i, field) in fields.iter().enumerate() {
                        let dt = field.data_type();
                        match dt {
                            DataType::Utf8 => append_value(
                                b.field_builder::<StringBuilder>(i).ok_or_else(|| {
                                    Error::SchemaMismatch {
                                        message: format!("Field {} builder missing", field.name()),
                                    }
                                })?,
                                dt,
                                None,
                            )?,
                            DataType::Int64 => append_value(
                                b.field_builder::<Int64Builder>(i).ok_or_else(|| {
                                    Error::SchemaMismatch {
                                        message: format!("Field {} builder missing", field.name()),
                                    }
                                })?,
                                dt,
                                None,
                            )?,
                            DataType::Float64 => append_value(
                                b.field_builder::<Float64Builder>(i).ok_or_else(|| {
                                    Error::SchemaMismatch {
                                        message: format!("Field {} builder missing", field.name()),
                                    }
                                })?,
                                dt,
                                None,
                            )?,
                            DataType::Boolean => append_value(
                                b.field_builder::<BooleanBuilder>(i).ok_or_else(|| {
                                    Error::SchemaMismatch {
                                        message: format!("Field {} builder missing", field.name()),
                                    }
                                })?,
                                dt,
                                None,
                            )?,
                            DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, _) => {
                                append_value(
                                    b.field_builder::<TimestampMillisecondBuilder>(i)
                                        .ok_or_else(|| Error::SchemaMismatch {
                                            message: format!(
                                                "Field {} builder missing",
                                                field.name()
                                            ),
                                        })?,
                                    dt,
                                    None,
                                )?
                            }
                            DataType::List(_) => append_value(
                                b.field_builder::<ListBuilder<Box<dyn ArrayBuilder>>>(i)
                                    .ok_or_else(|| Error::SchemaMismatch {
                                        message: format!("Field {} builder missing", field.name()),
                                    })?,
                                dt,
                                None,
                            )?,
                            DataType::Struct(inner_fields) => append_value(
                                b.field_builder::<StructBuilder>(i).ok_or_else(|| {
                                    Error::SchemaMismatch {
                                        message: format!("Field {} builder missing", field.name()),
                                    }
                                })?,
                                &DataType::Struct(inner_fields.clone()),
                                None,
                            )?,
                            _ => {}
                        }
                    }
                    b.append(false);
                }
            }
        }
        _ => {
            // Ignored types
        }
    }
    Ok(())
}

impl OcsfDecoder {
    fn decode(&self, data: &[u8], metadata: ArcEventMetadata) -> EncodeResult<EventBatch> {
        if let Some(limit) = self.config.max_size
            && data.len() > limit
        {
            return Err(EncodeError::MessageTooLarge {
                size: data.len(),
                limit,
            });
        }

        let document = sonic_rs::from_slice::<sonic_rs::Value>(data)
            .map_err(|e| Error::SonicParse { source: e })?;

        // Determine which schema to use
        let schema = if let Some(config_version) = &self.config.version {
            generated::get_schema_by_version(config_version).ok_or_else(|| {
                Error::UnsupportedVersion {
                    version: config_version.clone(),
                }
            })?
        } else {
            // Inference: try to get version from metadata.version in payload
            let inferred_version = document
                .get("metadata")
                .and_then(|m| m.get("version"))
                .and_then(|v| v.as_str())
                .unwrap_or("v1.8.0"); // Default to latest if not found

            generated::get_schema_by_version(inferred_version)
                .unwrap_or_else(generated::get_latest_schema)
        };

        let mut builders: Vec<Box<dyn ArrayBuilder>> = schema
            .fields()
            .iter()
            .map(|f| arrow_array::builder::make_builder(f.data_type(), 1))
            .collect();

        let mut unmapped = sonic_rs::json!({});

        if let Some(obj) = document.as_object() {
            for (i, field) in schema.fields().iter().enumerate() {
                let name = field.name();
                if name == "_unmapped" {
                    continue;
                }

                let value = obj.get(name);
                append_value(&mut builders[i], field.data_type(), value)?;
            }

            if let Some(unmapped_obj) = unmapped.as_object_mut() {
                for (k, v) in obj.iter() {
                    if schema.column_with_name(k).is_none() {
                        unmapped_obj.insert(k, v.clone());
                    }
                }
            }
        } else {
            for (i, field) in schema.fields().iter().enumerate() {
                if field.name() != "_unmapped" {
                    append_value(&mut builders[i], field.data_type(), None)?;
                }
            }
        }

        if let Ok(unmapped_idx) = schema.index_of("_unmapped") {
            let unmapped_str = sonic_rs::to_string(&unmapped).unwrap_or_else(|_| "{}".to_string());
            let b = builders[unmapped_idx]
                .as_any_mut()
                .downcast_mut::<StringBuilder>()
                .ok_or_else(|| Error::SchemaMismatch {
                    message: "Expected StringBuilder for _unmapped field".to_string(),
                })?;
            b.append_value(unmapped_str);
        }

        let arrays: Vec<std::sync::Arc<dyn arrow_array::Array>> =
            builders.into_iter().map(|mut b| b.finish()).collect();

        let record_batch = arrow_array::RecordBatch::try_new(schema, arrays)
            .map_err(|e| Error::ArrowEncode { source: e })?;
        let event_batch = EventBatch::new_with_xid(record_batch, metadata)
            .map_err(|e| Error::KineticCore { source: e })?;

        Ok(event_batch)
    }
}

/// Configuration for the OCSF Encoder.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default, ComponentDoc)]
#[component(type = "encoder", name = "ocsf")]
pub struct OcsfEncoderConfig {}

impl EncoderConfig for OcsfEncoderConfig {
    fn build(&self) -> EncodeResult<Arc<dyn Encoder>> {
        Ok(Arc::new(OcsfEncoder::default()))
    }
}

/// The OCSF Encoder.
#[derive(Default)]
pub struct OcsfEncoder {}

impl Encoder for OcsfEncoder {
    fn encode(&self, batch: &EventBatch) -> EncodeResult<Bytes> {
        self.encode_inner(batch).map_err(Into::into)
    }
}

impl OcsfEncoder {
    fn encode_inner(&self, batch: &EventBatch) -> Result<Bytes> {
        // Output Arrow IPC stream for performance.
        let mut buf = Vec::new();
        {
            let mut writer =
                arrow_ipc::writer::StreamWriter::try_new(&mut buf, batch.payload.schema().as_ref())
                    .map_err(|e| Error::ArrowEncode { source: e })?;

            writer
                .write(&batch.payload)
                .map_err(|e| Error::ArrowEncode { source: e })?;

            writer
                .finish()
                .map_err(|e| Error::ArrowEncode { source: e })?;
        }

        Ok(Bytes::from(buf))
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use kinetic_core::metadata::{ComponentId, EventMetadata};
    use std::sync::Arc;

    #[test]
    fn test_schema_generation_v1_8_0() {
        let schema = generated::get_schema_v1_8_0();
        assert!(schema.column_with_name("message").is_some());
        assert!(schema.column_with_name("_unmapped").is_some());
    }

    #[test]
    fn test_schema_lookup() {
        assert!(generated::get_schema_by_version("v1.3.0").is_some());
        assert!(generated::get_schema_by_version("1.8.0").is_some());
        assert!(generated::get_schema_by_version("v1.8.0").is_some());
        assert!(generated::get_schema_by_version("nonexistent").is_none());
    }

    #[test]
    fn test_ocsf_decoder_inference() {
        let config = OcsfDecoderConfig {
            version: None,
            max_size: None,
        };
        let decoder = config.build(generated::get_latest_schema()).unwrap();

        let metadata = Arc::new(EventMetadata::new(
            "test_pipeline",
            ComponentId::from("test_source"),
        ));
        let result = decoder.decode(
            b"{\"message\": \"test event\", \"metadata\": {\"version\": \"1.3.0\"}, \"unknown_field\": 123}",
            metadata,
        );
        assert!(result.is_ok());
        let batch = result.unwrap();

        let msg_col = batch
            .payload
            .column_by_name("message")
            .ok_or_else(|| {
                format!(
                    "Column 'message' not found in batch. Available columns: {:?}",
                    batch
                        .payload
                        .schema()
                        .fields()
                        .iter()
                        .map(|f| f.name())
                        .collect::<Vec<_>>()
                )
            })
            .unwrap();
        let msg_arr = msg_col
            .as_any()
            .downcast_ref::<arrow_array::StringArray>()
            .unwrap();
        assert_eq!(msg_arr.value(0), "test event");

        let unmapped_col = batch.payload.column_by_name("_unmapped").unwrap();
        let unmapped_arr = unmapped_col
            .as_any()
            .downcast_ref::<arrow_array::StringArray>()
            .unwrap();
        assert!(unmapped_arr.value(0).contains("\"unknown_field\":123"));
    }

    #[test]
    fn test_ocsf_encoder() {
        let config = OcsfEncoderConfig {};
        let encoder = config.build().unwrap();

        let schema = generated::get_latest_schema();
        let empty_record_batch = arrow_array::RecordBatch::new_empty(schema);
        let metadata = Arc::new(EventMetadata::new(
            "test_pipeline",
            ComponentId::from("test_source"),
        ));
        let event_batch = EventBatch::new(empty_record_batch, metadata).unwrap();

        let result = encoder.encode(&event_batch);
        assert!(result.is_ok());
        let bytes = result.unwrap();
        assert!(!bytes.is_empty());
    }
}