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::collections::HashMap;
7use std::net::IpAddr;
8use std::sync::{Arc, Mutex, RwLock};
9use std::time::{Duration, Instant};
10use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, BufReader, BufWriter};
11use tokio::sync::watch;
12use tracing::{debug, error, info, warn};
13
14/// Tracks per-IP authentication failure counts for rate limiting.
15pub type AuthRateLimiter = Arc<Mutex<HashMap<IpAddr, (u32, Instant)>>>;
16
17/// Maximum query text length accepted from the wire (1 MB).
18const MAX_QUERY_LENGTH: usize = 1024 * 1024;
19
20/// Maximum number of auth failures per IP within the rate-limit window.
21const MAX_AUTH_FAILURES: u32 = 5;
22
23/// Window during which auth failures are counted (60 seconds).
24const AUTH_FAILURE_WINDOW: Duration = Duration::from_secs(60);
25
26/// Create a new shared rate limiter.
27pub fn new_rate_limiter() -> AuthRateLimiter {
28    Arc::new(Mutex::new(HashMap::new()))
29}
30
31/// Check whether an IP is rate-limited and record a failure if requested.
32/// Returns `true` if the IP should be rejected.
33fn is_rate_limited(limiter: &AuthRateLimiter, ip: IpAddr) -> bool {
34    let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
35    // Clean up stale entries while we have the lock.
36    let now = Instant::now();
37    map.retain(|_, (_, ts)| now.duration_since(*ts) < AUTH_FAILURE_WINDOW);
38
39    if let Some((count, _)) = map.get(&ip) {
40        *count >= MAX_AUTH_FAILURES
41    } else {
42        false
43    }
44}
45
46/// Record an auth failure for the given IP.
47fn record_auth_failure(limiter: &AuthRateLimiter, ip: IpAddr) {
48    let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
49    let now = Instant::now();
50    let entry = map.entry(ip).or_insert((0, now));
51    // Reset counter if the window has elapsed.
52    if now.duration_since(entry.1) >= AUTH_FAILURE_WINDOW {
53        *entry = (1, now);
54    } else {
55        entry.0 += 1;
56    }
57}
58
59/// Clear the failure counter on successful auth.
60fn clear_auth_failures(limiter: &AuthRateLimiter, ip: IpAddr) {
61    let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
62    map.remove(&ip);
63}
64
65/// Constant-time password comparison. Hashes both inputs to fixed-size
66/// SHA-256 digests so neither length nor content leaks through timing.
67fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
68    use sha2::{Digest, Sha256};
69    let ha = Sha256::digest(a);
70    let hb = Sha256::digest(b);
71    let mut diff = 0u8;
72    for (x, y) in ha.iter().zip(hb.iter()) {
73        diff |= x ^ y;
74    }
75    diff == 0
76}
77
78/// Error messages that are safe to forward to the client verbatim.
79const SAFE_ERROR_PREFIXES: &[&str] = &[
80    "table not found",
81    "column not found",
82    "parse error",
83    "type mismatch",
84    "unknown table",
85    "unknown column",
86    "unknown function",
87    "syntax error",
88    "expected",
89    "unexpected",
90    "missing",
91    "duplicate",
92    "invalid",
93    "cannot",
94    "no such",
95    "already exists",
96];
97
98/// Sanitize an error message before sending it to the client.
99/// Known safe errors are passed through; everything else is replaced
100/// with a generic message to avoid leaking internal details.
101fn sanitize_error(e: &str) -> String {
102    let lower = e.to_lowercase();
103    for prefix in SAFE_ERROR_PREFIXES {
104        if lower.starts_with(prefix) {
105            return e.to_string();
106        }
107    }
108    "query execution error".into()
109}
110
111/// Options for a single connection, bundled to keep `handle_connection`'s
112/// argument list short.
113pub struct ConnOpts<'a> {
114    pub engine: Arc<RwLock<Engine>>,
115    pub expected_password: Option<String>,
116    pub shutdown_rx: &'a mut watch::Receiver<bool>,
117    pub idle_timeout: Duration,
118    pub query_timeout: Duration,
119    pub rate_limiter: Option<&'a AuthRateLimiter>,
120    pub peer_addr: Option<std::net::SocketAddr>,
121}
122
123/// Execute a query against the engine under the RwLock. Read-only
124/// statements acquire `.read()` so concurrent SELECTs can scan in
125/// parallel; mutations acquire `.write()`.
126fn dispatch_query(engine: &Arc<RwLock<Engine>>, query: &str) -> Result<QueryResult, String> {
127    let stmt_result = parser::parse(query).map_err(|e| e.to_string());
128
129    let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
130    if can_try_read {
131        let res = {
132            let eng = engine.read().map_err(|e| format!("lock poisoned: {e}"))?;
133            eng.execute_powql_readonly(query)
134        };
135        match res {
136            Ok(r) => return Ok(r),
137            Err(e) if e == READONLY_NEEDS_WRITE => {
138                // Escalate: fall through to the write path below.
139            }
140            Err(e) => return Err(e),
141        }
142    }
143
144    let mut eng = engine.write().map_err(|e| format!("lock poisoned: {e}"))?;
145    eng.execute_powql(query)
146}
147
148pub async fn handle_connection<S>(stream: S, opts: ConnOpts<'_>)
149where
150    S: AsyncRead + AsyncWrite + Unpin,
151{
152    let ConnOpts {
153        engine,
154        expected_password,
155        shutdown_rx,
156        idle_timeout,
157        query_timeout,
158        rate_limiter,
159        peer_addr,
160    } = opts;
161
162    let peer = peer_addr
163        .map(|a| a.to_string())
164        .unwrap_or_else(|| "unknown".into());
165    let peer_ip = peer_addr.map(|a| a.ip());
166
167    let (reader, writer) = tokio::io::split(stream);
168    let mut reader = BufReader::new(reader);
169    let mut writer = BufWriter::new(writer);
170
171    // Wait for Connect message (with idle timeout).
172    // Accept Ping messages before authentication so load balancers can
173    // health-check without completing a full CONNECT handshake.
174    let connect_msg = loop {
175        match tokio::time::timeout(idle_timeout, Message::read_from(&mut reader)).await {
176            Ok(Ok(Some(Message::Ping))) => {
177                debug!(peer = %peer, "pre-auth ping");
178                let pong = Message::Pong;
179                if pong.write_to(&mut writer).await.is_err() {
180                    return;
181                }
182                if writer.flush().await.is_err() {
183                    return;
184                }
185                continue;
186            }
187            Ok(Ok(Some(msg))) => break msg,
188            Ok(Ok(None)) => {
189                debug!(peer = %peer, "client closed before CONNECT");
190                return;
191            }
192            Ok(Err(e)) => {
193                error!(peer = %peer, error = %e, "error reading CONNECT");
194                return;
195            }
196            Err(_) => {
197                warn!(peer = %peer, "idle timeout waiting for CONNECT");
198                return;
199            }
200        }
201    };
202
203    match connect_msg {
204        Message::Connect { db_name, password } => {
205            // Check rate limiting before verifying credentials.
206            if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
207                if is_rate_limited(limiter, ip) {
208                    warn!(peer = %peer, "rate limited: too many auth failures");
209                    let err = Message::Error {
210                        message: "too many auth failures, try again later".into(),
211                    };
212                    err.write_to(&mut writer).await.ok();
213                    writer.flush().await.ok();
214                    return;
215                }
216            }
217
218            if let Some(expected) = &expected_password {
219                if !password
220                    .as_deref()
221                    .is_some_and(|p| constant_time_eq(p.as_bytes(), expected.as_bytes()))
222                {
223                    warn!(peer = %peer, db = %db_name, "auth rejected: bad password");
224                    // Record the failure for rate limiting.
225                    if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
226                        record_auth_failure(limiter, ip);
227                    }
228                    let err = Message::Error {
229                        message: "authentication failed".into(),
230                    };
231                    err.write_to(&mut writer).await.ok();
232                    writer.flush().await.ok();
233                    return;
234                }
235            }
236            // Auth succeeded — clear any prior failure count.
237            if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
238                clear_auth_failures(limiter, ip);
239            }
240            info!(peer = %peer, db = %db_name, "client connected");
241            let ok = Message::ConnectOk {
242                version: env!("CARGO_PKG_VERSION").into(),
243            };
244            if ok.write_to(&mut writer).await.is_err() {
245                return;
246            }
247            if writer.flush().await.is_err() {
248                return;
249            }
250        }
251        _ => {
252            warn!(peer = %peer, "first message was not CONNECT");
253            let err = Message::Error {
254                message: "expected CONNECT".into(),
255            };
256            err.write_to(&mut writer).await.ok();
257            writer.flush().await.ok();
258            return;
259        }
260    }
261
262    // Main query loop with idle timeout and shutdown awareness.
263    loop {
264        let msg = tokio::select! {
265            // Read next message with idle timeout.
266            result = tokio::time::timeout(idle_timeout, Message::read_from(&mut reader)) => {
267                match result {
268                    Ok(Ok(Some(msg))) => msg,
269                    Ok(Ok(None)) => break,
270                    Ok(Err(e)) => {
271                        error!(peer = %peer, error = %e, "read error");
272                        break;
273                    }
274                    Err(_) => {
275                        info!(peer = %peer, "idle timeout, closing connection");
276                        let err = Message::Error { message: "idle timeout".into() };
277                        err.write_to(&mut writer).await.ok();
278                        writer.flush().await.ok();
279                        break;
280                    }
281                }
282            }
283            // If server is shutting down, notify client and close.
284            _ = shutdown_rx.changed() => {
285                if *shutdown_rx.borrow() {
286                    info!(peer = %peer, "server shutting down, closing connection");
287                    let err = Message::Error { message: "server shutting down".into() };
288                    err.write_to(&mut writer).await.ok();
289                    writer.flush().await.ok();
290                    break;
291                }
292                continue;
293            }
294        };
295
296        let response = match msg {
297            Message::Ping => {
298                debug!(peer = %peer, "ping");
299                Message::Pong
300            }
301            Message::Query { query } => {
302                if query.len() > MAX_QUERY_LENGTH {
303                    Message::Error {
304                        message: format!(
305                            "query too large: {} bytes (max {})",
306                            query.len(),
307                            MAX_QUERY_LENGTH
308                        ),
309                    }
310                } else {
311                    debug!(peer = %peer, query = %query, "received query");
312                    let result = tokio::task::spawn_blocking({
313                        let engine = engine.clone();
314                        let query = query.clone();
315                        move || dispatch_query(&engine, &query)
316                    });
317                    match tokio::time::timeout(query_timeout, result).await {
318                        Ok(Ok(Ok(result))) => query_result_to_message(result),
319                        Ok(Ok(Err(e))) => Message::Error {
320                            message: sanitize_error(&e),
321                        },
322                        Ok(Err(e)) => Message::Error {
323                            message: format!("internal error: {e}"),
324                        },
325                        Err(_) => {
326                            warn!(peer = %peer, query = %query, "query timeout exceeded");
327                            Message::Error {
328                                message: "query timeout exceeded".into(),
329                            }
330                        }
331                    }
332                }
333            }
334            Message::Disconnect => {
335                debug!(peer = %peer, "received DISCONNECT");
336                break;
337            }
338            _ => Message::Error {
339                message: "unexpected message type".into(),
340            },
341        };
342
343        if response.write_to(&mut writer).await.is_err() {
344            break;
345        }
346        if writer.flush().await.is_err() {
347            break;
348        }
349    }
350
351    info!(peer = %peer, "client disconnected");
352}
353
354fn query_result_to_message(result: QueryResult) -> Message {
355    match result {
356        QueryResult::Rows { columns, rows } => {
357            let str_rows: Vec<Vec<String>> = rows
358                .iter()
359                .map(|row| row.iter().map(value_to_display).collect())
360                .collect();
361            Message::ResultRows {
362                columns,
363                rows: str_rows,
364            }
365        }
366        QueryResult::Scalar(val) => Message::ResultScalar {
367            value: value_to_display(&val),
368        },
369        QueryResult::Modified(n) => Message::ResultOk { affected: n },
370        QueryResult::Created(_name) => Message::ResultOk { affected: 0 },
371        QueryResult::Executed { .. } => Message::ResultOk { affected: 0 },
372    }
373}
374
375fn value_to_display(v: &Value) -> String {
376    match v {
377        Value::Int(n)      => n.to_string(),
378        Value::Float(n)    => format!("{n}"),
379        Value::Bool(b)     => b.to_string(),
380        Value::Str(s)      => s.clone(),
381        Value::DateTime(t) => format!("{t}"),
382        Value::Uuid(u)     => format!("{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}",
383            u[0], u[1], u[2], u[3], u[4], u[5], u[6], u[7],
384            u[8], u[9], u[10], u[11], u[12], u[13], u[14], u[15]),
385        Value::Bytes(b)    => format!("<{} bytes>", b.len()),
386        Value::Empty       => "{}".into(),
387    }
388}
powdb_server/handler.rs

powdb_server/
handler.rs
