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

powdb_server/
protocol.rs
