rocketmq_protocol/protocol/body/
message_codec.rs

1// Copyright 2023 The RocketMQ Rust Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8
9use std::collections::HashMap;
10use std::net::Ipv4Addr;
11use std::net::Ipv6Addr;
12use std::net::SocketAddr;
13use std::net::SocketAddrV4;
14use std::net::SocketAddrV6;
15
16use bytes::Buf;
17use bytes::Bytes;
18use cheetah_string::CheetahString;
19use rocketmq_error::Error;
20use rocketmq_error::Result;
21
22use rocketmq_model::common::sys_flag::message_sys_flag::MessageSysFlag;
23
24pub const MESSAGE_MAGIC_CODE: i32 = -626_843_481;
25pub const MESSAGE_MAGIC_CODE_V2: i32 = -626_843_477;
26pub const NAME_VALUE_SEPARATOR: u8 = 1;
27pub const PROPERTY_SEPARATOR: u8 = 2;
28
29#[derive(Clone, Debug, Eq, PartialEq)]
30pub struct DecodedMessageFrame {
31    pub store_size: i32,
32    pub magic_code: i32,
33    pub body_crc: u32,
34    pub queue_id: i32,
35    pub flag: i32,
36    pub queue_offset: i64,
37    pub commit_log_offset: i64,
38    pub sys_flag: i32,
39    pub born_timestamp: i64,
40    pub born_host: SocketAddr,
41    pub store_timestamp: i64,
42    pub store_host: SocketAddr,
43    pub reconsume_times: i32,
44    pub prepared_transaction_offset: i64,
45    pub body: Bytes,
46    pub topic: CheetahString,
47    pub properties_string: CheetahString,
48    pub properties: HashMap<CheetahString, CheetahString>,
49}
50
51pub fn decode_message_frame(input: &mut Bytes) -> Result<DecodedMessageFrame> {
52    let initial_len = input.remaining();
53    let store_size = read_i32(input)?;
54    if store_size <= 0 || store_size as usize > initial_len {
55        return Err(invalid("invalid total message size"));
56    }
57    let magic_code = read_i32(input)?;
58    if magic_code != MESSAGE_MAGIC_CODE && magic_code != MESSAGE_MAGIC_CODE_V2 {
59        return Err(invalid("unknown message magic code"));
60    }
61    let body_crc = read_u32(input)?;
62    let queue_id = read_i32(input)?;
63    let flag = read_i32(input)?;
64    let queue_offset = read_i64(input)?;
65    let commit_log_offset = read_i64(input)?;
66    let sys_flag = read_i32(input)?;
67    let born_timestamp = read_i64(input)?;
68    let born_host = read_host(input, sys_flag & MessageSysFlag::BORNHOST_V6_FLAG != 0)?;
69    let store_timestamp = read_i64(input)?;
70    let store_host = read_host(input, sys_flag & MessageSysFlag::STOREHOSTADDRESS_V6_FLAG != 0)?;
71    let reconsume_times = read_i32(input)?;
72    let prepared_transaction_offset = read_i64(input)?;
73    let body = read_sized_i32(input)?;
74    let topic_len = if magic_code == MESSAGE_MAGIC_CODE_V2 {
75        let value = read_i16(input)?;
76        usize::try_from(value).map_err(|_| invalid("negative topic length"))?
77    } else {
78        usize::from(read_u8(input)?)
79    };
80    let topic_bytes = take(input, topic_len)?;
81    let topic = std::str::from_utf8(&topic_bytes)
82        .map(CheetahString::from_slice)
83        .map_err(invalid_source)?;
84    let property_len = usize::try_from(read_i16(input)?).map_err(|_| invalid("negative properties length"))?;
85    let property_bytes = take(input, property_len)?;
86    let properties_string = CheetahString::from(String::from_utf8_lossy(&property_bytes).as_ref());
87    let properties = decode_properties(&property_bytes);
88
89    Ok(DecodedMessageFrame {
90        store_size,
91        magic_code,
92        body_crc,
93        queue_id,
94        flag,
95        queue_offset,
96        commit_log_offset,
97        sys_flag,
98        born_timestamp,
99        born_host,
100        store_timestamp,
101        store_host,
102        reconsume_times,
103        prepared_transaction_offset,
104        body,
105        topic,
106        properties_string,
107        properties,
108    })
109}
110
111pub fn decode_properties(bytes: &[u8]) -> HashMap<CheetahString, CheetahString> {
112    bytes
113        .split(|byte| *byte == PROPERTY_SEPARATOR)
114        .filter_map(|entry| {
115            let separator = entry.iter().position(|byte| *byte == NAME_VALUE_SEPARATOR)?;
116            if separator == 0 || separator + 1 >= entry.len() {
117                return None;
118            }
119            let key = String::from_utf8_lossy(&entry[..separator]);
120            let value = String::from_utf8_lossy(&entry[separator + 1..]);
121            Some((CheetahString::from(key.as_ref()), CheetahString::from(value.as_ref())))
122        })
123        .collect()
124}
125
126fn invalid(reason: &'static str) -> Error {
127    crate::error::invalid_body("decode_message_frame", reason)
128}
129
130fn invalid_source(source: impl std::error::Error + Send + Sync + 'static) -> Error {
131    crate::error::invalid_body_source("decode_message_frame", source)
132}
133
134fn ensure(input: &Bytes, len: usize) -> Result<()> {
135    if input.remaining() < len {
136        Err(invalid("truncated message frame"))
137    } else {
138        Ok(())
139    }
140}
141
142fn read_u8(input: &mut Bytes) -> Result<u8> {
143    ensure(input, 1)?;
144    Ok(input.get_u8())
145}
146fn read_i16(input: &mut Bytes) -> Result<i16> {
147    ensure(input, 2)?;
148    Ok(input.get_i16())
149}
150fn read_i32(input: &mut Bytes) -> Result<i32> {
151    ensure(input, 4)?;
152    Ok(input.get_i32())
153}
154fn read_u32(input: &mut Bytes) -> Result<u32> {
155    ensure(input, 4)?;
156    Ok(input.get_u32())
157}
158fn read_i64(input: &mut Bytes) -> Result<i64> {
159    ensure(input, 8)?;
160    Ok(input.get_i64())
161}
162fn take(input: &mut Bytes, len: usize) -> Result<Bytes> {
163    ensure(input, len)?;
164    Ok(input.split_to(len))
165}
166fn read_sized_i32(input: &mut Bytes) -> Result<Bytes> {
167    let len = usize::try_from(read_i32(input)?).map_err(|_| invalid("negative body length"))?;
168    take(input, len)
169}
170
171fn read_host(input: &mut Bytes, ipv6: bool) -> Result<SocketAddr> {
172    if ipv6 {
173        let bytes = take(input, 16)?;
174        let mut address = [0_u8; 16];
175        address.copy_from_slice(&bytes);
176        let port = read_i32(input)? as u16;
177        Ok(SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::from(address), port, 0, 0)))
178    } else {
179        let bytes = take(input, 4)?;
180        let mut address = [0_u8; 4];
181        address.copy_from_slice(&bytes);
182        let port = read_i32(input)? as u16;
183        Ok(SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::from(address), port)))
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use bytes::BufMut;
190    use bytes::BytesMut;
191
192    use super::*;
193
194    fn golden_frame() -> Bytes {
195        let mut frame = BytesMut::new();
196        frame.put_i32(0);
197        frame.put_i32(MESSAGE_MAGIC_CODE);
198        frame.put_u32(7);
199        frame.put_i32(3);
200        frame.put_i32(9);
201        frame.put_i64(11);
202        frame.put_i64(13);
203        frame.put_i32(0);
204        frame.put_i64(17);
205        frame.put_slice(&[127, 0, 0, 1]);
206        frame.put_i32(10911);
207        frame.put_i64(19);
208        frame.put_slice(&[127, 0, 0, 2]);
209        frame.put_i32(10912);
210        frame.put_i32(2);
211        frame.put_i64(23);
212        frame.put_i32(4);
213        frame.put_slice(b"body");
214        frame.put_u8(5);
215        frame.put_slice(b"topic");
216        frame.put_i16(8);
217        frame.put_slice(b"k\x01value\x02");
218        let len = frame.len() as i32;
219        frame[..4].copy_from_slice(&len.to_be_bytes());
220        frame.freeze()
221    }
222
223    #[test]
224    fn decodes_frozen_v1_frame() {
225        let decoded = decode_message_frame(&mut golden_frame()).unwrap();
226        assert_eq!(decoded.topic, "topic");
227        assert_eq!(decoded.body, Bytes::from_static(b"body"));
228        assert_eq!(decoded.properties.get("k").map(CheetahString::as_str), Some("value"));
229        assert_eq!(decoded.born_host.to_string(), "127.0.0.1:10911");
230    }
231
232    #[test]
233    fn rejects_truncated_and_unknown_frames() {
234        let mut truncated = golden_frame().slice(..20);
235        assert!(decode_message_frame(&mut truncated).is_err());
236        let mut unknown = golden_frame();
237        let mut bytes = BytesMut::from(unknown.as_ref());
238        bytes[4..8].copy_from_slice(&0_i32.to_be_bytes());
239        unknown = bytes.freeze();
240        assert!(decode_message_frame(&mut unknown).is_err());
241    }
242}
rocketmq_protocol/protocol/body/message_codec.rs

rocketmq_protocol/protocol/body/
message_codec.rs
