rocketmq_protocol/protocol/header_codec/
field_source.rs

1// Copyright 2026 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//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use bytes::Bytes;
16use cheetah_string::CheetahString;
17
18use super::private::FieldSourceSealed;
19use crate::HeaderMap;
20
21const KEY_LENGTH_BYTES: usize = 2;
22const VALUE_LENGTH_BYTES: usize = 4;
23const MAX_INITIAL_MAP_CAPACITY: usize = 1024;
24
25/// A validated, borrowed view of remoting command extension fields.
26///
27/// This trait is sealed because source implementations must guarantee that
28/// every visited key and value is valid immutable UTF-8 for the full source
29/// borrow. Returning `false` from the visitor stops the scan early.
30pub trait HeaderFieldSource: FieldSourceSealed {
31    /// Visits fields without allocating owned key/value strings.
32    fn visit_fields_while<'a>(&'a self, visitor: &mut dyn FnMut(&'a str, &'a str) -> bool);
33
34    /// Produces the owned compatibility representation.
35    fn to_header_map(&self) -> HeaderMap;
36}
37
38impl FieldSourceSealed for HeaderMap {}
39
40impl HeaderFieldSource for HeaderMap {
41    #[inline]
42    fn visit_fields_while<'a>(&'a self, visitor: &mut dyn FnMut(&'a str, &'a str) -> bool) {
43        for (key, value) in self {
44            if !visitor(key.as_str(), value.as_str()) {
45                break;
46            }
47        }
48    }
49
50    #[inline]
51    fn to_header_map(&self) -> HeaderMap {
52        self.clone()
53    }
54}
55
56#[cold]
57#[inline(never)]
58fn malformed_binary_fields(reason: &'static str) -> rocketmq_error::Error {
59    crate::error::serialization_decode_failed("binary-header-fields", reason)
60}
61
62/// A validated, immutable ROCKETMQ extension-field payload.
63///
64/// Construction validates the complete payload before it is retained by a
65/// remoting command. Subsequent scans therefore need no allocation and cannot
66/// observe mutable bytes.
67#[derive(Clone)]
68pub(crate) struct BinaryHeaderFields {
69    payload: Bytes,
70    entry_count: usize,
71}
72
73impl BinaryHeaderFields {
74    /// Validates and retains one complete extension-field payload.
75    pub(crate) fn new(payload: Bytes) -> rocketmq_error::Result<Self> {
76        let entry_count = Self::validate(&payload)?;
77        Ok(Self { payload, entry_count })
78    }
79
80    pub(crate) const fn len(&self) -> usize {
81        self.entry_count
82    }
83
84    /// Materializes the compatibility map. The payload has already been
85    /// validated, so iteration uses its immutable representation invariant.
86    pub(crate) fn materialize(&self) -> HeaderMap {
87        let mut map = HeaderMap::with_capacity(self.entry_count.min(MAX_INITIAL_MAP_CAPACITY));
88        for (key, value) in self.iter() {
89            map.insert(CheetahString::from_slice(key), CheetahString::from_slice(value));
90        }
91        map
92    }
93
94    #[inline]
95    fn iter(&self) -> BinaryHeaderFieldIter<'_> {
96        BinaryHeaderFieldIter {
97            payload: &self.payload,
98            cursor: 0,
99        }
100    }
101
102    fn validate(payload: &[u8]) -> rocketmq_error::Result<usize> {
103        let mut cursor = 0usize;
104        let mut entry_count = 0usize;
105        while cursor < payload.len() {
106            let key_length = Self::read_u16(payload, &mut cursor)?;
107            if key_length == 0 {
108                return Err(malformed_binary_fields("extension-field key is empty"));
109            }
110            Self::read_utf8(payload, &mut cursor, key_length, "truncated extension-field key")?;
111
112            let value_length = Self::read_i32(payload, &mut cursor)?;
113            if value_length < 0 {
114                return Err(malformed_binary_fields("extension-field value length is negative"));
115            }
116            let value = Self::read_utf8(
117                payload,
118                &mut cursor,
119                value_length as usize,
120                "truncated extension-field value",
121            )?;
122
123            // Java-compatible ROCKETMQ map decoding treats zero-length values
124            // as absent rather than storing an empty string.
125            if !value.is_empty() {
126                entry_count = entry_count.saturating_add(1);
127            }
128        }
129        Ok(entry_count)
130    }
131
132    fn read_u16(payload: &[u8], cursor: &mut usize) -> rocketmq_error::Result<usize> {
133        let bytes = Self::take(payload, cursor, KEY_LENGTH_BYTES, "missing extension-field key length")?;
134        Ok(u16::from_be_bytes([bytes[0], bytes[1]]) as usize)
135    }
136
137    fn read_i32(payload: &[u8], cursor: &mut usize) -> rocketmq_error::Result<i32> {
138        let bytes = Self::take(
139            payload,
140            cursor,
141            VALUE_LENGTH_BYTES,
142            "missing extension-field value length",
143        )?;
144        Ok(i32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
145    }
146
147    #[inline(never)]
148    fn read_utf8<'a>(
149        payload: &'a [u8],
150        cursor: &mut usize,
151        length: usize,
152        truncated_reason: &'static str,
153    ) -> rocketmq_error::Result<&'a str> {
154        let bytes = Self::take(payload, cursor, length, truncated_reason)?;
155        std::str::from_utf8(bytes).map_err(|_| malformed_binary_fields("extension-field text is not valid UTF-8"))
156    }
157
158    fn take<'a>(
159        payload: &'a [u8],
160        cursor: &mut usize,
161        length: usize,
162        reason: &'static str,
163    ) -> rocketmq_error::Result<&'a [u8]> {
164        let end = cursor
165            .checked_add(length)
166            .ok_or_else(|| malformed_binary_fields(reason))?;
167        let bytes = payload
168            .get(*cursor..end)
169            .ok_or_else(|| malformed_binary_fields(reason))?;
170        *cursor = end;
171        Ok(bytes)
172    }
173}
174
175impl FieldSourceSealed for BinaryHeaderFields {}
176
177impl HeaderFieldSource for BinaryHeaderFields {
178    #[inline]
179    fn visit_fields_while<'a>(&'a self, visitor: &mut dyn FnMut(&'a str, &'a str) -> bool) {
180        for (key, value) in self.iter() {
181            if !visitor(key, value) {
182                break;
183            }
184        }
185    }
186
187    #[inline]
188    fn to_header_map(&self) -> HeaderMap {
189        self.materialize()
190    }
191}
192
193struct BinaryHeaderFieldIter<'a> {
194    payload: &'a [u8],
195    cursor: usize,
196}
197
198impl<'a> BinaryHeaderFieldIter<'a> {
199    #[inline]
200    fn take(&mut self, length: usize) -> Option<&'a [u8]> {
201        let start = self.cursor;
202        let end = start.checked_add(length)?;
203        let bytes = self.payload.get(start..end)?;
204        self.cursor = end;
205        Some(bytes)
206    }
207
208    #[inline]
209    fn read_u16(&mut self) -> Option<usize> {
210        let bytes: [u8; KEY_LENGTH_BYTES] = self.take(KEY_LENGTH_BYTES)?.try_into().ok()?;
211        Some(u16::from_be_bytes(bytes) as usize)
212    }
213
214    #[inline]
215    fn read_i32(&mut self) -> Option<i32> {
216        let bytes: [u8; VALUE_LENGTH_BYTES] = self.take(VALUE_LENGTH_BYTES)?.try_into().ok()?;
217        Some(i32::from_be_bytes(bytes))
218    }
219
220    #[inline]
221    fn read_utf8(&mut self, length: usize) -> Option<&'a str> {
222        std::str::from_utf8(self.take(length)?).ok()
223    }
224}
225
226impl<'a> Iterator for BinaryHeaderFieldIter<'a> {
227    type Item = (&'a str, &'a str);
228
229    #[inline]
230    fn next(&mut self) -> Option<Self::Item> {
231        while self.cursor < self.payload.len() {
232            let key_length = self.read_u16()?;
233            let key = self.read_utf8(key_length)?;
234            let value_length = usize::try_from(self.read_i32()?).ok()?;
235            let value = self.read_utf8(value_length)?;
236            if !value.is_empty() {
237                return Some((key, value));
238            }
239        }
240        None
241    }
242}
243
244#[cfg(test)]
245mod tests {
246    use bytes::BufMut;
247    use bytes::BytesMut;
248
249    use super::*;
250
251    fn entry(out: &mut BytesMut, key: &[u8], value: &[u8]) {
252        out.put_u16(key.len() as u16);
253        out.extend_from_slice(key);
254        out.put_i32(value.len() as i32);
255        out.extend_from_slice(value);
256    }
257
258    #[test]
259    fn validates_and_materializes_duplicate_and_empty_values() {
260        let mut payload = BytesMut::new();
261        entry(&mut payload, b"key", b"first");
262        entry(&mut payload, b"empty", b"");
263        entry(&mut payload, b"key", b"last");
264
265        let fields = BinaryHeaderFields::new(payload.freeze()).unwrap();
266        let map = fields.materialize();
267
268        assert_eq!(map.len(), 1);
269        assert_eq!(map.get("key").map(CheetahString::as_str), Some("last"));
270        assert!(!map.contains_key("empty"));
271    }
272
273    #[test]
274    fn validates_and_materializes_two_binary_extension_fields() {
275        let mut payload = BytesMut::new();
276        entry(&mut payload, b"alpha", b"first");
277        entry(&mut payload, b"zeta", b"last");
278
279        let fields = BinaryHeaderFields::new(payload.freeze()).expect("valid binary extension fields");
280        let map = fields.materialize();
281
282        assert_eq!(map.get("alpha").map(CheetahString::as_str), Some("first"));
283        assert_eq!(map.get("zeta").map(CheetahString::as_str), Some("last"));
284    }
285
286    #[test]
287    fn rejects_truncated_negative_empty_key_and_invalid_utf8_payloads() {
288        let invalid_payloads = [
289            Bytes::from_static(&[0]),
290            Bytes::from_static(&[0, 1]),
291            Bytes::from_static(&[0, 1, b'k', 0, 0, 0]),
292            Bytes::from_static(&[0, 1, b'k', 0xff, 0xff, 0xff, 0xff]),
293            Bytes::from_static(&[0, 0, 0, 0, 0, 1, b'v']),
294            Bytes::from_static(&[0, 1, 0xff, 0, 0, 0, 1, b'v']),
295            Bytes::from_static(&[0, 1, b'k', 0, 0, 0, 1, 0xff]),
296            Bytes::from_static(&[0, 1, b'k', 0, 0, 4, 0]),
297        ];
298
299        for payload in invalid_payloads {
300            assert!(BinaryHeaderFields::new(payload).is_err());
301        }
302    }
303
304    #[test]
305    fn iterator_fails_closed_if_an_internal_payload_invariant_is_broken() {
306        let malformed_payloads = [
307            Bytes::from_static(&[0, 2, b'k']),
308            Bytes::from_static(&[0, 1, b'k', 0xff, 0xff, 0xff, 0xff]),
309            Bytes::from_static(&[0, 1, 0xff, 0, 0, 0, 1, b'v']),
310        ];
311
312        for payload in malformed_payloads {
313            let fields = BinaryHeaderFields {
314                payload,
315                entry_count: 1,
316            };
317            assert!(fields.iter().next().is_none());
318        }
319    }
320}
rocketmq_protocol/protocol/header_codec/field_source.rs

rocketmq_protocol/protocol/header_codec/
field_source.rs
