1use crate::protocol::Message;
2use powdb_auth::{Permission, Role, UserStore};
3use powdb_query::executor::{is_read_only_statement, Engine};
4use powdb_query::parser;
5use powdb_query::result::{QueryError, QueryResult};
6use powdb_storage::types::Value;
7use std::collections::HashMap;
8use std::net::IpAddr;
9use std::sync::{Arc, Mutex, RwLock};
10use std::time::{Duration, Instant};
11use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, BufReader, BufWriter};
12use tokio::sync::watch;
13use tracing::{debug, error, info, warn};
14use zeroize::Zeroizing;
15
16pub type AuthRateLimiter = Arc<Mutex<HashMap<IpAddr, (u32, Instant)>>>;
18
19const MAX_QUERY_LENGTH: usize = 1024 * 1024;
21
22const WRITE_TIMEOUT: Duration = Duration::from_secs(30);
25
26const MAX_AUTH_FAILURES: u32 = 5;
28
29const AUTH_FAILURE_WINDOW: Duration = Duration::from_secs(60);
31
32pub fn new_rate_limiter() -> AuthRateLimiter {
34 Arc::new(Mutex::new(HashMap::new()))
35}
36
37fn is_rate_limited(limiter: &AuthRateLimiter, ip: IpAddr) -> bool {
40 let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
41 let now = Instant::now();
43 map.retain(|_, (_, ts)| now.duration_since(*ts) < AUTH_FAILURE_WINDOW);
44
45 if let Some((count, _)) = map.get(&ip) {
46 *count >= MAX_AUTH_FAILURES
47 } else {
48 false
49 }
50}
51
52fn record_auth_failure(limiter: &AuthRateLimiter, ip: IpAddr) {
54 let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
55 let now = Instant::now();
56 let entry = map.entry(ip).or_insert((0, now));
57 if now.duration_since(entry.1) >= AUTH_FAILURE_WINDOW {
59 *entry = (1, now);
60 } else {
61 entry.0 += 1;
62 }
63}
64
65fn clear_auth_failures(limiter: &AuthRateLimiter, ip: IpAddr) {
67 let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
68 map.remove(&ip);
69}
70
71fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
74 use sha2::{Digest, Sha256};
75 let ha = Sha256::digest(a);
76 let hb = Sha256::digest(b);
77 let mut diff = 0u8;
78 for (x, y) in ha.iter().zip(hb.iter()) {
79 diff |= x ^ y;
80 }
81 diff == 0
82}
83
84#[derive(Debug, Clone, PartialEq, Eq)]
88pub struct Principal {
89 pub name: String,
90 pub role: String,
91}
92
93fn role_can_write(role: &str) -> bool {
97 Role::builtin(role).is_some_and(|r| r.allows(Permission::Write))
98}
99
100fn check_statement_permitted(
109 principal: Option<&Principal>,
110 stmt: &powdb_query::ast::Statement,
111) -> Result<(), QueryError> {
112 let Some(p) = principal else {
113 return Ok(());
116 };
117 if is_read_only_statement(stmt) || role_can_write(&p.role) {
118 return Ok(());
119 }
120 Err(QueryError::Execution(format!(
121 "permission denied: role '{}' cannot execute write statements",
122 p.role
123 )))
124}
125
126#[derive(Debug, Clone, PartialEq, Eq)]
128pub enum AuthOutcome {
129 Authenticated { principal: Option<Principal> },
133 Rejected,
136}
137
138pub fn authenticate_connect(
150 users: &UserStore,
151 expected_password: Option<&str>,
152 username: Option<&str>,
153 password: Option<&str>,
154) -> AuthOutcome {
155 if !users.is_empty() {
156 let Some(name) = username else {
158 return AuthOutcome::Rejected;
159 };
160 let Some(candidate) = password else {
161 return AuthOutcome::Rejected;
162 };
163 match users.authenticate(name, candidate) {
164 Some(user) => AuthOutcome::Authenticated {
165 principal: Some(Principal {
166 name: user.name.clone(),
167 role: user.role.clone(),
168 }),
169 },
170 None => AuthOutcome::Rejected,
171 }
172 } else {
173 match expected_password {
175 Some(expected) => {
176 if password.is_some_and(|p| constant_time_eq(p.as_bytes(), expected.as_bytes())) {
177 AuthOutcome::Authenticated { principal: None }
178 } else {
179 AuthOutcome::Rejected
180 }
181 }
182 None => AuthOutcome::Authenticated { principal: None },
183 }
184 }
185}
186
187const SAFE_ERROR_PREFIXES: &[&str] = &[
189 "table not found",
190 "column not found",
191 "parse error",
192 "type mismatch",
193 "unknown table",
194 "unknown column",
195 "unknown function",
196 "syntax error",
197 "expected",
198 "unexpected",
199 "missing",
200 "duplicate",
201 "invalid",
202 "cannot",
203 "no such",
204 "already exists",
205 "permission denied",
206 "row too large",
207];
208
209fn sanitize_error(e: &str) -> String {
213 let lower = e.to_lowercase();
214 for prefix in SAFE_ERROR_PREFIXES {
215 if lower.starts_with(prefix) {
216 return e.to_string();
217 }
218 }
219 "query execution error".into()
220}
221
222async fn write_msg<W: AsyncWrite + Unpin>(writer: &mut BufWriter<W>, msg: &Message) -> bool {
225 let write_fut = async {
226 if msg.write_to(writer).await.is_err() {
227 return false;
228 }
229 writer.flush().await.is_ok()
230 };
231 tokio::time::timeout(WRITE_TIMEOUT, write_fut)
232 .await
233 .unwrap_or_default()
234}
235
236pub struct ConnOpts<'a> {
239 pub engine: Arc<RwLock<Engine>>,
240 pub expected_password: Option<Zeroizing<String>>,
243 pub users: Arc<UserStore>,
247 pub shutdown_rx: &'a mut watch::Receiver<bool>,
248 pub idle_timeout: Duration,
249 pub query_timeout: Duration,
250 pub rate_limiter: Option<&'a AuthRateLimiter>,
251 pub peer_addr: Option<std::net::SocketAddr>,
252}
253
254fn dispatch_query(
263 engine: &Arc<RwLock<Engine>>,
264 query: &str,
265 principal: Option<&Principal>,
266) -> Result<QueryResult, QueryError> {
267 let stmt_result = parser::parse(query).map_err(|e| e.to_string());
268
269 if let Ok(stmt) = &stmt_result {
273 check_statement_permitted(principal, stmt)?;
274 }
275
276 let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
277 if can_try_read {
278 let res = {
279 let eng = engine
280 .read()
281 .map_err(|e| QueryError::Execution(format!("lock poisoned: {e}")))?;
282 eng.execute_powql_readonly(query)
283 };
284 match res {
285 Ok(r) => return Ok(r),
286 Err(QueryError::ReadonlyNeedsWrite) => {
287 }
289 Err(e) => return Err(e),
290 }
291 }
292
293 let mut eng = engine
294 .write()
295 .map_err(|e| QueryError::Execution(format!("lock poisoned: {e}")))?;
296 eng.execute_powql(query)
297}
298
299pub async fn handle_connection<S>(stream: S, opts: ConnOpts<'_>)
300where
301 S: AsyncRead + AsyncWrite + Unpin,
302{
303 let ConnOpts {
304 engine,
305 expected_password,
306 users,
307 shutdown_rx,
308 idle_timeout,
309 query_timeout,
310 rate_limiter,
311 peer_addr,
312 } = opts;
313
314 let peer = peer_addr
315 .map(|a| a.to_string())
316 .unwrap_or_else(|| "unknown".into());
317 let peer_ip = peer_addr.map(|a| a.ip());
318
319 let (reader, writer) = tokio::io::split(stream);
320 let mut reader = BufReader::new(reader);
321 let mut writer = BufWriter::new(writer);
322
323 let connect_msg = loop {
328 match tokio::time::timeout(idle_timeout, Message::read_from_preauth(&mut reader)).await {
329 Ok(Ok(Some(Message::Ping))) => {
330 debug!(peer = %peer, "pre-auth ping");
331 if !write_msg(&mut writer, &Message::Pong).await {
332 return;
333 }
334 continue;
335 }
336 Ok(Ok(Some(msg))) => break msg,
337 Ok(Ok(None)) => {
338 debug!(peer = %peer, "client closed before CONNECT");
339 return;
340 }
341 Ok(Err(e)) => {
342 error!(peer = %peer, error = %e, "error reading CONNECT");
343 return;
344 }
345 Err(_) => {
346 warn!(peer = %peer, "idle timeout waiting for CONNECT");
347 return;
348 }
349 }
350 };
351
352 let principal: Option<Principal>;
355 match connect_msg {
356 Message::Connect {
357 db_name,
358 password,
359 username,
360 } => {
361 if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
363 if is_rate_limited(limiter, ip) {
364 warn!(peer = %peer, "rate limited: too many auth failures");
365 let err = Message::Error {
366 message: "too many auth failures, try again later".into(),
367 };
368 write_msg(&mut writer, &err).await;
369 return;
370 }
371 }
372
373 let outcome = authenticate_connect(
374 &users,
375 expected_password.as_ref().map(|p| p.as_str()),
376 username.as_deref(),
377 password.as_ref().map(|p| p.as_str()),
378 );
379
380 match outcome {
381 AuthOutcome::Rejected => {
382 warn!(peer = %peer, db = %db_name, "auth rejected");
383 if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
385 record_auth_failure(limiter, ip);
386 }
387 let err = Message::Error {
388 message: "authentication failed".into(),
389 };
390 write_msg(&mut writer, &err).await;
391 return;
392 }
393 AuthOutcome::Authenticated {
394 principal: auth_principal,
395 } => {
396 if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
398 clear_auth_failures(limiter, ip);
399 }
400 match &auth_principal {
401 Some(p) => {
402 info!(peer = %peer, db = %db_name, user = %p.name, role = %p.role, "authenticated");
403 }
404 None => {
405 info!(peer = %peer, db = %db_name, "client connected");
406 }
407 }
408 principal = auth_principal;
409 }
410 }
411
412 let ok = Message::ConnectOk {
413 version: env!("CARGO_PKG_VERSION").into(),
414 };
415 if !write_msg(&mut writer, &ok).await {
416 return;
417 }
418 }
419 _ => {
420 warn!(peer = %peer, "first message was not CONNECT");
421 let err = Message::Error {
422 message: "expected CONNECT".into(),
423 };
424 write_msg(&mut writer, &err).await;
425 return;
426 }
427 }
428
429 loop {
431 let msg = tokio::select! {
432 result = tokio::time::timeout(idle_timeout, Message::read_from(&mut reader)) => {
434 match result {
435 Ok(Ok(Some(msg))) => msg,
436 Ok(Ok(None)) => break,
437 Ok(Err(e)) => {
438 error!(peer = %peer, error = %e, "read error");
439 break;
440 }
441 Err(_) => {
442 info!(peer = %peer, "idle timeout, closing connection");
443 let err = Message::Error { message: "idle timeout".into() };
444 write_msg(&mut writer, &err).await;
445 break;
446 }
447 }
448 }
449 _ = shutdown_rx.changed() => {
451 if *shutdown_rx.borrow() {
452 info!(peer = %peer, "server shutting down, closing connection");
453 let err = Message::Error { message: "server shutting down".into() };
454 write_msg(&mut writer, &err).await;
455 break;
456 }
457 continue;
458 }
459 };
460
461 let response = match msg {
462 Message::Ping => {
463 debug!(peer = %peer, "ping");
464 Message::Pong
465 }
466 Message::Query { query } => {
467 if query.len() > MAX_QUERY_LENGTH {
468 Message::Error {
469 message: format!(
470 "query too large: {} bytes (max {})",
471 query.len(),
472 MAX_QUERY_LENGTH
473 ),
474 }
475 } else {
476 debug!(peer = %peer, query = %query, "received query");
477 let handle = tokio::task::spawn_blocking({
478 let engine = engine.clone();
479 let query = query.clone();
480 let principal = principal.clone();
481 move || dispatch_query(&engine, &query, principal.as_ref())
482 });
483 let abort_handle = handle.abort_handle();
484 match tokio::time::timeout(query_timeout, handle).await {
485 Ok(Ok(Ok(result))) => query_result_to_message(result),
486 Ok(Ok(Err(e))) => Message::Error {
487 message: sanitize_error(&e.to_string()),
488 },
489 Ok(Err(e)) => Message::Error {
490 message: format!("internal error: {e}"),
491 },
492 Err(_) => {
493 abort_handle.abort();
494 warn!(peer = %peer, query = %query, "query timeout exceeded");
495 Message::Error {
496 message: "query timeout exceeded".into(),
497 }
498 }
499 }
500 }
501 }
502 Message::Disconnect => {
503 debug!(peer = %peer, "received DISCONNECT");
504 break;
505 }
506 _ => Message::Error {
507 message: "unexpected message type".into(),
508 },
509 };
510
511 if !write_msg(&mut writer, &response).await {
512 break;
513 }
514 }
515
516 info!(peer = %peer, "client disconnected");
517}
518
519fn query_result_to_message(result: QueryResult) -> Message {
520 match result {
521 QueryResult::Rows { columns, rows } => {
522 let str_rows: Vec<Vec<String>> = rows
523 .iter()
524 .map(|row| row.iter().map(value_to_display).collect())
525 .collect();
526 Message::ResultRows {
527 columns,
528 rows: str_rows,
529 }
530 }
531 QueryResult::Scalar(val) => Message::ResultScalar {
532 value: value_to_display(&val),
533 },
534 QueryResult::Modified(n) => Message::ResultOk { affected: n },
535 QueryResult::Created(name) => Message::ResultMessage {
536 message: format!("type {name} created"),
537 },
538 QueryResult::Executed { message } => Message::ResultMessage { message },
539 }
540}
541
542fn value_to_display(v: &Value) -> String {
543 match v {
544 Value::Int(n) => n.to_string(),
545 Value::Float(n) => format!("{n}"),
546 Value::Bool(b) => b.to_string(),
547 Value::Str(s) => s.clone(),
548 Value::DateTime(t) => format!("{t}"),
549 Value::Uuid(u) => format!("{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}",
550 u[0], u[1], u[2], u[3], u[4], u[5], u[6], u[7],
551 u[8], u[9], u[10], u[11], u[12], u[13], u[14], u[15]),
552 Value::Bytes(b) => format!("<{} bytes>", b.len()),
553 Value::Empty => "null".into(),
559 }
560}
561
562#[cfg(test)]
563mod tests {
564 use super::*;
565
566 #[test]
569 fn null_serializes_as_null_bareword_on_wire() {
570 assert_eq!(value_to_display(&Value::Empty), "null");
571 }
572
573 fn parsed(q: &str) -> powdb_query::ast::Statement {
576 parser::parse(q).unwrap()
577 }
578
579 fn principal(role: &str) -> Option<Principal> {
580 Some(Principal {
581 name: "u".into(),
582 role: role.into(),
583 })
584 }
585
586 #[test]
587 fn readonly_can_read_but_not_write() {
588 let p = principal("readonly");
589 assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
591 assert!(check_statement_permitted(p.as_ref(), &parsed("count(User)")).is_ok());
592 assert!(check_statement_permitted(p.as_ref(), &parsed("explain User")).is_ok());
593 for q in [
595 r#"insert User { name := "x" }"#,
596 "User filter .id = 1 update { age := 2 }",
597 "User filter .id = 1 delete",
598 "drop User",
599 "alter User add column c: str",
600 "type T { required id: int }",
601 "begin",
602 "commit",
603 "rollback",
604 ] {
605 let err = check_statement_permitted(p.as_ref(), &parsed(q))
606 .expect_err(&format!("must deny: {q}"));
607 assert!(
608 err.to_string().contains("permission denied"),
609 "unexpected error for {q}: {err}"
610 );
611 }
612 }
613
614 #[test]
615 fn readwrite_and_admin_have_full_query_access() {
616 for role in ["readwrite", "admin"] {
617 let p = principal(role);
618 assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
619 assert!(check_statement_permitted(
620 p.as_ref(),
621 &parsed(r#"insert User { name := "x" }"#)
622 )
623 .is_ok());
624 assert!(check_statement_permitted(p.as_ref(), &parsed("drop User")).is_ok());
625 }
626 }
627
628 #[test]
629 fn unknown_role_fails_closed_for_writes() {
630 let p = principal("mystery");
631 assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
632 assert!(
633 check_statement_permitted(p.as_ref(), &parsed(r#"insert User { name := "x" }"#))
634 .is_err()
635 );
636 }
637
638 #[test]
639 fn no_principal_means_full_access() {
640 assert!(check_statement_permitted(None, &parsed("drop User")).is_ok());
642 assert!(check_statement_permitted(None, &parsed(r#"insert User { name := "x" }"#)).is_ok());
643 }
644
645 fn store_with_alice() -> UserStore {
646 let mut s = UserStore::new();
647 s.create_user("alice", "pw", "readwrite").unwrap();
648 s
649 }
650
651 #[test]
654 fn empty_store_no_password_is_open() {
655 let s = UserStore::new();
656 assert_eq!(
657 authenticate_connect(&s, None, None, None),
658 AuthOutcome::Authenticated { principal: None }
659 );
660 assert_eq!(
662 authenticate_connect(&s, None, Some("x"), Some("y")),
663 AuthOutcome::Authenticated { principal: None }
664 );
665 }
666
667 #[test]
668 fn empty_store_correct_shared_password_succeeds() {
669 let s = UserStore::new();
670 assert_eq!(
671 authenticate_connect(&s, Some("pw"), None, Some("pw")),
672 AuthOutcome::Authenticated { principal: None }
673 );
674 }
675
676 #[test]
677 fn empty_store_wrong_shared_password_rejected() {
678 let s = UserStore::new();
679 assert_eq!(
680 authenticate_connect(&s, Some("pw"), None, Some("bad")),
681 AuthOutcome::Rejected
682 );
683 }
684
685 #[test]
686 fn empty_store_missing_password_rejected_when_expected() {
687 let s = UserStore::new();
688 assert_eq!(
689 authenticate_connect(&s, Some("pw"), None, None),
690 AuthOutcome::Rejected
691 );
692 }
693
694 #[test]
695 fn empty_store_ignores_username_for_shared_password() {
696 let s = UserStore::new();
699 assert_eq!(
700 authenticate_connect(&s, Some("pw"), Some("whoever"), Some("pw")),
701 AuthOutcome::Authenticated { principal: None }
702 );
703 }
704
705 #[test]
708 fn user_auth_success_binds_principal() {
709 let s = store_with_alice();
710 assert_eq!(
711 authenticate_connect(&s, None, Some("alice"), Some("pw")),
712 AuthOutcome::Authenticated {
713 principal: Some(Principal {
714 name: "alice".into(),
715 role: "readwrite".into(),
716 })
717 }
718 );
719 }
720
721 #[test]
722 fn user_auth_wrong_password_rejected() {
723 let s = store_with_alice();
724 assert_eq!(
725 authenticate_connect(&s, None, Some("alice"), Some("bad")),
726 AuthOutcome::Rejected
727 );
728 }
729
730 #[test]
731 fn user_auth_unknown_user_rejected() {
732 let s = store_with_alice();
733 assert_eq!(
734 authenticate_connect(&s, None, Some("mallory"), Some("pw")),
735 AuthOutcome::Rejected
736 );
737 }
738
739 #[test]
740 fn user_auth_missing_username_rejected() {
741 let s = store_with_alice();
742 assert_eq!(
743 authenticate_connect(&s, None, None, Some("pw")),
744 AuthOutcome::Rejected
745 );
746 }
747
748 #[test]
749 fn user_auth_missing_password_rejected() {
750 let s = store_with_alice();
751 assert_eq!(
752 authenticate_connect(&s, Some("pw"), Some("alice"), None),
753 AuthOutcome::Rejected
754 );
755 }
756
757 #[test]
758 fn user_auth_ignores_shared_password_when_users_present() {
759 let s = store_with_alice();
762 assert_eq!(
763 authenticate_connect(&s, Some("shared"), None, Some("shared")),
764 AuthOutcome::Rejected
765 );
766 }
767}