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

powdb_server/
protocol.rs
