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

powdb_server/
handler.rs
