Skip to main content

krabka_protocol/primitives/
string_bytes.rs

1use bytes::{Buf, BufMut, Bytes};
2
3use crate::{
4    ProtocolError,
5    primitives::{
6        fixed::{get_i16, get_i32, put_i16, put_i32},
7        varint::{compact_length, get_uvarint, put_uvarint, usize_uvarint_len, uvarint_len},
8    },
9};
10
11// ---- STRING (non-flexible) ----
12// Wire: INT16 length (>=0), then `length` bytes UTF-8. -1 = null.
13
14/// # Errors
15/// Returns [`ProtocolError::LengthOverflow`] when the length does not fit its
16/// prefix.
17pub fn put_string<B: BufMut>(buf: &mut B, s: &str) -> Result<(), ProtocolError> {
18    let len = i16::try_from(s.len()).map_err(|_| ProtocolError::LengthOverflow {
19        what: "STRING",
20        len: s.len(),
21    })?;
22    put_i16(buf, len);
23    buf.put_slice(s.as_bytes());
24    Ok(())
25}
26
27/// # Errors
28/// Returns [`ProtocolError::LengthOverflow`] when the length does not fit its
29/// prefix.
30pub fn put_nullable_string<B: BufMut>(buf: &mut B, s: Option<&str>) -> Result<(), ProtocolError> {
31    match s {
32        None => {
33            put_i16(buf, -1);
34            Ok(())
35        }
36        Some(s) => put_string(buf, s),
37    }
38}
39
40/// # Errors
41/// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
42pub fn get_string_owned<B: Buf>(buf: &mut B) -> Result<String, ProtocolError> {
43    get_nullable_string_owned(buf)?
44        .ok_or(ProtocolError::InvalidValue("non-nullable STRING was null"))
45}
46
47/// # Errors
48/// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
49pub fn get_nullable_string_owned<B: Buf>(buf: &mut B) -> Result<Option<String>, ProtocolError> {
50    // A negative length is null.
51    let Ok(n) = usize::try_from(get_i16(buf)?) else {
52        return Ok(None);
53    };
54    if buf.remaining() < n {
55        return Err(ProtocolError::UnexpectedEof {
56            needed: n - buf.remaining(),
57        });
58    }
59    let mut v = vec![0u8; n];
60    buf.copy_to_slice(&mut v);
61    let s = String::from_utf8(v).map_err(|e| ProtocolError::InvalidUtf8(e.utf8_error()))?;
62    Ok(Some(s))
63}
64
65#[must_use]
66pub fn string_len(s: &str) -> usize {
67    2 + s.len()
68}
69#[must_use]
70pub fn nullable_string_len(s: Option<&str>) -> usize {
71    2 + s.map_or(0, str::len)
72}
73
74// ---- COMPACT_STRING (flexible) ----
75// Wire: UVARINT length+1 (0 = null), then `length` UTF-8 bytes.
76
77/// # Errors
78/// Returns [`ProtocolError::LengthOverflow`] when the length does not fit its
79/// prefix.
80pub fn put_compact_string<B: BufMut>(buf: &mut B, s: &str) -> Result<(), ProtocolError> {
81    put_uvarint(buf, compact_length(s.len(), "COMPACT_STRING")?);
82    buf.put_slice(s.as_bytes());
83    Ok(())
84}
85
86/// # Errors
87/// Returns [`ProtocolError::LengthOverflow`] when the length does not fit its
88/// prefix.
89pub fn put_compact_nullable_string<B: BufMut>(
90    buf: &mut B,
91    s: Option<&str>,
92) -> Result<(), ProtocolError> {
93    match s {
94        None => {
95            put_uvarint(buf, 0);
96            Ok(())
97        }
98        Some(s) => put_compact_string(buf, s),
99    }
100}
101
102/// # Errors
103/// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
104pub fn get_compact_string_owned<B: Buf>(buf: &mut B) -> Result<String, ProtocolError> {
105    get_compact_nullable_string_owned(buf)?.ok_or(ProtocolError::InvalidValue(
106        "non-nullable COMPACT_STRING was null",
107    ))
108}
109
110/// # Errors
111/// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
112pub fn get_compact_nullable_string_owned<B: Buf>(
113    buf: &mut B,
114) -> Result<Option<String>, ProtocolError> {
115    let raw = get_uvarint(buf)?;
116    if raw == 0 {
117        return Ok(None);
118    }
119    let n = (raw - 1) as usize;
120    if buf.remaining() < n {
121        return Err(ProtocolError::UnexpectedEof {
122            needed: n - buf.remaining(),
123        });
124    }
125    let mut v = vec![0u8; n];
126    buf.copy_to_slice(&mut v);
127    let s = String::from_utf8(v).map_err(|e| ProtocolError::InvalidUtf8(e.utf8_error()))?;
128    Ok(Some(s))
129}
130
131#[must_use]
132pub fn compact_string_len(s: &str) -> usize {
133    usize_uvarint_len(s.len().saturating_add(1)).saturating_add(s.len())
134}
135#[must_use]
136pub fn compact_nullable_string_len(s: Option<&str>) -> usize {
137    s.map_or_else(|| uvarint_len(0), compact_string_len)
138}
139
140// ---- BYTES / COMPACT_BYTES ----
141// BYTES: INT32 length, `length` bytes. -1 = null.
142// COMPACT_BYTES: UVARINT length+1 (0=null), `length` bytes.
143
144/// # Errors
145/// Returns [`ProtocolError::LengthOverflow`] when the length does not fit its
146/// prefix.
147pub fn put_bytes<B: BufMut>(buf: &mut B, b: &[u8]) -> Result<(), ProtocolError> {
148    let len = i32::try_from(b.len()).map_err(|_| ProtocolError::LengthOverflow {
149        what: "BYTES",
150        len: b.len(),
151    })?;
152    put_i32(buf, len);
153    buf.put_slice(b);
154    Ok(())
155}
156
157/// # Errors
158/// Returns [`ProtocolError::LengthOverflow`] when the length does not fit its
159/// prefix.
160pub fn put_nullable_bytes<B: BufMut>(buf: &mut B, b: Option<&[u8]>) -> Result<(), ProtocolError> {
161    match b {
162        None => {
163            put_i32(buf, -1);
164            Ok(())
165        }
166        Some(b) => put_bytes(buf, b),
167    }
168}
169
170/// # Errors
171/// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
172pub fn get_bytes_owned<B: Buf>(buf: &mut B) -> Result<Bytes, ProtocolError> {
173    get_nullable_bytes_owned(buf)?.ok_or(ProtocolError::InvalidValue("non-nullable BYTES was null"))
174}
175
176/// # Errors
177/// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
178pub fn get_nullable_bytes_owned<B: Buf>(buf: &mut B) -> Result<Option<Bytes>, ProtocolError> {
179    // A negative length is null.
180    let Ok(n) = usize::try_from(get_i32(buf)?) else {
181        return Ok(None);
182    };
183    if buf.remaining() < n {
184        return Err(ProtocolError::UnexpectedEof {
185            needed: n - buf.remaining(),
186        });
187    }
188    let mut v = vec![0u8; n];
189    buf.copy_to_slice(&mut v);
190    Ok(Some(Bytes::from(v)))
191}
192
193#[must_use]
194pub fn bytes_len(b: &[u8]) -> usize {
195    4 + b.len()
196}
197#[must_use]
198pub fn nullable_bytes_len(b: Option<&[u8]>) -> usize {
199    4 + b.map_or(0, <[u8]>::len)
200}
201
202/// # Errors
203/// Returns [`ProtocolError::LengthOverflow`] when the length does not fit its
204/// prefix.
205pub fn put_compact_bytes<B: BufMut>(buf: &mut B, b: &[u8]) -> Result<(), ProtocolError> {
206    put_uvarint(buf, compact_length(b.len(), "COMPACT_BYTES")?);
207    buf.put_slice(b);
208    Ok(())
209}
210
211/// # Errors
212/// Returns [`ProtocolError::LengthOverflow`] when the length does not fit its
213/// prefix.
214pub fn put_compact_nullable_bytes<B: BufMut>(
215    buf: &mut B,
216    b: Option<&[u8]>,
217) -> Result<(), ProtocolError> {
218    match b {
219        None => {
220            put_uvarint(buf, 0);
221            Ok(())
222        }
223        Some(b) => put_compact_bytes(buf, b),
224    }
225}
226
227#[must_use]
228pub fn compact_bytes_len(b: &[u8]) -> usize {
229    usize_uvarint_len(b.len().saturating_add(1)).saturating_add(b.len())
230}
231
232/// Like `compact_bytes_len`, but takes the byte count directly rather than a slice.
233///
234/// Use this when you know the content size and do not need to materialise the
235/// buffer.
236#[must_use]
237pub fn compact_bytes_len_from_size(n: usize) -> usize {
238    usize_uvarint_len(n.saturating_add(1)).saturating_add(n)
239}
240#[must_use]
241pub fn compact_nullable_bytes_len(b: Option<&[u8]>) -> usize {
242    b.map_or_else(|| uvarint_len(0), compact_bytes_len)
243}
244
245/// # Errors
246/// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
247pub fn get_compact_bytes_owned<B: Buf>(buf: &mut B) -> Result<Bytes, ProtocolError> {
248    get_compact_nullable_bytes_owned(buf)?.ok_or(ProtocolError::InvalidValue(
249        "non-nullable COMPACT_BYTES was null",
250    ))
251}
252
253/// # Errors
254/// Returns the underlying protocol error when input is truncated, contains an invalid length or tag, or cannot be encoded for the selected version.
255pub fn get_compact_nullable_bytes_owned<B: Buf>(
256    buf: &mut B,
257) -> Result<Option<Bytes>, ProtocolError> {
258    let raw = get_uvarint(buf)?;
259    if raw == 0 {
260        return Ok(None);
261    }
262    let n = (raw - 1) as usize;
263    if buf.remaining() < n {
264        return Err(ProtocolError::UnexpectedEof {
265            needed: n - buf.remaining(),
266        });
267    }
268    let mut v = vec![0u8; n];
269    buf.copy_to_slice(&mut v);
270    Ok(Some(Bytes::from(v)))
271}
272
273#[cfg(test)]
274mod tests {
275
276    use bytes::BytesMut;
277
278    use super::*;
279
280    #[test]
281    fn string_roundtrip() {
282        let mut buf = BytesMut::new();
283        put_string(&mut buf, "kafka").unwrap();
284        // INT16(5) + bytes
285        assert2::assert!(&buf[..] == &[0x00, 0x05, b'k', b'a', b'f', b'k', b'a']);
286        let mut cur = &buf[..];
287        assert2::assert!(get_string_owned(&mut cur).unwrap() == "kafka");
288    }
289
290    #[test]
291    fn nullable_string_null() {
292        let mut buf = BytesMut::new();
293        put_nullable_string(&mut buf, None).unwrap();
294        assert2::assert!(&buf[..] == &[0xFF, 0xFF]);
295        let mut cur = &buf[..];
296        assert2::assert!(get_nullable_string_owned(&mut cur).unwrap() == None);
297    }
298
299    #[test]
300    fn compact_string_roundtrip() {
301        let mut buf = BytesMut::new();
302        put_compact_string(&mut buf, "kafka").unwrap();
303        // UVARINT(6) + bytes
304        assert2::assert!(&buf[..] == &[0x06, b'k', b'a', b'f', b'k', b'a']);
305        let mut cur = &buf[..];
306        assert2::assert!(get_compact_string_owned(&mut cur).unwrap() == "kafka");
307    }
308
309    #[test]
310    fn compact_nullable_string_null() {
311        let mut buf = BytesMut::new();
312        put_compact_nullable_string(&mut buf, None).unwrap();
313        assert2::assert!(&buf[..] == &[0x00]);
314        let mut cur = &buf[..];
315        assert2::assert!(get_compact_nullable_string_owned(&mut cur).unwrap() == None);
316    }
317
318    #[test]
319    fn empty_compact_string() {
320        let mut buf = BytesMut::new();
321        put_compact_string(&mut buf, "").unwrap();
322        assert2::assert!(&buf[..] == &[0x01]); // length = 1 means "0 bytes"
323        let mut cur = &buf[..];
324        assert2::assert!(get_compact_string_owned(&mut cur).unwrap() == "");
325    }
326
327    #[test]
328    fn bytes_roundtrip() {
329        let mut buf = BytesMut::new();
330        put_bytes(&mut buf, &[1, 2, 3]).unwrap();
331        let mut cur = &buf[..];
332        let out = get_bytes_owned(&mut cur).unwrap();
333        assert2::assert!(out.as_ref() == &[1, 2, 3]);
334    }
335
336    #[test]
337    fn invalid_utf8_is_rejected() {
338        // INT16(2) + invalid UTF-8 byte sequence
339        let bytes = [0x00, 0x02, 0xC3, 0x28];
340        let mut cur = &bytes[..];
341        assert2::assert!(matches!(
342            get_string_owned(&mut cur),
343            Err(ProtocolError::InvalidUtf8(_))
344        ));
345    }
346}