powdb_server/
handler.rs

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

powdb_server/
handler.rs
