rocketmq_protocol/protocol/header_codec/
sink.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::BufMut;
16use bytes::BytesMut;
17use cheetah_string::CheetahString;
18
19use super::private::Sealed;
20use super::write_json_string;
21use super::HeaderFieldContext;
22use super::HeaderValue;
23use super::ProtocolContractViolation;
24use crate::protocol::command_custom_header::HeaderMap;
25
26/// A statically dispatched destination for typed header fields.
27///
28/// This trait is sealed. Protocol-owned implementations preserve identical
29/// field semantics while allowing LLVM to specialize map and future binary
30/// writes without a trait object in the per-field path.
31pub trait EncodeSink: Sealed {
32    /// Writes one canonical field.
33    ///
34    /// The caller is responsible for header validation and Java range checks
35    /// before this method. Sinks must not duplicate those policy decisions.
36    ///
37    /// # Errors
38    ///
39    /// Returns a classified codec error when the destination wire format cannot
40    /// represent the key or value. Map writes are currently infallible.
41    fn write<V: HeaderValue>(
42        &mut self,
43        key: &'static str,
44        value: &V,
45        context: HeaderFieldContext,
46    ) -> Result<(), ProtocolContractViolation>;
47}
48
49/// An [`EncodeSink`] that appends typed fields directly to a [`HeaderMap`].
50pub struct MapSink<'a> {
51    out: &'a mut HeaderMap,
52}
53
54impl<'a> MapSink<'a> {
55    /// Creates a map sink over an existing destination.
56    #[inline]
57    pub const fn new(out: &'a mut HeaderMap) -> Self {
58        Self { out }
59    }
60
61    /// Returns the borrowed destination after the sink is no longer needed.
62    #[inline]
63    pub fn into_inner(self) -> &'a mut HeaderMap {
64        self.out
65    }
66}
67
68impl Sealed for MapSink<'_> {}
69
70impl EncodeSink for MapSink<'_> {
71    #[inline]
72    fn write<V: HeaderValue>(
73        &mut self,
74        key: &'static str,
75        value: &V,
76        _context: HeaderFieldContext,
77    ) -> Result<(), ProtocolContractViolation> {
78        self.out
79            .insert(CheetahString::from_static_str(key), value.to_map_value());
80        Ok(())
81    }
82}
83
84/// An [`EncodeSink`] that writes one JSON extension-field object directly.
85///
86/// All RocketMQ extension-field values remain JSON strings, including scalar
87/// Rust fields. String escaping is performed directly into the destination.
88pub struct JsonSink<'a> {
89    out: &'a mut BytesMut,
90    first: bool,
91}
92
93impl<'a> JsonSink<'a> {
94    /// Starts a JSON object in `out`.
95    #[inline]
96    pub fn new(out: &'a mut BytesMut) -> Self {
97        out.extend_from_slice(b"{");
98        Self { out, first: true }
99    }
100
101    /// Completes the JSON object and returns the destination.
102    #[inline]
103    pub fn finish(self) -> &'a mut BytesMut {
104        self.out.extend_from_slice(b"}");
105        self.out
106    }
107}
108
109impl Sealed for JsonSink<'_> {}
110
111impl EncodeSink for JsonSink<'_> {
112    #[inline]
113    fn write<V: HeaderValue>(
114        &mut self,
115        key: &'static str,
116        value: &V,
117        _context: HeaderFieldContext,
118    ) -> Result<(), ProtocolContractViolation> {
119        if self.first {
120            self.first = false;
121        } else {
122            self.out.extend_from_slice(b",");
123        }
124        write_json_string(self.out, key);
125        self.out.extend_from_slice(b":");
126        value.write_json_string(self.out);
127        Ok(())
128    }
129}
130
131/// An [`EncodeSink`] that writes canonical extension fields directly to a
132/// ROCKETMQ binary payload.
133///
134/// Each field uses a big-endian `u16` key length followed by a big-endian
135/// signed Java `int` value length. Scalar values are appended through
136/// [`HeaderValue::write_ascii`] without allocating an intermediate string.
137pub struct BinarySink<'a> {
138    out: &'a mut BytesMut,
139}
140
141impl<'a> BinarySink<'a> {
142    /// Creates a binary sink over an existing destination.
143    #[inline]
144    pub fn new(out: &'a mut BytesMut) -> Self {
145        Self { out }
146    }
147
148    /// Returns the borrowed destination after the sink is no longer needed.
149    #[inline]
150    pub fn into_inner(self) -> &'a mut BytesMut {
151        self.out
152    }
153}
154
155impl Sealed for BinarySink<'_> {}
156
157impl EncodeSink for BinarySink<'_> {
158    #[inline]
159    fn write<V: HeaderValue>(
160        &mut self,
161        key: &'static str,
162        value: &V,
163        context: HeaderFieldContext,
164    ) -> Result<(), ProtocolContractViolation> {
165        let key_len = u16::try_from(key.len()).map_err(|_| ProtocolContractViolation::KeyLengthOverflow {
166            header: context.header,
167            key: context.key,
168        })?;
169        if !V::ALWAYS_FITS_WIRE_LENGTH && value.encoded_len() > i32::MAX as usize {
170            return Err(ProtocolContractViolation::ValueLengthOverflow {
171                header: context.header,
172                key: context.key,
173            });
174        }
175
176        let pair_start = self.out.len();
177        self.out.put_u16(key_len);
178        self.out.extend_from_slice(key.as_bytes());
179
180        let value_len_offset = self.out.len();
181        self.out.put_i32(0);
182        let value_offset = self.out.len();
183        value.write_ascii(self.out);
184        let actual_value_len = self.out.len() - value_offset;
185        let actual_value_len = i32::try_from(actual_value_len).map_err(|_| {
186            self.out.truncate(pair_start);
187            ProtocolContractViolation::ValueLengthOverflow {
188                header: context.header,
189                key: context.key,
190            }
191        })?;
192        self.out[value_len_offset..value_len_offset + 4].copy_from_slice(&actual_value_len.to_be_bytes());
193        Ok(())
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200    use crate::protocol::header_codec::HeaderValueKind;
201
202    const CONTEXT: HeaderFieldContext =
203        HeaderFieldContext::new("ExampleHeader", "topic", HeaderValueKind::String, None);
204
205    #[test]
206    fn writes_into_existing_map_without_an_intermediate_map() {
207        let mut fields = HeaderMap::with_capacity(2);
208        fields.insert(
209            CheetahString::from_static_str("existing"),
210            CheetahString::from_static_str("value"),
211        );
212
213        let mut sink = MapSink::new(&mut fields);
214        sink.write("topic", &CheetahString::from("测试-topic"), CONTEXT)
215            .unwrap();
216        let fields = sink.into_inner();
217
218        assert_eq!(fields.len(), 2);
219        assert_eq!(fields.get("topic").map(CheetahString::as_str), Some("测试-topic"));
220        assert_eq!(fields.get("existing").map(CheetahString::as_str), Some("value"));
221    }
222
223    #[test]
224    fn canonical_write_replaces_an_existing_value() {
225        let mut fields = HeaderMap::new();
226        fields.insert(
227            CheetahString::from_static_str("topic"),
228            CheetahString::from_static_str("old"),
229        );
230
231        MapSink::new(&mut fields)
232            .write("topic", &String::from("new"), CONTEXT)
233            .unwrap();
234
235        assert_eq!(fields.get("topic").map(CheetahString::as_str), Some("new"));
236    }
237
238    #[test]
239    fn scalar_write_uses_the_header_value_canonical_form() {
240        let mut fields = HeaderMap::new();
241        let mut sink = MapSink::new(&mut fields);
242
243        sink.write(
244            "queueOffset",
245            &u64::MAX,
246            HeaderFieldContext::new("ExampleHeader", "queueOffset", HeaderValueKind::U64, None),
247        )
248        .unwrap();
249
250        assert_eq!(
251            fields.get("queueOffset").map(CheetahString::as_str),
252            Some("18446744073709551615")
253        );
254    }
255
256    #[test]
257    fn binary_sink_appends_canonical_pairs_without_replacing_existing_bytes() {
258        let mut out = BytesMut::from(&b"prefix"[..]);
259        let mut sink = BinarySink::new(&mut out);
260        sink.write("queueOffset", &-42_i64, CONTEXT).unwrap();
261        let out = sink.into_inner();
262
263        assert_eq!(out.len() - 6, 2 + 11 + 4 + 3);
264        assert_eq!(&out[..6], b"prefix");
265        assert_eq!(u16::from_be_bytes(out[6..8].try_into().unwrap()), 11);
266        assert_eq!(&out[8..19], b"queueOffset");
267        assert_eq!(i32::from_be_bytes(out[19..23].try_into().unwrap()), 3);
268        assert_eq!(&out[23..], b"-42");
269    }
270
271    #[test]
272    fn binary_sink_rejects_an_oversized_key_without_mutating_the_destination() {
273        let key = Box::leak("k".repeat(u16::MAX as usize + 1).into_boxed_str());
274        let mut out = BytesMut::from(&b"prefix"[..]);
275        let error = BinarySink::new(&mut out).write(key, &true, CONTEXT).unwrap_err();
276
277        assert!(matches!(error, ProtocolContractViolation::KeyLengthOverflow { .. }));
278        assert_eq!(out.as_ref(), b"prefix");
279    }
280
281    #[test]
282    fn json_sink_writes_string_scalars_and_escapes_text_without_allocating_values() {
283        let mut out = BytesMut::new();
284        let mut sink = JsonSink::new(&mut out);
285        sink.write("topic", &CheetahString::from("主题\"\\\n"), CONTEXT)
286            .unwrap();
287        sink.write(
288            "queueOffset",
289            &-42_i64,
290            HeaderFieldContext::new("ExampleHeader", "queueOffset", HeaderValueKind::I64, None),
291        )
292        .unwrap();
293        sink.write(
294            "enabled",
295            &true,
296            HeaderFieldContext::new("ExampleHeader", "enabled", HeaderValueKind::Bool, None),
297        )
298        .unwrap();
299        sink.finish();
300
301        let value: serde_json::Value = serde_json::from_slice(&out).unwrap();
302        assert_eq!(value["topic"], "主题\"\\\n");
303        assert_eq!(value["queueOffset"], "-42");
304        assert_eq!(value["enabled"], "true");
305    }
306}
rocketmq_protocol/protocol/header_codec/sink.rs

rocketmq_protocol/protocol/header_codec/
sink.rs
