powdb_server/
handler.rs

1use crate::protocol::Message;
2use powdb_query::executor::{is_read_only_statement, Engine, READONLY_NEEDS_WRITE};
3use powdb_query::parser;
4use powdb_query::result::QueryResult;
5use powdb_storage::types::Value;
6use std::sync::{Arc, RwLock};
7use std::time::Duration;
8use tokio::io::{AsyncWriteExt, BufReader, BufWriter};
9use tokio::net::TcpStream;
10use tokio::sync::watch;
11use tracing::{debug, error, info, warn};
12
13/// Constant-time byte comparison to prevent timing side-channel attacks
14/// on password verification. Returns `true` iff `a` and `b` are identical.
15fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
16    if a.len() != b.len() {
17        return false;
18    }
19    let mut diff = 0u8;
20    for (x, y) in a.iter().zip(b.iter()) {
21        diff |= x ^ y;
22    }
23    diff == 0
24}
25
26/// Error messages that are safe to forward to the client verbatim.
27const SAFE_ERROR_PREFIXES: &[&str] = &[
28    "table not found",
29    "column not found",
30    "parse error",
31    "type mismatch",
32    "unknown table",
33    "unknown column",
34    "unknown function",
35    "syntax error",
36    "expected",
37    "unexpected",
38    "missing",
39    "duplicate",
40    "invalid",
41    "cannot",
42    "no such",
43    "already exists",
44];
45
46/// Sanitize an error message before sending it to the client.
47/// Known safe errors are passed through; everything else is replaced
48/// with a generic message to avoid leaking internal details.
49fn sanitize_error(e: &str) -> String {
50    let lower = e.to_lowercase();
51    for prefix in SAFE_ERROR_PREFIXES {
52        if lower.starts_with(prefix) {
53            return e.to_string();
54        }
55    }
56    "query execution error".into()
57}
58
59/// Execute a query against the engine under the RwLock. Read-only
60/// statements acquire `.read()` so concurrent SELECTs can scan in
61/// parallel; mutations acquire `.write()`.
62fn dispatch_query(engine: &Arc<RwLock<Engine>>, query: &str) -> Result<QueryResult, String> {
63    let stmt_result = parser::parse(query).map_err(|e| e.to_string());
64
65    let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
66    if can_try_read {
67        let res = {
68            let eng = engine.read().map_err(|e| format!("lock poisoned: {e}"))?;
69            eng.execute_powql_readonly(query)
70        };
71        match res {
72            Ok(r) => return Ok(r),
73            Err(e) if e == READONLY_NEEDS_WRITE => {
74                // Escalate: fall through to the write path below.
75            }
76            Err(e) => return Err(e),
77        }
78    }
79
80    let mut eng = engine.write().map_err(|e| format!("lock poisoned: {e}"))?;
81    eng.execute_powql(query)
82}
83
84pub async fn handle_connection(
85    stream: TcpStream,
86    engine: Arc<RwLock<Engine>>,
87    expected_password: Option<String>,
88    shutdown_rx: &mut watch::Receiver<bool>,
89    idle_timeout: Duration,
90    query_timeout: Duration,
91) {
92    let peer = stream
93        .peer_addr()
94        .ok()
95        .map(|a| a.to_string())
96        .unwrap_or_else(|| "unknown".into());
97    let (reader, writer) = stream.into_split();
98    let mut reader = BufReader::new(reader);
99    let mut writer = BufWriter::new(writer);
100
101    // Wait for Connect message (with idle timeout).
102    let connect_msg =
103        match tokio::time::timeout(idle_timeout, Message::read_from(&mut reader)).await {
104            Ok(Ok(Some(msg))) => msg,
105            Ok(Ok(None)) => {
106                debug!(peer = %peer, "client closed before CONNECT");
107                return;
108            }
109            Ok(Err(e)) => {
110                error!(peer = %peer, error = %e, "error reading CONNECT");
111                return;
112            }
113            Err(_) => {
114                warn!(peer = %peer, "idle timeout waiting for CONNECT");
115                return;
116            }
117        };
118
119    match connect_msg {
120        Message::Connect { db_name, password } => {
121            if let Some(expected) = &expected_password {
122                if !password
123                    .as_deref()
124                    .is_some_and(|p| constant_time_eq(p.as_bytes(), expected.as_bytes()))
125                {
126                    warn!(peer = %peer, db = %db_name, "auth rejected: bad password");
127                    let err = Message::Error {
128                        message: "authentication failed".into(),
129                    };
130                    err.write_to(&mut writer).await.ok();
131                    writer.flush().await.ok();
132                    return;
133                }
134            }
135            info!(peer = %peer, db = %db_name, "client connected");
136            let ok = Message::ConnectOk {
137                version: "0.1.0".into(),
138            };
139            if ok.write_to(&mut writer).await.is_err() {
140                return;
141            }
142            if writer.flush().await.is_err() {
143                return;
144            }
145        }
146        _ => {
147            warn!(peer = %peer, "first message was not CONNECT");
148            let err = Message::Error {
149                message: "expected CONNECT".into(),
150            };
151            err.write_to(&mut writer).await.ok();
152            writer.flush().await.ok();
153            return;
154        }
155    }
156
157    // Main query loop with idle timeout and shutdown awareness.
158    loop {
159        let msg = tokio::select! {
160            // Read next message with idle timeout.
161            result = tokio::time::timeout(idle_timeout, Message::read_from(&mut reader)) => {
162                match result {
163                    Ok(Ok(Some(msg))) => msg,
164                    Ok(Ok(None)) => break,
165                    Ok(Err(e)) => {
166                        error!(peer = %peer, error = %e, "read error");
167                        break;
168                    }
169                    Err(_) => {
170                        info!(peer = %peer, "idle timeout, closing connection");
171                        let err = Message::Error { message: "idle timeout".into() };
172                        err.write_to(&mut writer).await.ok();
173                        writer.flush().await.ok();
174                        break;
175                    }
176                }
177            }
178            // If server is shutting down, notify client and close.
179            _ = shutdown_rx.changed() => {
180                if *shutdown_rx.borrow() {
181                    info!(peer = %peer, "server shutting down, closing connection");
182                    let err = Message::Error { message: "server shutting down".into() };
183                    err.write_to(&mut writer).await.ok();
184                    writer.flush().await.ok();
185                    break;
186                }
187                continue;
188            }
189        };
190
191        let response = match msg {
192            Message::Query { query } => {
193                debug!(peer = %peer, query = %query, "received query");
194                // Run query with timeout.
195                let result = tokio::task::spawn_blocking({
196                    let engine = engine.clone();
197                    let query = query.clone();
198                    move || dispatch_query(&engine, &query)
199                });
200                match tokio::time::timeout(query_timeout, result).await {
201                    Ok(Ok(Ok(result))) => query_result_to_message(result),
202                    Ok(Ok(Err(e))) => Message::Error {
203                        message: sanitize_error(&e),
204                    },
205                    Ok(Err(e)) => Message::Error {
206                        message: format!("internal error: {e}"),
207                    },
208                    Err(_) => {
209                        warn!(peer = %peer, query = %query, "query timeout exceeded");
210                        Message::Error {
211                            message: "query timeout exceeded".into(),
212                        }
213                    }
214                }
215            }
216            Message::Disconnect => {
217                debug!(peer = %peer, "received DISCONNECT");
218                break;
219            }
220            _ => Message::Error {
221                message: "unexpected message type".into(),
222            },
223        };
224
225        if response.write_to(&mut writer).await.is_err() {
226            break;
227        }
228        if writer.flush().await.is_err() {
229            break;
230        }
231    }
232
233    info!(peer = %peer, "client disconnected");
234}
235
236fn query_result_to_message(result: QueryResult) -> Message {
237    match result {
238        QueryResult::Rows { columns, rows } => {
239            let str_rows: Vec<Vec<String>> = rows
240                .iter()
241                .map(|row| row.iter().map(value_to_display).collect())
242                .collect();
243            Message::ResultRows {
244                columns,
245                rows: str_rows,
246            }
247        }
248        QueryResult::Scalar(val) => Message::ResultScalar {
249            value: value_to_display(&val),
250        },
251        QueryResult::Modified(n) => Message::ResultOk { affected: n },
252        QueryResult::Created(_name) => Message::ResultOk { affected: 0 },
253        QueryResult::Executed { .. } => Message::ResultOk { affected: 0 },
254    }
255}
256
257fn value_to_display(v: &Value) -> String {
258    match v {
259        Value::Int(n)      => n.to_string(),
260        Value::Float(n)    => format!("{n}"),
261        Value::Bool(b)     => b.to_string(),
262        Value::Str(s)      => s.clone(),
263        Value::DateTime(t) => format!("{t}"),
264        Value::Uuid(u)     => format!("{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}",
265            u[0], u[1], u[2], u[3], u[4], u[5], u[6], u[7],
266            u[8], u[9], u[10], u[11], u[12], u[13], u[14], u[15]),
267        Value::Bytes(b)    => format!("<{} bytes>", b.len()),
268        Value::Empty       => "{}".into(),
269    }
270}
powdb_server/handler.rs

powdb_server/
handler.rs
