powdb_server/
protocol.rs

1use tokio::io::{AsyncReadExt, AsyncWriteExt};
2
3const MSG_CONNECT: u8 = 0x01;
4const MSG_CONNECT_OK: u8 = 0x02;
5const MSG_QUERY: u8 = 0x03;
6const MSG_RESULT_ROWS: u8 = 0x07;
7const MSG_RESULT_SCALAR: u8 = 0x08;
8const MSG_RESULT_OK: u8 = 0x09;
9const MSG_ERROR: u8 = 0x0A;
10const MSG_DISCONNECT: u8 = 0x10;
11const MSG_PING: u8 = 0x11;
12const MSG_PONG: u8 = 0x12;
13
14/// Maximum payload size accepted from the wire (64 MB).
15const MAX_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
16
17/// Maximum payload size for pre-auth CONNECT messages (4 KB).
18/// Only a database name and password are needed before authentication.
19const MAX_CONNECT_PAYLOAD_SIZE: usize = 4096;
20
21/// Maximum number of columns allowed in a result set.
22const MAX_COLUMNS: usize = 4096;
23
24/// Maximum number of rows allowed in a single result message.
25const MAX_ROWS: usize = 10_000_000;
26
27#[derive(Debug, Clone)]
28pub enum Message {
29    Connect {
30        db_name: String,
31        password: Option<String>,
32    },
33    ConnectOk {
34        version: String,
35    },
36    Query {
37        query: String,
38    },
39    ResultRows {
40        columns: Vec<String>,
41        rows: Vec<Vec<String>>,
42    },
43    ResultScalar {
44        value: String,
45    },
46    ResultOk {
47        affected: u64,
48    },
49    Error {
50        message: String,
51    },
52    Disconnect,
53    Ping,
54    Pong,
55}
56
57impl Message {
58    /// Encode message into wire format: [type(1)][flags(1)][len(4)][payload]
59    pub fn encode(&self) -> Vec<u8> {
60        let (msg_type, payload) = match self {
61            Message::Connect { db_name, password } => {
62                let mut buf = encode_string(db_name);
63                // Password is encoded as a length-prefixed string. Empty (len=0) means None.
64                match password {
65                    Some(p) => buf.extend_from_slice(&encode_string(p)),
66                    None => buf.extend_from_slice(&0u32.to_le_bytes()),
67                }
68                (MSG_CONNECT, buf)
69            }
70            Message::ConnectOk { version } => (MSG_CONNECT_OK, encode_string(version)),
71            Message::Query { query } => (MSG_QUERY, encode_string(query)),
72            Message::ResultRows { columns, rows } => {
73                let mut buf = Vec::new();
74                buf.extend_from_slice(&(columns.len() as u16).to_le_bytes());
75                for col in columns {
76                    buf.extend_from_slice(&encode_string(col));
77                }
78                buf.extend_from_slice(&(rows.len() as u32).to_le_bytes());
79                for row in rows {
80                    for val in row {
81                        buf.extend_from_slice(&encode_string(val));
82                    }
83                }
84                (MSG_RESULT_ROWS, buf)
85            }
86            Message::ResultScalar { value } => (MSG_RESULT_SCALAR, encode_string(value)),
87            Message::ResultOk { affected } => (MSG_RESULT_OK, affected.to_le_bytes().to_vec()),
88            Message::Error { message } => (MSG_ERROR, encode_string(message)),
89            Message::Disconnect => (MSG_DISCONNECT, Vec::new()),
90            Message::Ping => (MSG_PING, Vec::new()),
91            Message::Pong => (MSG_PONG, Vec::new()),
92        };
93
94        let mut frame = Vec::with_capacity(6 + payload.len());
95        frame.push(msg_type);
96        frame.push(0); // flags
97        frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
98        frame.extend_from_slice(&payload);
99        frame
100    }
101
102    /// Decode message from wire format.
103    pub fn decode(data: &[u8]) -> Result<Message, String> {
104        if data.len() < 6 {
105            return Err("frame too short".into());
106        }
107        let msg_type = data[0];
108        let _flags = data[1];
109        let len_bytes: [u8; 4] = data[2..6]
110            .try_into()
111            .map_err(|_| "invalid header length field".to_string())?;
112        let payload_len = u32::from_le_bytes(len_bytes) as usize;
113        if 6 + payload_len > data.len() {
114            return Err("payload length exceeds frame".into());
115        }
116        let payload = &data[6..6 + payload_len];
117
118        match msg_type {
119            MSG_CONNECT => {
120                let mut pos = 0;
121                let db_name = decode_string(payload, &mut pos)?;
122                // Password is optional. If there are no more bytes, treat as None
123                // (backwards compatible with old clients that don't send a password).
124                let password = if pos < payload.len() {
125                    let p = decode_string(payload, &mut pos)?;
126                    if p.is_empty() {
127                        None
128                    } else {
129                        Some(p)
130                    }
131                } else {
132                    None
133                };
134                Ok(Message::Connect { db_name, password })
135            }
136            MSG_CONNECT_OK => {
137                let version = decode_string(payload, &mut 0)?;
138                Ok(Message::ConnectOk { version })
139            }
140            MSG_QUERY => {
141                let query = decode_string(payload, &mut 0)?;
142                Ok(Message::Query { query })
143            }
144            MSG_RESULT_ROWS => {
145                let mut pos = 0;
146                if pos + 2 > payload.len() {
147                    return Err("truncated column count".into());
148                }
149                let col_bytes: [u8; 2] = payload[pos..pos + 2]
150                    .try_into()
151                    .map_err(|_| "invalid column count bytes".to_string())?;
152                let col_count = u16::from_le_bytes(col_bytes) as usize;
153                pos += 2;
154                if col_count > MAX_COLUMNS {
155                    return Err("too many columns".into());
156                }
157                let mut columns = Vec::with_capacity(col_count);
158                for _ in 0..col_count {
159                    columns.push(decode_string(payload, &mut pos)?);
160                }
161                if pos + 4 > payload.len() {
162                    return Err("truncated row count".into());
163                }
164                let row_bytes: [u8; 4] = payload[pos..pos + 4]
165                    .try_into()
166                    .map_err(|_| "invalid row count bytes".to_string())?;
167                let row_count = u32::from_le_bytes(row_bytes) as usize;
168                pos += 4;
169                if row_count > MAX_ROWS {
170                    return Err("too many rows".into());
171                }
172                let mut rows = Vec::with_capacity(row_count);
173                for _ in 0..row_count {
174                    let mut row = Vec::with_capacity(col_count);
175                    for _ in 0..col_count {
176                        row.push(decode_string(payload, &mut pos)?);
177                    }
178                    rows.push(row);
179                }
180                Ok(Message::ResultRows { columns, rows })
181            }
182            MSG_RESULT_SCALAR => {
183                let value = decode_string(payload, &mut 0)?;
184                Ok(Message::ResultScalar { value })
185            }
186            MSG_RESULT_OK => {
187                if payload.len() < 8 {
188                    return Err("truncated result ok payload".into());
189                }
190                let aff_bytes: [u8; 8] = payload[0..8]
191                    .try_into()
192                    .map_err(|_| "invalid affected count bytes".to_string())?;
193                let affected = u64::from_le_bytes(aff_bytes);
194                Ok(Message::ResultOk { affected })
195            }
196            MSG_ERROR => {
197                let message = decode_string(payload, &mut 0)?;
198                Ok(Message::Error { message })
199            }
200            MSG_DISCONNECT => Ok(Message::Disconnect),
201            MSG_PING => Ok(Message::Ping),
202            MSG_PONG => Ok(Message::Pong),
203            _ => Err(format!("unknown message type: {msg_type:#x}")),
204        }
205    }
206
207    /// Write this message to an async writer.
208    pub async fn write_to<W: AsyncWriteExt + Unpin>(&self, writer: &mut W) -> std::io::Result<()> {
209        let bytes = self.encode();
210        writer.write_all(&bytes).await
211    }
212
213    /// Read a message from an async reader.
214    pub async fn read_from<R: AsyncReadExt + Unpin>(
215        reader: &mut R,
216    ) -> std::io::Result<Option<Message>> {
217        Self::read_from_with_limit(reader, MAX_PAYLOAD_SIZE).await
218    }
219
220    /// Read a pre-auth message with a smaller payload limit (4 KB).
221    /// Use this before authentication is complete to prevent oversized
222    /// CONNECT payloads from consuming server memory.
223    pub async fn read_from_preauth<R: AsyncReadExt + Unpin>(
224        reader: &mut R,
225    ) -> std::io::Result<Option<Message>> {
226        Self::read_from_with_limit(reader, MAX_CONNECT_PAYLOAD_SIZE).await
227    }
228
229    /// Read a message from an async reader with a configurable payload limit.
230    async fn read_from_with_limit<R: AsyncReadExt + Unpin>(
231        reader: &mut R,
232        max_payload: usize,
233    ) -> std::io::Result<Option<Message>> {
234        let mut header = [0u8; 6];
235        match reader.read_exact(&mut header).await {
236            Ok(_) => {}
237            Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
238            Err(e) => return Err(e),
239        }
240        let len_bytes: [u8; 4] = header[2..6].try_into().map_err(|_| {
241            std::io::Error::new(
242                std::io::ErrorKind::InvalidData,
243                "invalid header length field",
244            )
245        })?;
246        let payload_len = u32::from_le_bytes(len_bytes) as usize;
247        if payload_len > max_payload {
248            return Err(std::io::Error::new(
249                std::io::ErrorKind::InvalidData,
250                format!("payload too large: {payload_len} bytes (max {max_payload})"),
251            ));
252        }
253        let mut payload = vec![0u8; payload_len];
254        if payload_len > 0 {
255            reader.read_exact(&mut payload).await?;
256        }
257
258        let mut full = Vec::with_capacity(6 + payload_len);
259        full.extend_from_slice(&header);
260        full.extend_from_slice(&payload);
261
262        Message::decode(&full)
263            .map(Some)
264            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
265    }
266}
267
268fn encode_string(s: &str) -> Vec<u8> {
269    let mut buf = Vec::with_capacity(4 + s.len());
270    buf.extend_from_slice(&(s.len() as u32).to_le_bytes());
271    buf.extend_from_slice(s.as_bytes());
272    buf
273}
274
275fn decode_string(data: &[u8], pos: &mut usize) -> Result<String, String> {
276    if *pos + 4 > data.len() {
277        return Err("truncated string length".into());
278    }
279    let len_bytes: [u8; 4] = data[*pos..*pos + 4]
280        .try_into()
281        .map_err(|_| "invalid string length bytes".to_string())?;
282    let len = u32::from_le_bytes(len_bytes) as usize;
283    *pos += 4;
284    if *pos + len > data.len() {
285        return Err("truncated string data".into());
286    }
287    let s = String::from_utf8_lossy(&data[*pos..*pos + len]).into_owned();
288    *pos += len;
289    Ok(s)
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295
296    #[test]
297    fn test_encode_decode_query() {
298        let msg = Message::Query {
299            query: "User filter .age > 30".into(),
300        };
301        let bytes = msg.encode();
302        let decoded = Message::decode(&bytes).unwrap();
303        match decoded {
304            Message::Query { query } => assert_eq!(query, "User filter .age > 30"),
305            _ => panic!("expected Query"),
306        }
307    }
308
309    #[test]
310    fn test_encode_decode_result_rows() {
311        let msg = Message::ResultRows {
312            columns: vec!["name".into(), "age".into()],
313            rows: vec![
314                vec!["Alice".into(), "30".into()],
315                vec!["Bob".into(), "25".into()],
316            ],
317        };
318        let bytes = msg.encode();
319        let decoded = Message::decode(&bytes).unwrap();
320        match decoded {
321            Message::ResultRows { columns, rows } => {
322                assert_eq!(columns, vec!["name", "age"]);
323                assert_eq!(rows.len(), 2);
324            }
325            _ => panic!("expected ResultRows"),
326        }
327    }
328
329    #[test]
330    fn test_encode_decode_error() {
331        let msg = Message::Error {
332            message: "table not found".into(),
333        };
334        let bytes = msg.encode();
335        let decoded = Message::decode(&bytes).unwrap();
336        match decoded {
337            Message::Error { message } => assert_eq!(message, "table not found"),
338            _ => panic!("expected Error"),
339        }
340    }
341
342    #[test]
343    fn test_frame_length() {
344        let msg = Message::Query {
345            query: "User".into(),
346        };
347        let bytes = msg.encode();
348        assert!(bytes.len() >= 6);
349        let payload_len = u32::from_le_bytes(bytes[2..6].try_into().unwrap()) as usize;
350        assert_eq!(bytes.len(), 6 + payload_len);
351    }
352}
powdb_server/protocol.rs

powdb_server/
protocol.rs
