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

powdb_server/
handler.rs
