zus_common/

codec.rs

use {
  bytes::{Buf, BufMut, Bytes, BytesMut},
  tokio_util::codec::{Decoder, Encoder},
};

use crate::{
  compression::Compressor,
  encryption::DesEncryptor,
  error::{Result, ZusError},
  protocol::{RpcMessage, RpcProtocolHeader},
};

/// RPC Codec for encoding/decoding messages (replacing Java's ZNetProtocolCodecFactory)
pub struct RpcCodec {
  max_frame_length: usize,
  pub compressor: Compressor,
  encryptor: Option<DesEncryptor>,
}

impl RpcCodec {
  pub fn new() -> Self {
    Self {
      max_frame_length: 10 * 1024 * 1024, // 10MB default
      compressor: Compressor::new(),      // 4KB threshold, QuickLZ (matching C++/Java)
      encryptor: None,
    }
  }

  pub fn with_max_frame_length(max_frame_length: usize) -> Self {
    Self {
      max_frame_length,
      compressor: Compressor::new(),
      encryptor: None,
    }
  }

  pub fn with_compressor(compressor: Compressor) -> Self {
    Self {
      max_frame_length: 10 * 1024 * 1024,
      compressor,
      encryptor: None,
    }
  }

  pub fn with_config(max_frame_length: usize, compressor: Compressor) -> Self {
    Self {
      max_frame_length,
      compressor,
      encryptor: None,
    }
  }

  /// Create a new RpcCodec with encryption enabled.
  ///
  /// # Arguments
  /// * `key` - 8-byte DES key for encryption/decryption
  pub fn with_encryption(key: &[u8]) -> Result<Self> {
    let encryptor = DesEncryptor::try_new(key)?;
    Ok(Self {
      max_frame_length: 10 * 1024 * 1024,
      compressor: Compressor::new(),
      encryptor: Some(encryptor),
    })
  }

  /// Create a fully configured RpcCodec with all options.
  ///
  /// # Arguments
  /// * `max_frame_length` - Maximum frame size in bytes
  /// * `compressor` - Compressor for compression/decompression
  /// * `encryption_key` - Optional 8-byte DES key for encryption
  pub fn with_full_config(
    max_frame_length: usize,
    compressor: Compressor,
    encryption_key: Option<&[u8]>,
  ) -> Result<Self> {
    let encryptor = match encryption_key {
      | Some(key) => Some(DesEncryptor::try_new(key)?),
      | None => None,
    };
    Ok(Self {
      max_frame_length,
      compressor,
      encryptor,
    })
  }

  /// Set the encryption key. Pass `None` to disable encryption.
  pub fn set_encryption_key(&mut self, key: Option<&[u8]>) -> Result<()> {
    self.encryptor = match key {
      | Some(k) => Some(DesEncryptor::try_new(k)?),
      | None => None,
    };
    Ok(())
  }

  /// Check if encryption is enabled.
  pub fn is_encryption_enabled(&self) -> bool {
    self.encryptor.is_some()
  }
}

impl Default for RpcCodec {
  fn default() -> Self {
    Self::new()
  }
}

impl Decoder for RpcCodec {
  type Error = ZusError;
  type Item = RpcMessage;

  fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>> {
    // Need at least header size
    if src.len() < RpcProtocolHeader::HEADER_SIZE {
      return Ok(None);
    }

    // Peek at the header to get body length
    let mut peek_buf = src.clone();
    let header = RpcProtocolHeader::decode(&mut peek_buf)?;

    // Check if we have the full message
    let total_length = RpcProtocolHeader::HEADER_SIZE + header.body_length as usize;
    if src.len() < total_length {
      // Reserve space for the full message
      src.reserve(total_length - src.len());
      return Ok(None);
    }

    // Check max frame length
    if total_length > self.max_frame_length {
      return Err(ZusError::Protocol(format!(
        "Frame too large: {} > {}",
        total_length, self.max_frame_length
      )));
    }

    // Now decode the actual header
    let header = RpcProtocolHeader::decode(src)?;

    // Read body (possibly encrypted and/or compressed)
    let raw_body = src.split_to(header.body_length as usize).freeze();

    // Verify CRC on raw body (before decryption/decompression)
    // This matches Java behavior which computes CRC on received data
    if !header.verify_datacrc(&raw_body) {
      return Err(ZusError::CrcMismatch);
    }

    // Step 1: Decrypt if encrypted flag is set
    // Order: wire data -> decrypt -> decompress -> plaintext
    let decrypted_body = if header.is_encrypted() {
      match &self.encryptor {
        | Some(enc) => {
          let decrypted = enc.decrypt(&raw_body)?;
          Bytes::from(decrypted)
        }
        | None => {
          return Err(ZusError::Encryption(
            "Received encrypted message but no encryption key configured".to_string(),
          ));
        }
      }
    } else {
      raw_body
    };

    // Step 2: Decode body based on message type
    // For requests/notifications: Parse method first, then decompress params if needed
    // For responses: Decompress entire body first
    let (method, body) = if header.msg_type == zus_proto::constants::MSG_TYPE_REQ
      || header.msg_type == zus_proto::constants::MSG_TYPE_NOTIFY
    {
      // Request/Notification wire format: [method_len(4)][method_bytes][params_len(4)][params_bytes]
      // If compressed, only params_bytes are compressed (matching Java behavior)

      let mut body_buf = decrypted_body;

      // Parse method (always uncompressed)
      if body_buf.len() < 4 {
        return Err(ZusError::Protocol("Invalid method name length".to_string()));
      }
      let method_len = body_buf.get_u32() as usize;
      if body_buf.len() < method_len {
        return Err(ZusError::Protocol("Invalid method name".to_string()));
      }
      let method = body_buf.split_to(method_len);

      // Parse params length
      if body_buf.len() < 4 {
        return Err(ZusError::Protocol("Invalid params length".to_string()));
      }
      let params_len = body_buf.get_u32() as usize;
      if body_buf.len() < params_len {
        return Err(ZusError::Protocol("Invalid params data".to_string()));
      }

      // Get params (possibly compressed)
      let raw_params = body_buf.split_to(params_len);

      // Decompress params if compressed flag is set
      let params = if header.is_compressed() {
        self.compressor.decompress(&raw_params)?
      } else {
        raw_params
      };

      (method, params)
    } else {
      // Response - handle encrypted length prefix, then decompress if needed
      let body_data = if header.is_encrypted() {
        // For encrypted responses, first 4 bytes are original length
        let mut body_buf = decrypted_body;
        if body_buf.len() < 4 {
          return Err(ZusError::Protocol(
            "Invalid encrypted response: missing length prefix".to_string(),
          ));
        }
        let original_len = body_buf.get_u32() as usize;
        if body_buf.len() < original_len {
          return Err(ZusError::Protocol(format!(
            "Invalid encrypted response: expected {} bytes, got {}",
            original_len,
            body_buf.len()
          )));
        }
        body_buf.split_to(original_len)
      } else {
        decrypted_body
      };

      // Decompress entire body if compressed
      let body = if header.is_compressed() {
        self.compressor.decompress(&body_data)?
      } else {
        body_data
      };
      (Bytes::new(), body)
    };

    Ok(Some(RpcMessage { header, method, body }))
  }
}

impl Encoder<RpcMessage> for RpcCodec {
  type Error = ZusError;

  fn encode(&mut self, mut item: RpcMessage, dst: &mut BytesMut) -> Result<()> {
    // Encode body based on message type:
    // For requests/notifications: [method_len(4)][method_bytes][params_len(4)][params_bytes]
    //   - Only params_bytes are compressed (matching Java behavior)
    //   - This allows Java decoder to parse method before decompression
    // For responses: just [params_bytes] (no length prefix)
    //   - Entire body can be compressed
    //
    // Processing order: plaintext -> compress -> encrypt -> wire

    let is_request = !item.method.is_empty();
    let mut full_body = BytesMut::new();
    let was_compressed: bool;

    if is_request {
      // Request or notification
      // Step 1: Try to compress ONLY the params data (not method)
      let (compressed_params, params_compressed) = self.compressor.compress(&item.body)?;
      was_compressed = params_compressed;

      // Step 2: Encode as [method_len][method][params_len][compressed_or_uncompressed_params]
      full_body.put_u32(item.method.len() as u32);
      full_body.put(item.method);
      full_body.put_u32(compressed_params.len() as u32); // Length of possibly compressed params
      full_body.put(compressed_params);
    } else {
      // Response - compress entire body (no method)
      let (compressed_body, body_compressed) = self.compressor.compress(&item.body)?;
      was_compressed = body_compressed;
      full_body.put(compressed_body);
    }

    // Step 3: Encrypt the entire body if encryption is enabled
    // For responses (no method), prepend original body length to handle zero-padding
    let final_body = if let Some(ref enc) = self.encryptor {
      let data_to_encrypt = if is_request {
        // For requests, the format already includes length prefixes
        full_body.freeze()
      } else {
        // For responses, prepend original length (4 bytes) to handle padding
        let mut with_len = BytesMut::with_capacity(4 + full_body.len());
        with_len.put_u32(full_body.len() as u32);
        with_len.put(full_body);
        with_len.freeze()
      };
      let encrypted = enc.encrypt(&data_to_encrypt)?;
      item.header.set_encrypted(true);
      Bytes::from(encrypted)
    } else {
      item.header.set_encrypted(false);
      full_body.freeze()
    };

    // Update header with final body length and compression flag
    item.header.body_length = final_body.len() as u32;
    item.header.set_compressed(was_compressed);

    // Calculate CRC on the actual data being sent (encrypted and/or compressed)
    // This matches Java behavior which computes CRC on the received body data
    item.header.datacrc = RpcProtocolHeader::calculate_datacrc(&final_body);

    // Reserve space
    let total_length = RpcProtocolHeader::HEADER_SIZE + final_body.len();
    dst.reserve(total_length);

    // Encode header
    item.header.encode(dst);

    // Encode final body (encrypted and/or compressed)
    dst.put(final_body);

    Ok(())
  }
}

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

  #[test]
  fn test_codec_roundtrip() {
    let mut codec = RpcCodec::new();

    let method = Bytes::from("test.method");
    let body = Bytes::from("hello");
    let msg = RpcMessage::new_request(1, method.clone(), body.clone());

    let mut buf = BytesMut::new();
    codec.encode(msg, &mut buf).unwrap();

    let decoded = codec.decode(&mut buf).unwrap().unwrap();
    assert_eq!(decoded.header.sequence, 1);
    assert_eq!(decoded.method, method);
    assert_eq!(decoded.body, body);
  }

  #[test]
  fn test_codec_with_encryption_roundtrip() {
    let key = b"12345678";
    let mut codec = RpcCodec::with_encryption(key).unwrap();

    let method = Bytes::from("test.method");
    let body = Bytes::from("hello encrypted world");
    let msg = RpcMessage::new_request(1, method.clone(), body.clone());

    let mut buf = BytesMut::new();
    codec.encode(msg, &mut buf).unwrap();

    // Verify encryption flag is set in header
    let mut peek_buf = buf.clone();
    let header = RpcProtocolHeader::decode(&mut peek_buf).unwrap();
    assert!(header.is_encrypted());

    let decoded = codec.decode(&mut buf).unwrap().unwrap();
    assert_eq!(decoded.header.sequence, 1);
    assert_eq!(decoded.method, method);
    assert_eq!(decoded.body, body);
  }

  #[test]
  fn test_codec_encrypted_response() {
    let key = b"testkey!";
    let mut codec = RpcCodec::with_encryption(key).unwrap();

    // Test with any body size - padding is handled by prepending original length
    let body = Bytes::from("response data");
    let msg = RpcMessage::new_response(42, body.clone());

    let mut buf = BytesMut::new();
    codec.encode(msg, &mut buf).unwrap();

    let decoded = codec.decode(&mut buf).unwrap().unwrap();
    assert_eq!(decoded.header.sequence, 42);
    assert!(decoded.method.is_empty());
    assert_eq!(decoded.body, body);
  }

  #[test]
  fn test_codec_encryption_key_mismatch() {
    let key1 = b"12345678";
    let key2 = b"87654321";

    let mut encoder = RpcCodec::with_encryption(key1).unwrap();
    let mut decoder = RpcCodec::with_encryption(key2).unwrap();

    let method = Bytes::from("test.method");
    let body = Bytes::from("hello");
    let msg = RpcMessage::new_request(1, method, body);

    let mut buf = BytesMut::new();
    encoder.encode(msg, &mut buf).unwrap();

    // Decoding with wrong key should fail (likely produce garbage or error)
    let result = decoder.decode(&mut buf);
    // Due to block cipher properties, decryption with wrong key produces garbage
    // which likely causes protocol parsing errors
    assert!(result.is_err() || result.unwrap().is_none());
  }

  #[test]
  fn test_codec_encrypted_without_key_configured() {
    // Encode with encryption
    let key = b"12345678";
    let mut encoder = RpcCodec::with_encryption(key).unwrap();

    let method = Bytes::from("test.method");
    let body = Bytes::from("hello");
    let msg = RpcMessage::new_request(1, method, body);

    let mut buf = BytesMut::new();
    encoder.encode(msg, &mut buf).unwrap();

    // Try to decode without encryption key configured
    let mut decoder = RpcCodec::new();
    let result = decoder.decode(&mut buf);
    assert!(result.is_err());

    if let Err(ZusError::Encryption(msg)) = result {
      assert!(msg.contains("no encryption key configured"));
    } else {
      panic!("Expected Encryption error");
    }
  }

  #[test]
  fn test_codec_set_encryption_key() {
    let mut codec = RpcCodec::new();
    assert!(!codec.is_encryption_enabled());

    codec.set_encryption_key(Some(b"12345678")).unwrap();
    assert!(codec.is_encryption_enabled());

    codec.set_encryption_key(None).unwrap();
    assert!(!codec.is_encryption_enabled());
  }

  #[test]
  fn test_codec_full_config() {
    let compressor = Compressor::new();
    let key = b"mykey123";

    let codec = RpcCodec::with_full_config(5 * 1024 * 1024, compressor, Some(key)).unwrap();
    assert!(codec.is_encryption_enabled());

    let codec2 = RpcCodec::with_full_config(5 * 1024 * 1024, Compressor::new(), None).unwrap();
    assert!(!codec2.is_encryption_enabled());
  }
}