use crate::common::compression::{CompressionAlgorithm, CompressionUtil};
use crate::common::encryption::{EncryptionAlgorithm, EncryptionUtil};
use crate::common::error::Result;
use crate::common::protocol::{Frame, SerializationFormat};
use crate::common::serializer::SerializationUtil;
use lazy_static::lazy_static;
pub const MIN_COMPRESSION_PAYLOAD_BYTES: usize = 512;
lazy_static! {
pub static ref PRE_NEGOTIATION_PARSER: MessageParser = MessageParser::new(
SerializationFormat::Json,
CompressionAlgorithm::None,
EncryptionAlgorithm::None,
);
}
#[derive(Debug, Clone)]
pub struct MessageParser {
default_format: SerializationFormat,
default_compression: CompressionAlgorithm,
default_encryption: EncryptionAlgorithm,
custom_format_name: Option<String>,
}
impl MessageParser {
pub fn new(
format: SerializationFormat,
compression: CompressionAlgorithm,
encryption: EncryptionAlgorithm,
) -> Self {
Self {
default_format: format,
default_compression: compression,
default_encryption: encryption,
custom_format_name: None,
}
}
pub fn with_custom_format(
format_name: &str,
compression: CompressionAlgorithm,
encryption: EncryptionAlgorithm,
) -> Self {
Self {
default_format: SerializationFormat::Json, default_compression: compression,
default_encryption: encryption,
custom_format_name: Some(format_name.to_string()),
}
}
pub fn new_with_format_compression(
format: SerializationFormat,
compression: CompressionAlgorithm,
) -> Self {
Self::new(format, compression, EncryptionAlgorithm::None)
}
pub fn protobuf() -> Self {
Self::new(
SerializationFormat::Protobuf,
CompressionAlgorithm::None,
EncryptionAlgorithm::None,
)
}
pub fn json() -> Self {
Self::new(
SerializationFormat::Json,
CompressionAlgorithm::None,
EncryptionAlgorithm::None,
)
}
pub fn default_format(&self) -> SerializationFormat {
self.default_format
}
pub fn default_compression(&self) -> CompressionAlgorithm {
self.default_compression.clone()
}
pub fn default_encryption(&self) -> EncryptionAlgorithm {
self.default_encryption.clone()
}
pub fn parse(&self, data: &[u8]) -> Result<Frame> {
self.parse_with_fallback(data, true)
}
pub fn parse_with_fallback(&self, data: &[u8], allow_fallback: bool) -> Result<Frame> {
let decrypted = self.decrypt_data_with_fallback(data, allow_fallback)?;
let decompressed = self.decompress_data_with_fallback(&decrypted, allow_fallback)?;
self.parse_decompressed_with_fallback(&decompressed, allow_fallback)
}
pub fn parse_with_format(&self, data: &[u8], format: SerializationFormat) -> Result<Frame> {
let decrypted = self.decrypt_data(data)?;
let decompressed = self.decompress_data(&decrypted)?;
let serializer = if let Some(custom_name) = &self.custom_format_name {
SerializationUtil::get_serializer_by_name(custom_name)
} else {
SerializationUtil::get_serializer(format)
}
.ok_or_else(|| {
let format_info = if let Some(name) = &self.custom_format_name {
format!("custom format '{}'", name)
} else {
format!("{:?}", format)
};
crate::common::error::FlareError::deserialization_error(format!(
"Serializer not found: {}",
format_info
))
})?;
serializer.deserialize(&decompressed)
}
pub fn serialize(&self, frame: &Frame) -> Result<Vec<u8>> {
self.serialize_with_format(
frame,
self.default_format,
self.default_compression.clone(),
self.default_encryption.clone(),
)
}
pub fn serialize_with_format(
&self,
frame: &Frame,
format: SerializationFormat,
compression: CompressionAlgorithm,
encryption: EncryptionAlgorithm,
) -> Result<Vec<u8>> {
let serializer = if let Some(custom_name) = &self.custom_format_name {
SerializationUtil::get_serializer_by_name(custom_name)
} else {
SerializationUtil::get_serializer(format)
}
.ok_or_else(|| {
let format_info = if let Some(name) = &self.custom_format_name {
format!("custom format '{}'", name)
} else {
format!("{:?}", format)
};
crate::common::error::FlareError::encoding_error(format!(
"Serializer not found: {}",
format_info
))
})?;
let data = serializer.serialize(frame)?;
let compressed = if Self::should_compress_payload(data.len(), &compression) {
CompressionUtil::compress(&data, compression)?
} else {
data
};
self.encrypt_data(&compressed, encryption)
}
pub fn serialize_with_format_compression(
&self,
frame: &Frame,
format: SerializationFormat,
compression: CompressionAlgorithm,
) -> Result<Vec<u8>> {
self.serialize_with_format(frame, format, compression, self.default_encryption.clone())
}
pub fn get_compression_from_frame(frame: &Frame) -> CompressionAlgorithm {
frame
.metadata
.get("compression")
.and_then(|bytes| std::str::from_utf8(bytes).ok())
.and_then(CompressionAlgorithm::from_str)
.unwrap_or(CompressionAlgorithm::None)
}
pub fn should_compress_payload(payload_len: usize, compression: &CompressionAlgorithm) -> bool {
let _ = payload_len;
*compression != CompressionAlgorithm::None
}
pub fn get_format_from_frame(frame: &Frame) -> Option<SerializationFormat> {
frame
.metadata
.get("format")
.and_then(|bytes| std::str::from_utf8(bytes).ok())
.and_then(|s| {
if s.eq_ignore_ascii_case("protobuf") {
Some(SerializationFormat::Protobuf)
} else if s.eq_ignore_ascii_case("json") {
Some(SerializationFormat::Json)
} else {
None
}
})
}
pub fn get_encryption_from_frame(frame: &Frame) -> EncryptionAlgorithm {
frame
.metadata
.get("encryption")
.and_then(|bytes| std::str::from_utf8(bytes).ok())
.and_then(EncryptionAlgorithm::from_str)
.unwrap_or(EncryptionAlgorithm::None)
}
fn decrypt_data_with_fallback(&self, data: &[u8], allow_fallback: bool) -> Result<Vec<u8>> {
if self.default_encryption == EncryptionAlgorithm::None {
return Ok(data.to_vec());
}
let encryptor_name = self.default_encryption.as_str();
let encryptor = EncryptionUtil::find(&encryptor_name).ok_or_else(|| {
let registered = EncryptionUtil::list_registered();
let error_msg = format!(
"Encryptor '{}' not found. Registered: {:?}",
encryptor_name, registered
);
tracing::error!("{}", error_msg);
crate::common::error::FlareError::deserialization_error(error_msg)
})?;
match encryptor.decrypt(data) {
Ok(decrypted) => Ok(decrypted),
Err(e) => {
if allow_fallback {
tracing::trace!(
"解密失败,尝试作为未加密数据处理: encryption={:?}, data_len={}",
self.default_encryption,
data.len()
);
Ok(data.to_vec())
} else {
Err(crate::common::error::FlareError::deserialization_error(
format!(
"解密失败: encryption={:?}, error={}, data_len={}",
self.default_encryption,
e,
data.len()
),
))
}
}
}
}
fn decrypt_data(&self, data: &[u8]) -> Result<Vec<u8>> {
self.decrypt_data_with_fallback(data, true)
}
fn encrypt_data(&self, data: &[u8], encryption: EncryptionAlgorithm) -> Result<Vec<u8>> {
if encryption == EncryptionAlgorithm::None {
return Ok(data.to_vec());
}
let encryptor_name = encryption.as_str();
let encryptor = EncryptionUtil::find(&encryptor_name).ok_or_else(|| {
let registered = EncryptionUtil::list_registered();
let error_msg = format!(
"Encryptor '{}' not found. Registered: {:?}",
encryptor_name, registered
);
tracing::error!("{}", error_msg);
crate::common::error::FlareError::encoding_error(error_msg)
})?;
encryptor.encrypt(data)
}
fn decompress_data(&self, data: &[u8]) -> Result<Vec<u8>> {
self.decompress_data_with_fallback(data, true)
}
fn decompress_data_with_fallback(&self, data: &[u8], allow_fallback: bool) -> Result<Vec<u8>> {
if self.default_compression == CompressionAlgorithm::None {
return Ok(data.to_vec());
}
match CompressionUtil::auto_decompress(data) {
Ok((decompressed, detected_algorithm)) => {
if detected_algorithm != CompressionAlgorithm::None {
Ok(decompressed)
} else {
if allow_fallback {
tracing::trace!(
"自动检测未发现压缩,按阈值策略作为未压缩数据处理: compression={:?}, data_len={}",
self.default_compression,
data.len()
);
Ok(data.to_vec())
} else {
Err(crate::common::error::FlareError::deserialization_error(
format!(
"解压缩失败(严格模式): 配置了压缩 {:?} 但数据未压缩",
self.default_compression
),
))
}
}
}
Err(e) => {
if allow_fallback {
tracing::trace!(
"解压缩失败,尝试作为未压缩数据处理: compression={:?}, data_len={}",
self.default_compression,
data.len()
);
Ok(data.to_vec())
} else {
Err(crate::common::error::FlareError::deserialization_error(
format!(
"解压缩失败(严格模式): compression={:?}, error={}",
self.default_compression, e
),
))
}
}
}
}
#[allow(dead_code)]
fn parse_decompressed(&self, decompressed: &[u8]) -> Result<Frame> {
self.parse_decompressed_with_fallback(decompressed, true)
}
fn parse_decompressed_with_fallback(
&self,
decompressed: &[u8],
allow_fallback: bool,
) -> Result<Frame> {
let detected_serializers = SerializationUtil::auto_detect(decompressed);
for serializer in detected_serializers {
if let Ok(frame) = serializer.deserialize(decompressed) {
return Ok(frame);
}
}
if allow_fallback {
self.try_all_serializers(decompressed)
} else {
let serializer = if let Some(custom_name) = &self.custom_format_name {
SerializationUtil::get_serializer_by_name(custom_name)
} else {
SerializationUtil::get_serializer(self.default_format)
}
.ok_or_else(|| {
let format_info = if let Some(name) = &self.custom_format_name {
format!("custom format '{}'", name)
} else {
format!("format {:?}", self.default_format)
};
crate::common::error::FlareError::deserialization_error(format!(
"Serializer not found for {}",
format_info
))
})?;
serializer.deserialize(decompressed).map_err(|e| {
let format_info = if let Some(name) = &self.custom_format_name {
format!("custom format '{}'", name)
} else {
format!("format {:?}", self.default_format)
};
crate::common::error::FlareError::deserialization_error(format!(
"反序列化失败(严格模式): {}, error={}",
format_info, e
))
})
}
}
fn try_all_serializers(&self, data: &[u8]) -> Result<Frame> {
[SerializationFormat::Protobuf, SerializationFormat::Json]
.iter()
.find_map(|&format| {
SerializationUtil::get_serializer(format)
.and_then(|serializer| serializer.deserialize(data).ok())
})
.ok_or_else(|| {
crate::common::error::FlareError::deserialization_error(
"Failed to parse message: no compatible serializer found".to_string(),
)
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::common::protocol::{FrameBuilder, ping};
#[test]
fn test_parse_protobuf() {
let parser = MessageParser::protobuf();
let frame = FrameBuilder::new()
.with_command(crate::common::protocol::Command {
r#type: Some(
crate::common::protocol::flare::core::commands::command::Type::System(ping()),
),
})
.build();
let data = parser.serialize(&frame).unwrap();
let parsed = parser.parse(&data).unwrap();
assert_eq!(parsed.message_id, frame.message_id);
}
#[test]
fn test_parse_json() {
let parser = &PRE_NEGOTIATION_PARSER;
let frame = FrameBuilder::new()
.with_command(crate::common::protocol::Command {
r#type: Some(
crate::common::protocol::flare::core::commands::command::Type::System(ping()),
),
})
.build();
let data = parser.serialize(&frame).unwrap();
let parsed = parser.parse(&data).unwrap();
assert_eq!(parsed.message_id, frame.message_id);
}
#[test]
#[cfg(feature = "compression-gzip")]
fn small_payload_uses_negotiated_compression_and_strict_parse_accepts_it() {
let parser = MessageParser::new(
SerializationFormat::Protobuf,
CompressionAlgorithm::Gzip,
EncryptionAlgorithm::None,
);
let frame = FrameBuilder::new()
.with_command(crate::common::protocol::Command {
r#type: Some(
crate::common::protocol::flare::core::commands::command::Type::System(ping()),
),
})
.build();
let data = parser.serialize(&frame).unwrap();
let (_, detected) = CompressionUtil::auto_decompress(&data).unwrap();
assert_eq!(detected, CompressionAlgorithm::Gzip);
let parsed = parser.parse_with_fallback(&data, false).unwrap();
assert_eq!(parsed.message_id, frame.message_id);
}
#[test]
#[cfg(feature = "compression-gzip")]
fn large_payload_uses_negotiated_compression() {
let parser = MessageParser::new(
SerializationFormat::Protobuf,
CompressionAlgorithm::Gzip,
EncryptionAlgorithm::None,
);
let frame = FrameBuilder::new()
.with_metadata(
"padding".to_string(),
vec![b'x'; MIN_COMPRESSION_PAYLOAD_BYTES * 2],
)
.build();
let data = parser.serialize(&frame).unwrap();
let (_, detected) = CompressionUtil::auto_decompress(&data).unwrap();
assert_eq!(detected, CompressionAlgorithm::Gzip);
let parsed = parser.parse_with_fallback(&data, false).unwrap();
assert_eq!(parsed.message_id, frame.message_id);
}
#[test]
fn test_get_encryption_from_frame() {
use crate::common::protocol::FrameBuilder;
use std::collections::HashMap;
let frame = FrameBuilder::new().build();
assert_eq!(
MessageParser::get_encryption_from_frame(&frame),
EncryptionAlgorithm::None
);
let mut metadata = HashMap::new();
metadata.insert("encryption".to_string(), b"aes256gcm".to_vec());
let frame = FrameBuilder::new()
.with_metadata("encryption".to_string(), b"aes256gcm".to_vec())
.build();
assert_eq!(
MessageParser::get_encryption_from_frame(&frame),
EncryptionAlgorithm::Aes256Gcm
);
let frame = FrameBuilder::new()
.with_metadata("encryption".to_string(), b"invalid".to_vec())
.build();
assert_eq!(
MessageParser::get_encryption_from_frame(&frame),
EncryptionAlgorithm::Custom("invalid".to_string())
);
}
}