iridium_stomp/codec.rs
1use bytes::{Buf, BufMut, BytesMut};
2use std::io;
3use tokio_util::codec::{Decoder, Encoder};
4
5use crate::frame::Frame;
6use crate::parser::{DEFAULT_MAX_FRAME_SIZE, parse_frame_slice_bounded, unescape_header_value};
7
8/// Escape a STOMP 1.2 header value for wire transmission.
9///
10/// Per STOMP 1.2 spec, the following characters must be escaped:
11/// - backslash (0x5c) → `\\`
12/// - carriage return (0x0d) → `\r`
13/// - line feed (0x0a) → `\n`
14/// - colon (0x3a) → `\c` (primarily for header names, but we escape in values too for safety)
15fn escape_header_value(input: &str) -> String {
16 let mut result = String::with_capacity(input.len());
17 for ch in input.chars() {
18 match ch {
19 '\\' => result.push_str("\\\\"),
20 '\r' => result.push_str("\\r"),
21 '\n' => result.push_str("\\n"),
22 ':' => result.push_str("\\c"),
23 _ => result.push(ch),
24 }
25 }
26 result
27}
28
29/// (parser-based implementation uses `src` directly; header parsing is
30/// delegated to the `parser` module.)
31/// Items produced or consumed by the codec.
32///
33/// A `StompItem` is either a decoded `Frame` or a `Heartbeat` marker
34/// representing a single LF received on the wire.
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub enum StompItem {
37 /// A decoded STOMP frame (command + headers + body)
38 Frame(Frame),
39 /// A single heartbeat pulse (LF)
40 Heartbeat,
41}
42
43/// `StompCodec` implements `tokio_util::codec::{Decoder, Encoder}` for the
44/// STOMP wire protocol.
45///
46/// Responsibilities:
47/// - Decode incoming bytes into `StompItem::Frame` or `StompItem::Heartbeat`.
48/// - Support both NUL-terminated frames and frames using the `content-length`
49/// header (STOMP 1.2) for binary bodies containing NUL bytes.
50/// - Encode `StompItem` back into bytes for the wire format and emit
51/// `content-length` when necessary.
52pub struct StompCodec {
53 // No internal buffer: we parse directly from the provided `src` buffer.
54 /// Largest frame this codec will decode before returning an error, in
55 /// bytes. Bounds both an oversized `content-length` and a frame that never
56 /// terminates, so neither can exhaust memory.
57 max_frame_size: usize,
58}
59
60impl StompCodec {
61 /// Create a new `StompCodec` bounding frames at
62 /// [`DEFAULT_MAX_FRAME_SIZE`].
63 pub fn new() -> Self {
64 Self {
65 max_frame_size: DEFAULT_MAX_FRAME_SIZE,
66 }
67 }
68
69 /// Create a `StompCodec` that rejects any frame larger than
70 /// `max_frame_size` bytes.
71 pub fn with_max_frame_size(max_frame_size: usize) -> Self {
72 Self { max_frame_size }
73 }
74}
75
76impl Default for StompCodec {
77 fn default() -> Self {
78 Self::new()
79 }
80}
81
82impl Decoder for StompCodec {
83 type Item = StompItem;
84 type Error = io::Error;
85 /// Decode bytes from `src` into a `StompItem`.
86 ///
87 /// Parameters
88 /// - `src`: a mutable reference to the read buffer containing bytes from the
89 /// transport. The decoder may consume bytes from this buffer (using
90 /// methods like `advance` or `split_to`) when it successfully decodes a
91 /// frame. If there are not enough bytes to form a complete frame, this
92 /// method should return `Ok(None)` and leave `src` in the same state.
93 ///
94 /// Returns
95 /// - `Ok(Some(StompItem))` when a full item (frame or heartbeat) was
96 /// decoded and bytes were consumed from `src` accordingly.
97 /// - `Ok(None)` when more bytes are required to decode a complete item.
98 /// - `Err(io::Error)` on protocol or data errors (invalid UTF-8, malformed
99 /// frames, missing NUL after a content-length body, etc.).
100 fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
101 // Heartbeat: a STOMP 1.2 EOL, which is a bare LF or a CRLF. We parse
102 // directly from `src` (BytesMut is contiguous, so `chunk()` is the whole
103 // buffer); there is no separate accumulation buffer.
104 match src.chunk().first() {
105 Some(b'\n') => {
106 src.advance(1);
107 return Ok(Some(StompItem::Heartbeat));
108 }
109 Some(b'\r') => match src.chunk().get(1) {
110 Some(b'\n') => {
111 src.advance(2);
112 return Ok(Some(StompItem::Heartbeat));
113 }
114 // A lone CR so far: wait for the next byte to tell CRLF apart
115 // from the (malformed) start of a frame.
116 None => return Ok(None),
117 // CR followed by something else: let the frame parser judge it.
118 Some(_) => {}
119 },
120 _ => {}
121 }
122
123 let chunk = src.chunk();
124 match parse_frame_slice_bounded(chunk, self.max_frame_size) {
125 Ok(Some((cmd_bytes, headers, body, consumed))) => {
126 // Enforce the bound on a *complete* frame too, not only on the
127 // incomplete (`Ok(None)`) path. A whole oversized frame can
128 // arrive in one read — its body bounded inside the parser, but
129 // large headers or framing could still push the total over the
130 // limit — so reject on total wire size here as the authoritative
131 // guard.
132 if consumed > self.max_frame_size {
133 return Err(io::Error::new(
134 io::ErrorKind::InvalidData,
135 format!(
136 "frame size {} exceeds maximum frame size {}",
137 consumed, self.max_frame_size
138 ),
139 ));
140 }
141 // advance src by consumed
142 src.advance(consumed);
143
144 // build owned Frame
145 let command = String::from_utf8(cmd_bytes).map_err(|e| {
146 io::Error::new(
147 io::ErrorKind::InvalidData,
148 format!("invalid utf8 in command: {}", e),
149 )
150 })?;
151 // convert headers Vec<(Vec<u8>,Vec<u8>)> -> Vec<(String,String)>
152 // and unescape per STOMP 1.2 spec
153 let mut hdrs: Vec<(String, String)> = Vec::new();
154 for (k, v) in headers {
155 // Unescape header key
156 let k_unescaped = unescape_header_value(&k).map_err(|e| {
157 io::Error::new(
158 io::ErrorKind::InvalidData,
159 format!("invalid escape in header key: {}", e),
160 )
161 })?;
162 let ks = String::from_utf8(k_unescaped).map_err(|e| {
163 io::Error::new(
164 io::ErrorKind::InvalidData,
165 format!("invalid utf8 in header key: {}", e),
166 )
167 })?;
168 // Unescape header value
169 let v_unescaped = unescape_header_value(&v).map_err(|e| {
170 io::Error::new(
171 io::ErrorKind::InvalidData,
172 format!("invalid escape in header value: {}", e),
173 )
174 })?;
175 let vs = String::from_utf8(v_unescaped).map_err(|e| {
176 io::Error::new(
177 io::ErrorKind::InvalidData,
178 format!("invalid utf8 in header value: {}", e),
179 )
180 })?;
181 hdrs.push((ks, vs));
182 }
183
184 let body = body.unwrap_or_default();
185
186 let frame = Frame {
187 command,
188 headers: hdrs,
189 body,
190 };
191 Ok(Some(StompItem::Frame(frame)))
192 }
193 Ok(None) => {
194 // Incomplete frame. If we have already buffered the maximum and
195 // still cannot parse a whole frame, the frame is oversized (or
196 // never NUL-terminated) — bound it here rather than buffering
197 // without limit. The content-length path is bounded inside the
198 // parser; this covers the NUL-terminated body that never ends.
199 if src.len() > self.max_frame_size {
200 return Err(io::Error::new(
201 io::ErrorKind::InvalidData,
202 format!("frame exceeds maximum frame size {}", self.max_frame_size),
203 ));
204 }
205 Ok(None)
206 }
207 Err(e) => Err(io::Error::new(
208 io::ErrorKind::InvalidData,
209 format!("parse error: {}", e),
210 )),
211 }
212 }
213}
214
215impl Encoder<StompItem> for StompCodec {
216 type Error = io::Error;
217 /// Encode a `StompItem` into the provided destination buffer.
218 ///
219 /// Parameters
220 /// - `item`: the `StompItem` to encode. The encoder takes ownership of the
221 /// item (and any contained `Frame`) and may consume/mutate its contents.
222 /// - `dst`: destination buffer where encoded bytes should be appended.
223 /// This is the same `BytesMut` provided by the `tokio_util::codec`
224 /// framework (e.g. `Framed`). Do not replace or reassign `dst`; instead
225 /// append bytes into it using `BufMut` methods (`put_u8`,
226 /// `put_slice`, `extend_from_slice`, etc.). After `encode` returns the
227 /// contents of `dst` will be written to the underlying transport.
228 ///
229 /// Returns
230 /// - `Ok(())` on success, or `Err(io::Error)` on encoding-related errors.
231 fn encode(&mut self, item: StompItem, dst: &mut BytesMut) -> Result<(), Self::Error> {
232 match item {
233 StompItem::Heartbeat => {
234 dst.put_u8(b'\n');
235 }
236 StompItem::Frame(frame) => {
237 dst.extend_from_slice(frame.command.as_bytes());
238 dst.put_u8(b'\n');
239
240 let mut headers = frame.headers;
241 let has_cl = headers
242 .iter()
243 .any(|(k, _)| k.to_lowercase() == "content-length");
244 if !has_cl {
245 let include_cl =
246 frame.body.contains(&0) || std::str::from_utf8(&frame.body).is_err();
247 if include_cl {
248 headers.push(("content-length".to_string(), frame.body.len().to_string()));
249 }
250 }
251
252 for (k, v) in headers {
253 // Escape header name and value per STOMP 1.2 spec
254 let escaped_key = escape_header_value(&k);
255 let escaped_val = escape_header_value(&v);
256 dst.extend_from_slice(escaped_key.as_bytes());
257 dst.put_u8(b':');
258 dst.extend_from_slice(escaped_val.as_bytes());
259 dst.put_u8(b'\n');
260 }
261
262 dst.put_slice(b"\n");
263 dst.extend_from_slice(&frame.body);
264 dst.put_u8(0);
265 }
266 }
267
268 Ok(())
269 }
270}