use std::collections::HashMap;
use std::str;
use bytes::Buf;
use bytes::BufMut;
use bytes::Bytes;
use bytes::BytesMut;
use cheetah_string::CheetahString;
use crate::protocol::command_custom_header::CommandCustomHeader;
use crate::protocol::command_custom_header::HeaderEncodeCapability;
use crate::protocol::header_codec::BinaryHeaderFields;
use crate::protocol::header_field_merge::has_custom_ext_collision;
use crate::protocol::header_field_merge::merge_header_and_dynamic;
use crate::protocol::remoting_command::RemotingCommand;
use crate::protocol::LanguageCode;
use crate::ProtocolContractViolation;
const MAP_ENTRY_ENCODED_BYTES_ESTIMATE: usize = 16;
const MIN_MAP_CAPACITY: usize = 4;
const MAX_MAP_CAPACITY: usize = 1024;
fn decoding_error(required: usize, available: usize) -> rocketmq_error::Error {
crate::error::serialization_decode_failed("binary", format!("required {required} bytes, got {available}"))
}
fn trailing_header_error(remaining: usize) -> rocketmq_error::Error {
crate::error::serialization_decode_failed(
"binary",
format!("ROCKETMQ header has {remaining} trailing bytes after extension fields"),
)
}
enum SortedExtFields<'a> {
Direct(std::collections::hash_map::Iter<'a, CheetahString, CheetahString>),
Sorted(std::vec::IntoIter<(&'a CheetahString, &'a CheetahString)>),
}
impl<'a> Iterator for SortedExtFields<'a> {
type Item = (&'a CheetahString, &'a CheetahString);
fn next(&mut self) -> Option<Self::Item> {
match self {
Self::Direct(entries) => entries.find(|(key, _)| !key.is_empty()),
Self::Sorted(entries) => entries.next(),
}
}
}
fn sorted_ext_fields(map: &HashMap<CheetahString, CheetahString>) -> SortedExtFields<'_> {
if map.len() <= 1 {
return SortedExtFields::Direct(map.iter());
}
let mut entries = map.iter().filter(|(key, _)| !key.is_empty()).collect::<Vec<_>>();
entries.sort_unstable_by(|(left, _), (right, _)| left.as_str().cmp(right.as_str()));
SortedExtFields::Sorted(entries.into_iter())
}
pub struct RocketMQSerializable;
impl RocketMQSerializable {
pub(crate) const INITIAL_ENCODE_CAPACITY: usize = 256;
#[inline]
fn estimated_map_capacity(encoded_len: usize) -> usize {
(encoded_len / MAP_ENTRY_ENCODED_BYTES_ESTIMATE).clamp(MIN_MAP_CAPACITY, MAX_MAP_CAPACITY)
}
#[inline]
pub fn write_str(buf: &mut BytesMut, use_short_length: bool, s: &str) -> usize {
let bytes = s.as_bytes();
let len = bytes.len();
let length_size = if use_short_length {
buf.put_u16(len as u16);
2
} else {
buf.put_u32(len as u32);
4
};
buf.put_slice(bytes); length_size + len
}
#[inline]
pub fn read_str(
buf: &mut BytesMut,
use_short_length: bool,
limit: usize,
) -> rocketmq_error::Result<Option<CheetahString>> {
let len = if use_short_length {
if buf.remaining() < 2 {
return Err(decoding_error(2, buf.remaining()));
}
buf.get_u16() as usize
} else {
if buf.remaining() < 4 {
return Err(decoding_error(4, buf.remaining()));
}
buf.get_u32() as usize
};
if len == 0 {
return Ok(None);
}
if len > limit {
return Err(decoding_error(len, limit));
}
if buf.remaining() < len {
return Err(decoding_error(len, buf.remaining()));
}
let bytes = buf.split_to(len).freeze();
let value = CheetahString::try_copy_from_bytes(bytes)
.map_err(|error| crate::error::serialization_source("decode", "binary", error.into_parts().1))?;
Ok(Some(value))
}
#[inline]
pub fn rocketmq_protocol_encode(cmd: &mut RemotingCommand, buf: &mut BytesMut) -> usize {
let checkpoint = buf.len();
match Self::try_rocketmq_protocol_encode(cmd, buf) {
Ok(encoded) => encoded,
Err(_) => {
buf.truncate(checkpoint);
0
}
}
}
#[inline]
pub fn try_rocketmq_protocol_encode(
cmd: &RemotingCommand,
buf: &mut BytesMut,
) -> Result<usize, ProtocolContractViolation> {
Self::try_rocketmq_protocol_encode_with_capability(cmd, buf, cmd.custom_header_encode_capability())
}
pub(crate) fn try_rocketmq_protocol_encode_with_capability(
cmd: &RemotingCommand,
buf: &mut BytesMut,
capability: HeaderEncodeCapability,
) -> Result<usize, ProtocolContractViolation> {
buf.reserve(Self::INITIAL_ENCODE_CAPACITY);
let checkpoint = buf.len();
let result = Self::try_rocketmq_protocol_encode_inner(cmd, buf, capability);
if result.is_err() {
buf.truncate(checkpoint);
}
result
}
fn try_rocketmq_protocol_encode_inner(
cmd: &RemotingCommand,
buf: &mut BytesMut,
capability: HeaderEncodeCapability,
) -> Result<usize, ProtocolContractViolation> {
let begin_index = buf.len();
buf.put_u16(cmd.code() as u16); buf.put_u8(cmd.language().get_code()); buf.put_u16(cmd.version() as u16); buf.put_i32(cmd.opaque()); buf.put_i32(cmd.flag());
if let Some(remark) = cmd.remark() {
Self::write_str(buf, false, remark.as_str());
} else {
buf.put_i32(0);
}
let map_len_index = buf.len();
buf.put_i32(0);
match cmd.command_custom_header_ref() {
Some(header)
if capability == HeaderEncodeCapability::DirectBinary
&& !has_custom_ext_collision(header, cmd.ext_fields()) =>
{
header.encode_direct_binary(buf)?;
if let Some(ext_fields) = cmd.ext_fields() {
Self::write_ext_fields(buf, ext_fields)?;
}
}
Some(header) => {
let merged = merge_header_and_dynamic(header, cmd.ext_fields())?;
Self::write_ext_fields(buf, &merged)?;
}
None => {
if let Some(ext_fields) = cmd.ext_fields() {
Self::write_ext_fields(buf, ext_fields)?;
}
}
}
let current_length = buf.len();
let ext_fields_length = Self::checked_ext_fields_length(current_length - map_len_index - 4)?;
buf[map_len_index..map_len_index + 4].copy_from_slice(&ext_fields_length.to_be_bytes());
Ok(buf.len() - begin_index)
}
pub(crate) fn try_rocketmq_protocol_encode_direct(
cmd: &RemotingCommand,
header: &dyn CommandCustomHeader,
buf: &mut BytesMut,
) -> Result<usize, ProtocolContractViolation> {
const FIXED_HEADER_LENGTH: usize = 21;
const EXT_FIELDS_LENGTH_OFFSET: usize = 17;
let begin_index = buf.len();
let mut fixed = [0u8; FIXED_HEADER_LENGTH];
fixed[..2].copy_from_slice(&(cmd.code() as u16).to_be_bytes());
fixed[2] = cmd.language().get_code();
fixed[3..5].copy_from_slice(&(cmd.version() as u16).to_be_bytes());
fixed[5..9].copy_from_slice(&cmd.opaque().to_be_bytes());
fixed[9..13].copy_from_slice(&cmd.flag().to_be_bytes());
buf.extend_from_slice(&fixed);
header.encode_direct_binary(buf)?;
let ext_fields_length = Self::checked_ext_fields_length(buf.len() - begin_index - FIXED_HEADER_LENGTH)?;
let map_len_index = begin_index + EXT_FIELDS_LENGTH_OFFSET;
buf[map_len_index..map_len_index + 4].copy_from_slice(&ext_fields_length.to_be_bytes());
Ok(buf.len() - begin_index)
}
#[inline]
fn write_ext_fields(
buf: &mut BytesMut,
fields: &HashMap<CheetahString, CheetahString>,
) -> Result<(), ProtocolContractViolation> {
for (key, value) in sorted_ext_fields(fields) {
let key_length = Self::checked_dynamic_key_length(key.len())?;
let value_length = Self::checked_dynamic_value_length(value.len())?;
buf.put_u16(key_length);
buf.put_slice(key.as_bytes());
buf.put_i32(value_length);
buf.put_slice(value.as_bytes());
}
Ok(())
}
#[inline]
fn checked_ext_fields_length(length: usize) -> Result<i32, ProtocolContractViolation> {
i32::try_from(length).map_err(|_| ProtocolContractViolation::ExtensionFieldsLengthOverflow)
}
#[inline]
fn checked_dynamic_key_length(length: usize) -> Result<u16, ProtocolContractViolation> {
i16::try_from(length)
.map(|length| length as u16)
.map_err(|_| ProtocolContractViolation::DynamicKeyLengthOverflow)
}
#[inline]
fn checked_dynamic_value_length(length: usize) -> Result<i32, ProtocolContractViolation> {
i32::try_from(length).map_err(|_| ProtocolContractViolation::DynamicValueLengthOverflow)
}
pub fn rocket_mq_protocol_encode_bytes(cmd: &RemotingCommand) -> Bytes {
let remark_bytes = cmd.remark().map(|remark| remark.as_bytes().to_vec());
let remark_len = remark_bytes.as_ref().map_or(0, |v| v.len());
let ext_fields_bytes = if let Some(ext) = cmd.get_ext_fields() {
Self::map_serialize(ext)
} else {
None
};
let ext_len = ext_fields_bytes.as_ref().map_or(0, |v| v.len());
let total_len = Self::cal_total_len(remark_len, ext_len);
let mut header_buffer = BytesMut::with_capacity(total_len);
header_buffer.put_i16(cmd.code() as i16);
header_buffer.put_u8(cmd.language().get_code());
header_buffer.put_i16(cmd.version() as i16);
header_buffer.put_i32(cmd.opaque());
header_buffer.put_i32(cmd.flag());
if let Some(remark_bytes) = remark_bytes {
header_buffer.put_i32(remark_bytes.len() as i32);
header_buffer.put(remark_bytes.as_ref());
} else {
header_buffer.put_i32(0);
}
if let Some(ext_fields_bytes) = ext_fields_bytes {
header_buffer.put_i32(ext_fields_bytes.len() as i32);
header_buffer.put(ext_fields_bytes.as_ref());
} else {
header_buffer.put_i32(0);
}
header_buffer.freeze()
}
#[inline]
pub fn map_serialize(map: &HashMap<CheetahString, CheetahString>) -> Option<BytesMut> {
if map.is_empty() {
return None;
}
let mut total_length = 0;
let mut valid_entries = 0;
for (key, value) in map.iter() {
if !key.is_empty() {
total_length += 2 + key.len() + 4 + value.len();
valid_entries += 1;
}
}
if valid_entries == 0 {
return None;
}
let mut content = BytesMut::with_capacity(total_length);
for (key, value) in sorted_ext_fields(map) {
content.put_u16(key.len() as u16);
content.put_slice(key.as_bytes());
content.put_i32(value.len() as i32);
content.put_slice(value.as_bytes());
}
Some(content)
}
pub fn cal_total_len(remark_len: usize, ext_len: usize) -> usize {
2 + 1 + 2 + 4 + 4 + 4 + remark_len + 4 + ext_len }
pub fn rocket_mq_protocol_decode(
header_buffer: &mut BytesMut,
header_len: usize,
) -> rocketmq_error::Result<RemotingCommand> {
let available = header_buffer.remaining();
if available < header_len {
return Err(decoding_error(header_len, available));
}
if available > header_len {
return Err(trailing_header_error(available - header_len));
}
Self::rocket_mq_protocol_decode_bytes(header_buffer.split().freeze())
}
pub(crate) fn rocket_mq_protocol_decode_bytes(header: Bytes) -> rocketmq_error::Result<RemotingCommand> {
const FIXED_HEADER_LEN: usize = 13;
const LENGTH_FIELD_LEN: usize = 4;
if header.len() < FIXED_HEADER_LEN {
return Err(decoding_error(FIXED_HEADER_LEN, header.len()));
}
let code = i16::from_be_bytes([header[0], header[1]]) as i32;
let language = LanguageCode::from(header[2]);
let version = i16::from_be_bytes([header[3], header[4]]) as i32;
let opaque = i32::from_be_bytes([header[5], header[6], header[7], header[8]]);
let flag = i32::from_be_bytes([header[9], header[10], header[11], header[12]]);
let mut cursor = FIXED_HEADER_LEN;
if header.len() - cursor < LENGTH_FIELD_LEN {
return Err(decoding_error(LENGTH_FIELD_LEN, header.len() - cursor));
}
let remark_length = u32::from_be_bytes([
header[cursor],
header[cursor + 1],
header[cursor + 2],
header[cursor + 3],
]) as usize;
cursor += LENGTH_FIELD_LEN;
if remark_length > header.len() - cursor {
return Err(decoding_error(remark_length, header.len() - cursor));
}
let remark = if remark_length == 0 {
None
} else {
let end = cursor + remark_length;
let value = CheetahString::try_copy_from_bytes(header.slice(cursor..end))
.map_err(|error| crate::error::serialization_source("decode", "binary", error.into_parts().1))?;
cursor = end;
Some(value)
};
if header.len() - cursor < LENGTH_FIELD_LEN {
return Err(decoding_error(LENGTH_FIELD_LEN, header.len() - cursor));
}
let ext_fields_length = i32::from_be_bytes([
header[cursor],
header[cursor + 1],
header[cursor + 2],
header[cursor + 3],
]);
cursor += LENGTH_FIELD_LEN;
let payload = if ext_fields_length > 0 {
let ext_fields_length = ext_fields_length as usize;
if ext_fields_length > header.len() - cursor {
return Err(decoding_error(ext_fields_length, header.len() - cursor));
}
let end = cursor + ext_fields_length;
let payload = header.slice(cursor..end);
cursor = end;
payload
} else {
Bytes::new()
};
if cursor != header.len() {
return Err(trailing_header_error(header.len() - cursor));
}
let ext_fields = BinaryHeaderFields::new(payload)?;
Ok(RemotingCommand::from_binary_wire_parts(
code, language, version, opaque, flag, remark, ext_fields,
))
}
#[inline]
pub fn map_deserialize(
buffer: &mut BytesMut,
len: usize,
) -> rocketmq_error::Result<HashMap<CheetahString, CheetahString>> {
if len == 0 {
return Ok(HashMap::new());
}
if len > buffer.remaining() {
return Err(decoding_error(len, buffer.remaining()));
}
let estimated_entries = Self::estimated_map_capacity(len);
let mut map = HashMap::with_capacity(estimated_entries);
let target_remaining = buffer.remaining().saturating_sub(len);
while buffer.remaining() > target_remaining {
let key = Self::read_str(buffer, true, len)?.ok_or_else(|| decoding_error(0, 0))?;
if let Some(value) = Self::read_str(buffer, false, len)? {
map.insert(key, value);
}
}
Ok(map)
}
}
#[cfg(test)]
mod tests {
use bytes::BufMut;
use bytes::BytesMut;
use super::*;
fn minimal_header_without_ext_len() -> BytesMut {
let mut buf = BytesMut::new();
buf.put_i16(0);
buf.put_u8(LanguageCode::JAVA.get_code());
buf.put_i16(0);
buf.put_i32(0);
buf.put_i32(0);
buf.put_i32(0);
buf
}
#[test]
fn write_str_short_length() {
let mut buf = BytesMut::new();
let written = RocketMQSerializable::write_str(&mut buf, true, "test");
assert_eq!(written, 6);
assert_eq!(buf, BytesMut::from(&[0, 4, 116, 101, 115, 116][..]));
}
#[test]
fn write_str_long_length() {
let mut buf = BytesMut::new();
let written = RocketMQSerializable::write_str(&mut buf, false, "test");
assert_eq!(written, 8);
assert_eq!(buf, BytesMut::from(&[0, 0, 0, 4, 116, 101, 115, 116][..]));
}
#[test]
fn read_str_short_length() {
let mut buf = BytesMut::from(&[0, 4, 116, 101, 115, 116][..]);
let read = RocketMQSerializable::read_str(&mut buf, true, 10).unwrap();
assert_eq!(read, Some("test".into()));
}
#[test]
fn read_str_long_length() {
let mut buf = BytesMut::from(&[0, 0, 0, 4, 116, 101, 115, 116][..]);
let read = RocketMQSerializable::read_str(&mut buf, false, 10).unwrap();
assert_eq!(read, Some("test".into()));
}
#[test]
fn read_str_exceeds_limit() {
let mut buf = BytesMut::from(&[0, 0, 0, 4, 116, 101, 115, 116][..]);
let read = RocketMQSerializable::read_str(&mut buf, false, 2);
assert!(read.is_err());
}
#[test]
fn read_str_rejects_invalid_utf8() {
let mut buf = BytesMut::from(&[0, 2, 0xff, 0xfe][..]);
let read = RocketMQSerializable::read_str(&mut buf, true, 10);
assert!(read.is_err());
}
#[test]
fn map_serialize_empty() {
let map = HashMap::new();
let serialized = RocketMQSerializable::map_serialize(&map);
assert!(serialized.is_none());
}
#[test]
fn map_serialize_non_empty() {
let mut map = HashMap::new();
map.insert("key".into(), "value".into());
let serialized = RocketMQSerializable::map_serialize(&map).unwrap();
assert_eq!(
serialized,
BytesMut::from(&[0, 3, 107, 101, 121, 0, 0, 0, 5, 118, 97, 108, 117, 101][..])
);
}
#[test]
fn map_serialize_preserves_present_empty_value() {
let map = HashMap::from([("key".into(), CheetahString::new())]);
let serialized = RocketMQSerializable::map_serialize(&map).unwrap();
assert_eq!(serialized, BytesMut::from(&[0, 3, 107, 101, 121, 0, 0, 0, 0][..]));
}
#[test]
fn map_deserialize_empty() {
let mut buf = BytesMut::new();
let deserialized = RocketMQSerializable::map_deserialize(&mut buf, 0).unwrap();
assert!(deserialized.is_empty());
}
#[test]
fn map_deserialize_non_empty() {
let mut buf = BytesMut::from(&[0, 3, 107, 101, 121, 0, 0, 0, 5, 118, 97, 108, 117, 101][..]);
let deserialized = RocketMQSerializable::map_deserialize(&mut buf, 14).unwrap();
assert_eq!(deserialized, [("key".into(), "value".into())].iter().cloned().collect());
}
#[test]
fn map_capacity_estimate_covers_compact_request_headers_and_stays_bounded() {
assert_eq!(RocketMQSerializable::estimated_map_capacity(1), 4);
assert_eq!(RocketMQSerializable::estimated_map_capacity(292), 18);
assert_eq!(RocketMQSerializable::estimated_map_capacity(385), 24);
assert_eq!(RocketMQSerializable::estimated_map_capacity(usize::MAX), 1024);
}
#[test]
fn map_deserialize_normalizes_zero_length_value_to_absent() {
let mut buf = BytesMut::from(&[0, 3, 107, 101, 121, 0, 0, 0, 0][..]);
let deserialized = RocketMQSerializable::map_deserialize(&mut buf, 9).unwrap();
assert!(deserialized.is_empty());
assert!(buf.is_empty());
}
#[test]
fn extension_field_length_checks_limit_and_limit_plus_one() {
assert_eq!(
RocketMQSerializable::checked_ext_fields_length(i32::MAX as usize).unwrap(),
i32::MAX
);
assert!(matches!(
RocketMQSerializable::checked_ext_fields_length(i32::MAX as usize + 1),
Err(ProtocolContractViolation::ExtensionFieldsLengthOverflow)
));
assert_eq!(
RocketMQSerializable::checked_dynamic_value_length(i32::MAX as usize).unwrap(),
i32::MAX
);
assert!(matches!(
RocketMQSerializable::checked_dynamic_value_length(i32::MAX as usize + 1),
Err(ProtocolContractViolation::DynamicValueLengthOverflow)
));
}
#[test]
fn dynamic_key_length_java_compatible() {
let maximum = i16::MAX as usize;
assert_eq!(
RocketMQSerializable::checked_dynamic_key_length(maximum - 1).unwrap(),
(i16::MAX - 1) as u16
);
assert_eq!(
RocketMQSerializable::checked_dynamic_key_length(maximum).unwrap(),
i16::MAX as u16
);
assert!(matches!(
RocketMQSerializable::checked_dynamic_key_length(maximum + 1),
Err(ProtocolContractViolation::DynamicKeyLengthOverflow)
));
}
#[test]
fn rocketmq_protocol_decode_rejects_short_fixed_header_without_panic() {
let mut buf = BytesMut::from(&[0_u8; 12][..]);
if RocketMQSerializable::rocket_mq_protocol_decode(&mut buf, 12).is_ok() {
panic!("short fixed header should decode to error");
}
}
#[test]
fn rocketmq_protocol_decode_rejects_missing_ext_length_without_panic() {
let mut buf = minimal_header_without_ext_len();
let header_len = buf.len();
if RocketMQSerializable::rocket_mq_protocol_decode(&mut buf, header_len).is_ok() {
panic!("missing ext length should decode to error");
}
}
#[test]
fn rocketmq_protocol_decode_rejects_truncated_ext_fields_without_panic() {
let mut buf = minimal_header_without_ext_len();
buf.put_i32(10);
let header_len = buf.len();
if RocketMQSerializable::rocket_mq_protocol_decode(&mut buf, header_len).is_ok() {
panic!("truncated ext fields should decode to error");
}
}
#[test]
fn rocketmq_protocol_decode_rejects_bytes_after_declared_ext_fields() {
let mut buf = minimal_header_without_ext_len();
buf.put_i32(0);
buf.put_u8(1);
let header_len = buf.len();
assert!(RocketMQSerializable::rocket_mq_protocol_decode(&mut buf, header_len).is_err());
}
#[test]
fn bytes_decoder_preserves_remark_and_extension_fields() {
let mut buf = BytesMut::new();
buf.put_i16(10);
buf.put_u8(LanguageCode::RUST.get_code());
buf.put_i16(501);
buf.put_i32(7);
buf.put_i32(1);
RocketMQSerializable::write_str(&mut buf, false, "remark");
let ext_length_offset = buf.len();
buf.put_i32(0);
let ext_start = buf.len();
buf.put_u16(3);
buf.extend_from_slice(b"key");
buf.put_i32(5);
buf.extend_from_slice(b"value");
let ext_length = (buf.len() - ext_start) as i32;
buf[ext_length_offset..ext_length_offset + 4].copy_from_slice(&ext_length.to_be_bytes());
let command = RocketMQSerializable::rocket_mq_protocol_decode_bytes(buf.freeze()).unwrap();
assert_eq!(command.code(), 10);
assert_eq!(command.language(), LanguageCode::RUST);
assert_eq!(command.version(), 501);
assert_eq!(command.opaque(), 7);
assert_eq!(command.flag(), 1);
assert_eq!(command.remark().map(CheetahString::as_str), Some("remark"));
assert_eq!(command.ext_fields().unwrap().get("key").unwrap(), "value");
}
#[test]
fn binary_decode_does_not_generate_an_outbound_request_id() {
let mut buf = BytesMut::new();
buf.put_i16(105);
buf.put_u8(LanguageCode::JAVA.get_code());
buf.put_i16(321);
buf.put_i32(0x0102_0304);
buf.put_i32(3);
RocketMQSerializable::write_str(&mut buf, false, "wire");
buf.put_i32(0);
let before = crate::protocol::remoting_command::request_id_generation_count();
let command = RocketMQSerializable::rocket_mq_protocol_decode_bytes(buf.freeze()).unwrap();
assert_eq!(
crate::protocol::remoting_command::request_id_generation_count(),
before,
"binary decode must not consume an outbound correlation ID"
);
assert_eq!(command.code(), 105);
assert_eq!(command.language(), LanguageCode::JAVA);
assert_eq!(command.version(), 321);
assert_eq!(command.opaque(), 0x0102_0304);
assert_eq!(command.flag(), 3);
assert_eq!(command.remark().map(CheetahString::as_str), Some("wire"));
assert!(command.ext_fields().is_some_and(HashMap::is_empty));
assert_eq!(command.get_serialize_type(), crate::protocol::SerializeType::ROCKETMQ);
}
}