aerostream_client/
protocol.rs

1use bytes::{BufMut, Bytes, BytesMut};
2use crate::error::AeroError;
3
4/// AeroStream native protocol magic bytes: [0xAE, 0x01].
5pub const MAGIC: [u8; 2] = [0xAE, 0x01];
6
7/// Protocol header length in bytes: magic (2B) + command (1B) + body_len (4B).
8pub const HEADER_LEN: usize = 7;
9
10/// Protocol status code for successful operation.
11pub const STATUS_OK: u8 = 0;
12
13/// Protocol status code indicating no data available (empty response).
14pub const STATUS_EMPTY: u8 = 1;
15
16/// Protocol status code indicating data payload follows.
17pub const STATUS_DATA: u8 = 2;
18
19/// Protocol status code indicating authentication failure.
20pub const STATUS_AUTH_FAILED: u8 = 3;
21
22/// Protocol status code indicating out-of-order sequence number.
23pub const STATUS_OUT_OF_ORDER: u8 = 45;
24
25/// Native protocol command identifiers.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27#[repr(u8)]
28pub enum Command {
29    /// Command 0: Authentication handshake.
30    Auth = 0,
31    /// Command 1: Produce/Append message.
32    Produce = 1,
33    /// Command 2: Consumer Fetch (single/range).
34    Fetch = 2,
35    /// Command 3: Replica Fetch/Sync.
36    ReplicaFetch = 3,
37    /// Command 4: Multi-Entry Long-Polling Fetch.
38    FetchMulti = 4,
39}
40
41impl Command {
42    /// Convert byte to Command enum variant.
43    pub fn from_u8(b: u8) -> Result<Self, AeroError> {
44        match b {
45            0 => Ok(Command::Auth),
46            1 => Ok(Command::Produce),
47            2 => Ok(Command::Fetch),
48            3 => Ok(Command::ReplicaFetch),
49            4 => Ok(Command::FetchMulti),
50            other => Err(AeroError::ProtocolError(format!("Unknown command ID: {other}"))),
51        }
52    }
53}
54
55/// A consumed record with its log offset and payload bytes.
56#[derive(Debug, Clone, PartialEq, Eq)]
57pub struct Record {
58    /// Commit log offset of the message.
59    pub offset: u64,
60    /// Raw message payload.
61    pub payload: Bytes,
62}
63
64impl Record {
65    /// Construct a new Record.
66    pub fn new(offset: u64, payload: Bytes) -> Self {
67        Self { offset, payload }
68    }
69}
70
71/// Validate 2-byte protocol magic prefix.
72#[inline]
73pub fn is_valid_magic(b0: u8, b1: u8) -> bool {
74    b0 == MAGIC[0] && b1 == MAGIC[1]
75}
76
77/// Encode a standard 7-byte request header into a pre-allocated buffer.
78pub fn encode_header(cmd: Command, body_len: u32, buf: &mut BytesMut) {
79    buf.put_u8(MAGIC[0]);
80    buf.put_u8(MAGIC[1]);
81    buf.put_u8(cmd as u8);
82    buf.put_u32(body_len);
83}
84
85/// Encode Command 0 (Auth) request: `[token: UTF-8]`.
86pub fn encode_auth_request(token: &str) -> Bytes {
87    let token_bytes = token.as_bytes();
88    let mut buf = BytesMut::with_capacity(HEADER_LEN + token_bytes.len());
89    encode_header(Command::Auth, token_bytes.len() as u32, &mut buf);
90    buf.put_slice(token_bytes);
91    buf.freeze()
92}
93
94/// Encode Command 1 (Produce) request:
95/// `[topic_len: u16 BE][topic: UTF-8][partition: u32 BE][payload_len: u32 BE][payload: raw bytes]`.
96pub fn encode_produce_request(topic: &str, partition: u32, payload: &[u8]) -> Result<Bytes, AeroError> {
97    let topic_bytes = topic.as_bytes();
98    if topic_bytes.len() > u16::MAX as usize {
99        return Err(AeroError::ProtocolError(format!(
100            "Topic length {} exceeds u16::MAX",
101            topic_bytes.len()
102        )));
103    }
104    let body_len = 2 + topic_bytes.len() + 4 + 4 + payload.len();
105    if body_len > u32::MAX as usize {
106        return Err(AeroError::ProtocolError("Produce body length exceeds u32::MAX".into()));
107    }
108
109    let mut buf = BytesMut::with_capacity(HEADER_LEN + body_len);
110    encode_header(Command::Produce, body_len as u32, &mut buf);
111    buf.put_u16(topic_bytes.len() as u16);
112    buf.put_slice(topic_bytes);
113    buf.put_u32(partition);
114    buf.put_u32(payload.len() as u32);
115    buf.put_slice(payload);
116    Ok(buf.freeze())
117}
118
119/// Encode Command 2 (Fetch) request:
120/// `[topic_len: u16 BE][topic: UTF-8][partition: u32 BE][start_offset: u64 BE][max_bytes: u32 BE]`.
121pub fn encode_fetch_request(
122    topic: &str,
123    partition: u32,
124    start_offset: u64,
125    max_bytes: u32,
126) -> Result<Bytes, AeroError> {
127    let topic_bytes = topic.as_bytes();
128    if topic_bytes.len() > u16::MAX as usize {
129        return Err(AeroError::ProtocolError(format!(
130            "Topic length {} exceeds u16::MAX",
131            topic_bytes.len()
132        )));
133    }
134    let body_len = 2 + topic_bytes.len() + 4 + 8 + 4;
135    let mut buf = BytesMut::with_capacity(HEADER_LEN + body_len);
136    encode_header(Command::Fetch, body_len as u32, &mut buf);
137    buf.put_u16(topic_bytes.len() as u16);
138    buf.put_slice(topic_bytes);
139    buf.put_u32(partition);
140    buf.put_u64(start_offset);
141    buf.put_u32(max_bytes);
142    Ok(buf.freeze())
143}
144
145/// Encode Command 4 (FetchMulti) request:
146/// `[topic_len: u16 BE][topic: UTF-8][partition: u32 BE][start_offset: u64 BE][max_bytes: u32 BE][max_wait_ms: u32 BE]`.
147pub fn encode_fetch_multi_request(
148    topic: &str,
149    partition: u32,
150    start_offset: u64,
151    max_bytes: u32,
152    max_wait_ms: u32,
153) -> Result<Bytes, AeroError> {
154    let topic_bytes = topic.as_bytes();
155    if topic_bytes.len() > u16::MAX as usize {
156        return Err(AeroError::ProtocolError(format!(
157            "Topic length {} exceeds u16::MAX",
158            topic_bytes.len()
159        )));
160    }
161    let body_len = 2 + topic_bytes.len() + 4 + 8 + 4 + 4;
162    let mut buf = BytesMut::with_capacity(HEADER_LEN + body_len);
163    encode_header(Command::FetchMulti, body_len as u32, &mut buf);
164    buf.put_u16(topic_bytes.len() as u16);
165    buf.put_slice(topic_bytes);
166    buf.put_u32(partition);
167    buf.put_u64(start_offset);
168    buf.put_u32(max_bytes);
169    buf.put_u32(max_wait_ms);
170    Ok(buf.freeze())
171}
172
173/// Decode response prefix `[0xAE, 0x01, status]`.
174pub fn decode_status_prefix(bytes: &[u8; 3]) -> Result<u8, AeroError> {
175    if !is_valid_magic(bytes[0], bytes[1]) {
176        return Err(AeroError::ProtocolError(format!(
177            "Invalid magic prefix: [0x{:02X}, 0x{:02X}], expected [0xAE, 0x01]",
178            bytes[0], bytes[1]
179        )));
180    }
181    Ok(bytes[2])
182}
183
184/// Decode produce response: parses status and the 8-byte offset if successful.
185pub fn decode_produce_response(status: u8, offset_bytes: Option<&[u8; 8]>) -> Result<u64, AeroError> {
186    match status {
187        STATUS_OK => {
188            let bytes = offset_bytes.ok_or_else(|| {
189                AeroError::ProtocolError("Produce OK status missing 8-byte offset".into())
190            })?;
191            Ok(u64::from_be_bytes(*bytes))
192        }
193        STATUS_AUTH_FAILED => Err(AeroError::AuthFailed("Produce rejected: not authenticated".into())),
194        STATUS_OUT_OF_ORDER => Err(AeroError::OutOfOrder),
195        other => Err(AeroError::ProtocolError(format!("Produce rejected with status: {other}"))),
196    }
197}
198
199/// Decode Multi-Entry Fetch response entries and slice payloads.
200pub fn decode_multi_entries(
201    entry_count: u32,
202    index_data: &[u8],
203    mut payload_bytes: Bytes,
204) -> Result<Vec<Record>, AeroError> {
205    let expected_index_len = (entry_count as usize) * 12;
206    if index_data.len() < expected_index_len {
207        return Err(AeroError::ProtocolError(format!(
208            "Multi-fetch index data too short: got {} bytes, expected {}",
209            index_data.len(),
210            expected_index_len
211        )));
212    }
213
214    let mut records = Vec::with_capacity(entry_count as usize);
215    let mut offset_idx = 0;
216    for _ in 0..entry_count {
217        let off = u64::from_be_bytes(index_data[offset_idx..offset_idx + 8].try_into().unwrap());
218        let len = u32::from_be_bytes(index_data[offset_idx + 8..offset_idx + 12].try_into().unwrap()) as usize;
219        offset_idx += 12;
220
221        if payload_bytes.len() < len {
222            return Err(AeroError::ProtocolError(format!(
223                "Multi-fetch payload truncated: needed {} bytes, remaining {}",
224                len,
225                payload_bytes.len()
226            )));
227        }
228        let record_payload = payload_bytes.split_to(len);
229        records.push(Record::new(off, record_payload));
230    }
231
232    Ok(records)
233}
aerostream_client/protocol.rs

aerostream_client/
protocol.rs
