knafeh 1.1.0

use prost::Message;

use crate::codec::Codec;
use crate::error::KnafehError;

/// Protobuf codec — the default codec for Knafeh.
///
/// Wraps payloads in a simple `Payload` envelope (protobuf message with
/// a single `bytes data = 1` field) to ensure the wire format is always
/// a valid protobuf message. Validates the envelope on decode and rejects
/// mismatched message types.
#[derive(Debug, Clone, Default)]
pub struct ProtobufCodec;

impl ProtobufCodec {
    pub fn new() -> Self {
        Self
    }
}

impl Codec for ProtobufCodec {
    fn encode(&self, value: &[u8]) -> Result<Vec<u8>, KnafehError> {
        let payload = Payload {
            data: value.to_vec(),
        };
        Ok(payload.encode_to_vec())
    }

    fn decode(&self, data: &[u8]) -> Result<Vec<u8>, KnafehError> {
        // Decode directly from &[u8] — avoids an extra Bytes allocation.
        let payload = Payload::decode(data)
            .map_err(|e| KnafehError::Codec(format!("protobuf decode error: {e}")))?;

        // Reject mismatched messages: if the wire data is non-empty but
        // decoded to an empty payload, it likely has unknown fields from
        // a different message type that Payload silently ignores.
        if payload.data.is_empty() && !data.is_empty() {
            return Err(KnafehError::Codec(
                "protobuf decode error: mismatched or unknown payload envelope".to_string(),
            ));
        }

        Ok(payload.data)
    }

    fn content_type(&self) -> &str {
        "application/grpc+proto"
    }

    fn name(&self) -> &str {
        "protobuf"
    }
}

/// Raw protobuf codec — passes bytes through without an envelope.
///
/// Use this when you manage protobuf serialization yourself and don't
/// need the framework to wrap/unwrap payloads. The codec performs no
/// validation or interpretation of the bytes — any payload is accepted.
///
/// Note: the `Codec` trait returns owned `Vec<u8>`, so a copy still
/// occurs on encode/decode. "Raw" means no envelope, not zero-copy.
#[derive(Debug, Clone, Default)]
pub struct RawProtobufCodec;

impl RawProtobufCodec {
    pub fn new() -> Self {
        Self
    }
}

impl Codec for RawProtobufCodec {
    fn encode(&self, value: &[u8]) -> Result<Vec<u8>, KnafehError> {
        Ok(value.to_vec())
    }

    fn decode(&self, data: &[u8]) -> Result<Vec<u8>, KnafehError> {
        Ok(data.to_vec())
    }

    fn content_type(&self) -> &str {
        "application/grpc+proto"
    }

    fn name(&self) -> &str {
        "raw-protobuf"
    }
}

/// Simple protobuf envelope for wrapping arbitrary bytes.
///
/// ```protobuf
/// message Payload {
///   bytes data = 1;
/// }
/// ```
#[derive(Clone, PartialEq, Message)]
struct Payload {
    #[prost(bytes = "vec", tag = "1")]
    data: Vec<u8>,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_protobuf_codec_roundtrip() {
        let codec = ProtobufCodec::new();
        let input = b"hello protobuf";
        let encoded = codec.encode(input).unwrap();
        assert_ne!(encoded, input);
        let decoded = codec.decode(&encoded).unwrap();
        assert_eq!(decoded, input);
    }

    #[test]
    fn test_protobuf_codec_empty() {
        let codec = ProtobufCodec::new();
        let encoded = codec.encode(b"").unwrap();
        let decoded = codec.decode(&encoded).unwrap();
        assert_eq!(decoded, b"");
    }

    #[test]
    fn test_protobuf_codec_large() {
        let codec = ProtobufCodec::new();
        let input = vec![0xAB; 100_000];
        let encoded = codec.encode(&input).unwrap();
        let decoded = codec.decode(&encoded).unwrap();
        assert_eq!(decoded, input);
    }

    #[test]
    fn test_protobuf_codec_invalid_decode() {
        let codec = ProtobufCodec::new();
        // Invalid protobuf (truncated varint).
        let result = codec.decode(&[0x80, 0x80]);
        assert!(result.is_err());
    }

    #[test]
    fn test_protobuf_codec_mismatched_message() {
        let codec = ProtobufCodec::new();
        // A valid protobuf message with tag 2 (not tag 1) — Payload
        // ignores unknown fields and decodes to empty data.
        let result = codec.decode(&[0x12, 0x03, 0x66, 0x6f, 0x6f]);
        assert!(result.is_err());
    }

    #[test]
    fn test_raw_protobuf_codec_passthrough() {
        let codec = RawProtobufCodec::new();
        let input = b"raw bytes";
        let encoded = codec.encode(input).unwrap();
        assert_eq!(encoded, input);
        let decoded = codec.decode(&encoded).unwrap();
        assert_eq!(decoded, input);
    }
}