krabka_kafka_tap/
frame.rs

1//! Kafka wire framing, and request and response correlation. This module has
2//! no schema knowledge.
3
4use std::{collections::HashMap, fmt::Write as _};
5
6use krabka_ids::{ApiKey, ApiVersion};
7
8#[derive(Debug, PartialEq, Eq, Clone, Copy)]
9pub struct RequestPrefix {
10    pub api_key: ApiKey,
11    pub api_version: ApiVersion,
12    pub correlation_id: i32,
13}
14
15/// Parse the fixed request-header prefix that every request type has:
16/// `api_key: i16, api_version: i16, correlation_id: i32`.
17#[must_use]
18pub fn parse_request_prefix(body: &[u8]) -> Option<RequestPrefix> {
19    let [k0, k1, v0, v1, c0, c1, c2, c3, ..] = *body else {
20        return None;
21    };
22    Some(RequestPrefix {
23        api_key: ApiKey(i16::from_be_bytes([k0, k1])),
24        api_version: ApiVersion(i16::from_be_bytes([v0, v1])),
25        correlation_id: i32::from_be_bytes([c0, c1, c2, c3]),
26    })
27}
28
29/// Every response body starts with `correlation_id: i32`, before any tagged
30/// header. This is true for flexible and non-flexible responses alike.
31#[must_use]
32pub fn read_correlation_id(body: &[u8]) -> Option<i32> {
33    body.first_chunk().copied().map(i32::from_be_bytes)
34}
35
36/// Per-connection map from a correlation id to the (`api_key`, `api_version`) of
37/// the request that carried it. The relay uses the map to classify the
38/// responses.
39#[derive(Default)]
40pub struct Pending {
41    map: HashMap<i32, (ApiKey, ApiVersion)>,
42}
43
44impl Pending {
45    pub fn record(&mut self, correlation_id: i32, api_key: ApiKey, api_version: ApiVersion) {
46        self.map.insert(correlation_id, (api_key, api_version));
47    }
48    #[must_use]
49    pub fn take(&mut self, correlation_id: i32) -> Option<(ApiKey, ApiVersion)> {
50        self.map.remove(&correlation_id)
51    }
52}
53
54/// One captured frame, emitted by the relay to the recorder spool.
55#[derive(Debug, Clone, PartialEq, Eq)]
56pub struct CapturedFrame {
57    pub api_key: ApiKey,
58    pub version: ApiVersion,
59    pub is_request: bool,
60    /// Full frame body, that is the header and the message, without the
61    /// 4-byte length prefix.
62    pub body: Vec<u8>,
63}
64
65impl CapturedFrame {
66    /// The frame as one newline-terminated JSON record of the standalone
67    /// tap's spool, with the body in lowercase hex.
68    #[must_use]
69    pub fn spool_line(&self) -> String {
70        let mut hex = String::with_capacity(self.body.len() * 2);
71        for b in &self.body {
72            let _ = write!(hex, "{b:02x}");
73        }
74        format!(
75            "{{\"api_key\":{},\"version\":{},\"is_request\":{},\"body_hex\":\"{}\"}}\n",
76            self.api_key, self.version, self.is_request, hex
77        )
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84
85    fn req_frame(api_key: i16, api_version: i16, corr: i32, payload: &[u8]) -> Vec<u8> {
86        let mut b = Vec::new();
87        b.extend_from_slice(&api_key.to_be_bytes());
88        b.extend_from_slice(&api_version.to_be_bytes());
89        b.extend_from_slice(&corr.to_be_bytes());
90        b.extend_from_slice(payload);
91        b
92    }
93
94    #[test]
95    fn parses_request_header_prefix() {
96        let body = req_frame(18, 3, 7, &[0xaa, 0xbb]);
97        let p = parse_request_prefix(&body).unwrap();
98        assert2::assert!(
99            p == RequestPrefix {
100                api_key: ApiKey(18),
101                api_version: ApiVersion(3),
102                correlation_id: 7
103            }
104        );
105    }
106
107    /// `read_correlation_id` needs exactly four bytes, and its guard is the one
108    /// place that decides. A body of exactly four is the shortest legal one and
109    /// must parse; three must be refused rather than indexed into. Together
110    /// those separate `<` from `<=` and from `==`.
111    #[test]
112    fn correlation_id_needs_exactly_four_bytes() {
113        assert2::check!(read_correlation_id(&7_i32.to_be_bytes()) == Some(7));
114        assert2::check!(read_correlation_id(&[0, 0, 0]) == None);
115        assert2::check!(read_correlation_id(&[]) == None);
116        // A longer body is a real response: the id is its first four bytes and
117        // the rest is payload.
118        assert2::check!(read_correlation_id(&[0, 0, 0, 9, 0xaa, 0xbb]) == Some(9));
119    }
120
121    /// The length guard is a strict `<`, and four bytes is exactly enough. Read
122    /// as `<=` that body is refused; read as `==` it is refused while a
123    /// three-byte body falls through to index past its end.
124    #[test]
125    fn correlation_id_needs_four_bytes_and_four_is_enough() {
126        assert2::check!(read_correlation_id(&7i32.to_be_bytes()) == Some(7));
127        assert2::check!(read_correlation_id(&[0, 0, 0, 1, 0xff]) == Some(1));
128        assert2::check!(read_correlation_id(&[0, 0, 0]) == None);
129        assert2::check!(read_correlation_id(&[]) == None);
130    }
131
132    #[test]
133    fn correlates_response_by_id() {
134        let mut pending = Pending::default();
135        let body = req_frame(1, 11, 42, &[]);
136        let p = parse_request_prefix(&body).unwrap();
137        pending.record(p.correlation_id, p.api_key, p.api_version);
138        let mut resp = Vec::new();
139        resp.extend_from_slice(&42i32.to_be_bytes());
140        resp.extend_from_slice(&[0x01]);
141        let got = pending.take(read_correlation_id(&resp).unwrap()).unwrap();
142        assert2::assert!(got == (ApiKey(1), ApiVersion(11)));
143    }
144}
krabka_kafka_tap/frame.rs

krabka_kafka_tap/
frame.rs
