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

powdb_server/
protocol.rs
