Skip to main content

reddb_server/wire/postgres/
server.rs

1//! PostgreSQL wire-protocol listener (Phase 3.1 PG parity).
2//!
3//! Accepts TCP connections from PG-compatible clients, drives the startup
4//! handshake, and routes simple/extended-query frames into the existing
5//! `RedDBRuntime::execute_query` path. Results are adapted back into PG
6//! `RowDescription` + `DataRow` frames via `types::value_to_pg_wire_bytes`.
7
8use std::collections::HashMap;
9use std::sync::Arc;
10
11use tokio::io::{AsyncRead, AsyncWrite};
12use tokio::net::TcpListener;
13
14use super::catalog_views::translate_pg_catalog_query;
15use super::protocol::{
16    read_frame, read_startup, write_frame, write_raw_byte, BackendMessage, ColumnDescriptor,
17    DescribeTarget, FrontendMessage, PgWireError, TransactionStatus,
18};
19use super::types::{pg_param_to_param_value, value_to_pg_wire_bytes, PgOid};
20use crate::auth::Role;
21use crate::runtime::query_request::{
22    ParamValue, PreparedId, PreparedRegistry, QueryRequest, QueryRequestExecutor,
23};
24use crate::runtime::RedDBRuntime;
25use crate::storage::query::unified::{UnifiedRecord, UnifiedResult};
26use reddb_rql::ast::QueryExpr;
27use reddb_types::Value;
28
29/// Startup-tuned configuration for the PG wire listener.
30#[derive(Debug, Clone)]
31pub struct PgWireConfig {
32    /// TCP bind address ("host:port"). The caller is responsible for
33    /// keeping this disjoint from the native wire / gRPC / HTTP listeners.
34    pub bind_addr: String,
35    /// PG version string sent back in `ParameterStatus`. Many drivers
36    /// sniff this to enable/disable features. RedDB advertises a
37    /// recent-enough version to get the broadest client support.
38    pub server_version: String,
39    /// Optional TLS material. When set, an `SSLRequest` is answered with
40    /// `'S'` and the connection is upgraded via a rustls handshake before
41    /// the startup flow continues over the encrypted stream. The cert/key
42    /// surface is the *same* [`WireTlsConfig`] the native RedWire listener
43    /// uses — the two transports share one cert config rather than
44    /// inventing a parallel one. `None` leaves the historical behaviour:
45    /// `SSLRequest` is declined with `'N'` and clients continue in
46    /// plaintext.
47    pub tls: Option<crate::wire::tls::WireTlsConfig>,
48}
49
50/// The authenticated identity resolved during the PG-Wire startup
51/// handshake. Installed as the thread-local execution identity for the
52/// duration of each statement so every query dispatched over this
53/// connection flows through the same IAM/policy gates as native
54/// transports. `None` connections (loopback trust) run under the ambient
55/// embedded identity, preserving the pre-auth behaviour.
56#[derive(Debug, Clone)]
57struct PgAuthContext {
58    username: String,
59    role: Role,
60    tenant: Option<String>,
61}
62
63#[derive(Debug, Clone)]
64struct PgPreparedStatement {
65    sql: String,
66    prepared_id: PreparedId,
67    param_type_oids: Vec<u32>,
68}
69
70#[derive(Debug, Clone)]
71struct PgPortal {
72    sql: String,
73    prepared_id: PreparedId,
74    params: Vec<ParamValue>,
75    #[allow(dead_code)]
76    result_format_codes: Vec<i16>,
77    row_description_sent: bool,
78    result: Option<crate::runtime::OwnedQueryResult>,
79    row_offset: usize,
80}
81
82impl Default for PgWireConfig {
83    fn default() -> Self {
84        Self {
85            bind_addr: "127.0.0.1:5432".to_string(),
86            server_version: "15.0 (RedDB 3.1)".to_string(),
87            tls: None,
88        }
89    }
90}
91
92/// Run a synchronous, possibly long-parking runtime call without
93/// monopolising the async worker thread.
94///
95/// `QUEUE READ … WAIT <budget>` parks the calling thread on a condvar
96/// for up to the wait budget. Doing that directly on a tokio worker
97/// holds the worker for the whole budget, which serialises every other
98/// connection's query behind the parked waiter — a producer's
99/// `QUEUE PUSH` could not run to notify the waiter, so an enqueue during
100/// a WAIT would never be delivered — and it stalls the runtime's timer
101/// driver. `block_in_place` hands the remaining tasks on this worker to
102/// a sibling worker for the duration while keeping the current OS thread
103/// — and therefore the per-connection thread-locals `execute_query`
104/// relies on (current tenant, connection id, transaction context) —
105/// intact, which `spawn_blocking` would not.
106///
107/// `block_in_place` panics on a current-thread runtime, so we fall back
108/// to a direct call there: a single-threaded runtime has no sibling
109/// worker to hand work to and nothing else to starve (e.g. the
110/// `#[tokio::test]` harnesses that drive one connection at a time).
111fn run_runtime_blocking<T>(f: impl FnOnce() -> T) -> T {
112    use tokio::runtime::{Handle, RuntimeFlavor};
113    match Handle::try_current().map(|h| h.runtime_flavor()) {
114        Ok(RuntimeFlavor::MultiThread) => tokio::task::block_in_place(f),
115        _ => f(),
116    }
117}
118
119fn pg_frame_requires_execution(frame: &FrontendMessage) -> bool {
120    matches!(
121        frame,
122        FrontendMessage::Query(_)
123            | FrontendMessage::Parse(_)
124            | FrontendMessage::Describe(_)
125            | FrontendMessage::Execute(_)
126    )
127}
128
129async fn acquire_pg_execution(
130    runtime: &RedDBRuntime,
131) -> Result<crate::transport::execution_admission::ExecutionPermit, &'static str> {
132    use crate::transport::execution_admission::ExecutionAdmissionError;
133
134    let admission = runtime.execution_admission();
135    match tokio::time::timeout(runtime.execution_queue_timeout(), admission.acquire()).await {
136        Ok(Ok(permit)) => Ok(permit),
137        Ok(Err(ExecutionAdmissionError::QueueFull)) => {
138            Err("execution_queue_full: engine execution queue is full; retryable=true")
139        }
140        Ok(Err(ExecutionAdmissionError::Closed)) => {
141            Err("execution_admission_closed: engine execution admission is unavailable")
142        }
143        Err(_) => Err(
144            "execution_queue_timeout: deadline expired while waiting to execute; retryable=true",
145        ),
146    }
147}
148
149/// Holds the execution slot only until the response starts. PG handlers do
150/// all synchronous engine work before their first write; releasing at that
151/// boundary prevents a slow client from occupying engine capacity while rows
152/// drain through the socket.
153struct ExecutionResponseWriter<'a, S> {
154    inner: &'a mut S,
155    permit: Option<crate::transport::execution_admission::ExecutionPermit>,
156}
157
158impl<'a, S> ExecutionResponseWriter<'a, S> {
159    fn new(
160        inner: &'a mut S,
161        permit: Option<crate::transport::execution_admission::ExecutionPermit>,
162    ) -> Self {
163        Self { inner, permit }
164    }
165}
166
167impl<S: AsyncWrite + Unpin> AsyncWrite for ExecutionResponseWriter<'_, S> {
168    fn poll_write(
169        mut self: std::pin::Pin<&mut Self>,
170        cx: &mut std::task::Context<'_>,
171        buf: &[u8],
172    ) -> std::task::Poll<Result<usize, std::io::Error>> {
173        drop(self.permit.take());
174        std::pin::Pin::new(&mut *self.inner).poll_write(cx, buf)
175    }
176
177    fn poll_flush(
178        mut self: std::pin::Pin<&mut Self>,
179        cx: &mut std::task::Context<'_>,
180    ) -> std::task::Poll<Result<(), std::io::Error>> {
181        drop(self.permit.take());
182        std::pin::Pin::new(&mut *self.inner).poll_flush(cx)
183    }
184
185    fn poll_shutdown(
186        mut self: std::pin::Pin<&mut Self>,
187        cx: &mut std::task::Context<'_>,
188    ) -> std::task::Poll<Result<(), std::io::Error>> {
189        drop(self.permit.take());
190        std::pin::Pin::new(&mut *self.inner).poll_shutdown(cx)
191    }
192}
193
194/// Run `f` with the authenticated PG-Wire identity installed as the
195/// thread-local execution identity, restoring the prior state afterwards.
196///
197/// The runtime's query dispatch resolves `CURRENT_USER/ROLE/TENANT` and
198/// folds RLS/IAM policies from these thread-locals, so installing them
199/// here is what makes PG-Wire statements transport-invariant with the
200/// native wire. A `None` context leaves the ambient identity untouched.
201fn execute_with_pg_auth_context<T>(
202    auth_context: Option<&PgAuthContext>,
203    f: impl FnOnce() -> T,
204) -> T {
205    let Some(auth_context) = auth_context else {
206        return f();
207    };
208    let _guard = PgAuthContextGuard::install(auth_context);
209    f()
210}
211
212/// RAII guard that installs the PG-Wire auth identity + tenant on the
213/// current thread and restores the previous state on drop, so pooled
214/// runtime threads never leak an identity across statements.
215struct PgAuthContextGuard {
216    previous_tenant: Option<String>,
217}
218
219impl PgAuthContextGuard {
220    fn install(auth_context: &PgAuthContext) -> Self {
221        let previous_tenant = crate::runtime::execution_context::current_tenant();
222        match &auth_context.tenant {
223            Some(tenant) => crate::runtime::execution_context::set_current_tenant(tenant.clone()),
224            None => crate::runtime::execution_context::clear_current_tenant(),
225        }
226        crate::runtime::execution_context::set_current_auth_identity(
227            auth_context.username.clone(),
228            auth_context.role,
229        );
230        Self { previous_tenant }
231    }
232}
233
234impl Drop for PgAuthContextGuard {
235    fn drop(&mut self) {
236        crate::runtime::execution_context::clear_current_auth_identity();
237        match self.previous_tenant.take() {
238            Some(tenant) => crate::runtime::execution_context::set_current_tenant(tenant),
239            None => crate::runtime::execution_context::clear_current_tenant(),
240        }
241    }
242}
243
244/// Whether `bind_addr` ("host:port") resolves to a loopback host. Trust
245/// authentication survives only for loopback binds with no credentials
246/// configured; every other bind requires a verified password.
247fn is_loopback_bind(bind_addr: &str) -> bool {
248    crate::transport::is_loopback_bind(bind_addr)
249}
250
251/// Extract the cleartext password from a `PasswordMessage` payload (a
252/// single NUL-terminated C-string), returning `None` on malformed input.
253fn password_from_message(payload: &[u8]) -> Option<String> {
254    let nul = payload.iter().position(|&b| b == 0)?;
255    std::str::from_utf8(&payload[..nul])
256        .ok()
257        .map(|password| password.to_string())
258}
259
260/// Spawn the PG wire listener. Blocks until the listener errors out.
261/// Each connection is handled in its own tokio task.
262pub async fn start_pg_wire_listener(
263    config: PgWireConfig,
264    runtime: Arc<RedDBRuntime>,
265) -> Result<(), Box<dyn std::error::Error>> {
266    let listener = TcpListener::bind(&config.bind_addr).await?;
267    start_pg_wire_listener_on(listener, config, runtime).await
268}
269
270/// Serve PostgreSQL wire traffic on a listener bound during node bootstrap.
271pub async fn start_pg_wire_listener_on(
272    listener: TcpListener,
273    config: PgWireConfig,
274    runtime: Arc<RedDBRuntime>,
275) -> Result<(), Box<dyn std::error::Error>> {
276    // Posture: with no auth store enabled, PG-Wire falls back to trust
277    // auth for loopback binds only (see `authenticate_startup`). Surface a
278    // startup warning whenever that fallback is active so an operator never
279    // exposes an unauthenticated listener without noticing.
280    let trust_active = runtime
281        .auth_store()
282        .map(|store| !store.is_enabled())
283        .unwrap_or(true)
284        && is_loopback_bind(&config.bind_addr);
285    if trust_active {
286        tracing::warn!(
287            transport = "pg-wire",
288            bind = %config.bind_addr,
289            "PG-Wire trust authentication is active (no auth store configured); \
290             loopback-only. Configure auth before exposing this listener."
291        );
292    }
293    // Build the shared TLS acceptor once, up front, so a bad cert/key
294    // fails the whole listener rather than every individual connection.
295    // The acceptor is cheap to clone (an `Arc` internally) and is handed
296    // to each connection so `SSLRequest` can be answered with `'S'` and
297    // the rustls handshake performed before the startup flow resumes.
298    let tls_acceptor = match config.tls.as_ref() {
299        Some(tls_cfg) => Some(crate::wire::tls::build_tls_acceptor(tls_cfg)?),
300        None => None,
301    };
302    tracing::info!(
303        transport = "pg-wire",
304        bind = %config.bind_addr,
305        tls = tls_acceptor.is_some(),
306        "listener online"
307    );
308    let cfg = Arc::new(config);
309    let admission = crate::wire::admission::ConnectionAdmission::new("pg-wire");
310    loop {
311        let (stream, peer) = listener.accept().await?;
312        // Bound concurrent sessions. Without this the accept loop spawned a
313        // task per connection with no ceiling, so an unauthenticated peer
314        // could exhaust descriptors before presenting a credential.
315        let Some(permit) = admission.try_admit() else {
316            continue;
317        };
318        // Small result/control frames must not wait for a client's delayed ACK
319        // before the next frame (notably ReadyForQuery) can leave the server.
320        if let Err(error) = stream.set_nodelay(true) {
321            tracing::warn!(transport = "pg-wire", peer = %peer, err = %error,
322                "failed to configure connection");
323            continue;
324        }
325        let rt = Arc::clone(&runtime);
326        let cfg = Arc::clone(&cfg);
327        let acceptor = tls_acceptor.clone();
328        let peer_str = peer.to_string();
329        tokio::spawn(async move {
330            let _permit = permit;
331            if let Err(e) = handle_connection(stream, rt, cfg, acceptor).await {
332                tracing::warn!(
333                    transport = "pg-wire",
334                    peer = %peer_str,
335                    err = %e,
336                    "connection failed"
337                );
338            }
339        });
340    }
341}
342
343/// Drive one connection's lifetime: SSL negotiation → startup →
344/// authentication → query loop.
345///
346/// The first frame may be `SSLRequest` / `GSSENCRequest` (pre-auth
347/// negotiation) or a plain `Startup`. When a TLS acceptor is present an
348/// `SSLRequest` is answered with `'S'` and the socket is upgraded via a
349/// rustls handshake; the rest of the session (a fresh `Startup`,
350/// authentication, and the query loop) then runs over the encrypted
351/// stream. Without an acceptor — or for `GSSENCRequest`, which we never
352/// support — the request is declined with `'N'` and the client continues
353/// in plaintext on the same socket.
354pub(crate) async fn handle_connection<S>(
355    mut stream: S,
356    runtime: Arc<RedDBRuntime>,
357    config: Arc<PgWireConfig>,
358    tls_acceptor: Option<tokio_rustls::TlsAcceptor>,
359) -> Result<(), PgWireError>
360where
361    S: AsyncRead + AsyncWrite + Unpin + Send,
362{
363    // Bound the whole pre-session negotiation. `read_startup` and the rustls
364    // handshake both block on the peer, and neither had a deadline: a client
365    // that connected and sent nothing (or stalled mid-handshake) pinned a
366    // task and a connection slot indefinitely, before any credential was
367    // presented. The loop below can iterate (SSLRequest/GSSENCRequest are
368    // declined and the client retries), so each read is bounded rather than
369    // the loop as a whole.
370    loop {
371        let startup = crate::wire::admission::with_handshake_deadline(read_startup(&mut stream))
372            .await
373            .map_err(PgWireError::from)??;
374        match startup {
375            FrontendMessage::SslRequest => {
376                if let Some(acceptor) = tls_acceptor.as_ref() {
377                    // 'S' = SSL supported. Perform the rustls handshake and
378                    // hand the encrypted stream to the session driver.
379                    write_raw_byte(&mut stream, b'S').await?;
380                    let tls_stream =
381                        crate::wire::admission::with_handshake_deadline(acceptor.accept(stream))
382                            .await??;
383                    return serve_session_after_tls(tls_stream, runtime, config).await;
384                }
385                // 'N' = not supported — client continues in plaintext and
386                // re-sends a normal Startup on the same socket.
387                write_raw_byte(&mut stream, b'N').await?;
388                continue;
389            }
390            FrontendMessage::GssEncRequest => {
391                // GSSAPI encryption is never offered; decline and let the
392                // client retry with SSLRequest or plaintext Startup.
393                write_raw_byte(&mut stream, b'N').await?;
394                continue;
395            }
396            FrontendMessage::Startup(params) => {
397                return run_pg_session(stream, runtime, config, false, params).await;
398            }
399            FrontendMessage::Unknown { .. } => {
400                // CancelRequest: no response expected; drop the socket.
401                return Ok(());
402            }
403            other => {
404                return Err(PgWireError::Protocol(format!(
405                    "unexpected startup frame: {other:?}"
406                )));
407            }
408        }
409    }
410}
411
412/// Continue the startup flow after a successful TLS upgrade. The client
413/// now sends a fresh `Startup` over the encrypted stream; a second
414/// `SSLRequest` / `GSSENCRequest` cannot be honoured (no nested
415/// encryption) and is declined with `'N'`.
416async fn serve_session_after_tls<S>(
417    mut stream: tokio_rustls::server::TlsStream<S>,
418    runtime: Arc<RedDBRuntime>,
419    config: Arc<PgWireConfig>,
420) -> Result<(), PgWireError>
421where
422    S: AsyncRead + AsyncWrite + Unpin + Send,
423{
424    loop {
425        match read_startup(&mut stream).await? {
426            FrontendMessage::SslRequest | FrontendMessage::GssEncRequest => {
427                write_raw_byte(&mut stream, b'N').await?;
428                continue;
429            }
430            FrontendMessage::Startup(params) => {
431                return run_pg_session(stream, runtime, config, true, params).await;
432            }
433            FrontendMessage::Unknown { .. } => return Ok(()),
434            other => {
435                return Err(PgWireError::Protocol(format!(
436                    "unexpected startup frame: {other:?}"
437                )));
438            }
439        }
440    }
441}
442
443/// Authenticate a `Startup` and drive the query loop over `stream`.
444/// `over_tls` records whether the transport is encrypted, which gates the
445/// cleartext-password posture in [`authenticate_startup`].
446async fn run_pg_session<S>(
447    mut stream: S,
448    runtime: Arc<RedDBRuntime>,
449    config: Arc<PgWireConfig>,
450    over_tls: bool,
451    params: super::protocol::StartupParams,
452) -> Result<(), PgWireError>
453where
454    S: AsyncRead + AsyncWrite + Unpin + Send,
455{
456    let auth_context =
457        match authenticate_startup(&mut stream, &runtime, &config, over_tls, &params).await? {
458            Some(context) => context,
459            // Authentication failed: the error frame was already sent and
460            // the socket must close.
461            None => return Ok(()),
462        };
463
464    let mut prepared: HashMap<String, PgPreparedStatement> = HashMap::new();
465    let mut portals: HashMap<String, PgPortal> = HashMap::new();
466    let prepared_registry = PreparedRegistry::new();
467
468    // Main query loop.
469    loop {
470        let frame = match read_frame(&mut stream).await {
471            Ok(f) => f,
472            Err(PgWireError::Eof) => return Ok(()),
473            Err(e) => return Err(e),
474        };
475
476        // PG sessions are serial, but every session competes for the same
477        // runtime-wide engine budget as HTTP, gRPC and RedWire. Admission is
478        // taken before any parser/planner/storage work. The response writer
479        // releases it before the first socket write, so a slow client cannot
480        // occupy engine capacity while rows drain.
481        let mut execution_permit = if pg_frame_requires_execution(&frame) {
482            match acquire_pg_execution(&runtime).await {
483                Ok(permit) => Some(permit),
484                Err(message) => {
485                    // SQLSTATE 53000 = insufficient_resources. Keep the stable
486                    // RedDB reason code in the message for clients that expose
487                    // only the PostgreSQL ErrorResponse text.
488                    send_error(&mut stream, "53000", message).await?;
489                    write_frame(
490                        &mut stream,
491                        &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
492                    )
493                    .await?;
494                    continue;
495                }
496            }
497        } else {
498            None
499        };
500
501        match frame {
502            FrontendMessage::Query(sql) => {
503                let mut response =
504                    ExecutionResponseWriter::new(&mut stream, execution_permit.take());
505                handle_simple_query(&mut response, &runtime, auth_context.as_ref(), &sql).await?;
506            }
507            FrontendMessage::Parse(msg) => {
508                let mut response =
509                    ExecutionResponseWriter::new(&mut stream, execution_permit.take());
510                handle_parse(
511                    &mut response,
512                    &runtime,
513                    &prepared_registry,
514                    &mut prepared,
515                    &portals,
516                    msg,
517                )
518                .await?;
519            }
520            FrontendMessage::Bind(msg) => {
521                handle_bind(
522                    &mut stream,
523                    &prepared_registry,
524                    &prepared,
525                    &mut portals,
526                    msg,
527                )
528                .await?;
529            }
530            FrontendMessage::Describe(msg) => {
531                let mut response =
532                    ExecutionResponseWriter::new(&mut stream, execution_permit.take());
533                handle_describe(
534                    &mut response,
535                    &runtime,
536                    auth_context.as_ref(),
537                    &prepared_registry,
538                    &prepared,
539                    &mut portals,
540                    msg,
541                )
542                .await?;
543            }
544            FrontendMessage::Execute(msg) => {
545                let mut response =
546                    ExecutionResponseWriter::new(&mut stream, execution_permit.take());
547                handle_execute(
548                    &mut response,
549                    &runtime,
550                    auth_context.as_ref(),
551                    &prepared_registry,
552                    &mut portals,
553                    msg,
554                )
555                .await?;
556            }
557            FrontendMessage::Close(msg) => {
558                handle_close(
559                    &mut stream,
560                    &prepared_registry,
561                    &mut prepared,
562                    &mut portals,
563                    msg,
564                )
565                .await?;
566            }
567            FrontendMessage::Terminate => return Ok(()),
568            FrontendMessage::Flush => {
569                // Frames are written immediately; no additional marker is
570                // needed. ReadyForQuery belongs to Sync, not Flush.
571                continue;
572            }
573            FrontendMessage::Sync => {
574                write_frame(
575                    &mut stream,
576                    &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
577                )
578                .await?;
579            }
580            FrontendMessage::PasswordMessage(_) => {
581                // Should only arrive during auth. Ignore post-auth.
582                continue;
583            }
584            FrontendMessage::Unknown { tag, .. } => {
585                send_error(
586                    &mut stream,
587                    "0A000",
588                    &format!("unsupported frame tag 0x{tag:02x}"),
589                )
590                .await?;
591                write_frame(
592                    &mut stream,
593                    &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
594                )
595                .await?;
596            }
597            other => {
598                send_error(
599                    &mut stream,
600                    "0A000",
601                    &format!("unsupported frame {other:?}"),
602                )
603                .await?;
604                write_frame(
605                    &mut stream,
606                    &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
607                )
608                .await?;
609            }
610        }
611    }
612}
613
614async fn handle_parse<S>(
615    stream: &mut S,
616    runtime: &RedDBRuntime,
617    prepared_registry: &PreparedRegistry,
618    prepared: &mut HashMap<String, PgPreparedStatement>,
619    portals: &HashMap<String, PgPortal>,
620    msg: super::protocol::ParseMessage,
621) -> Result<(), PgWireError>
622where
623    S: AsyncWrite + Unpin,
624{
625    let inferred_param_type_oids = infer_pg_cast_param_type_oids(&msg.query);
626    let sql = rewrite_pg_parameter_casts(&msg.query);
627    let prepared_query = match prepared_registry.prepare(runtime, &sql) {
628        Ok(prepared_query) => prepared_query,
629        Err(err) => {
630            send_error(stream, "42601", &err.to_string()).await?;
631            return Ok(());
632        }
633    };
634    let mut param_type_oids = msg.param_type_oids;
635    if param_type_oids.is_empty() {
636        param_type_oids.resize(prepared_query.parameter_count, PgOid::Unknown.as_u32());
637    }
638    for (idx, oid) in inferred_param_type_oids {
639        if idx >= param_type_oids.len() {
640            param_type_oids.resize(idx + 1, PgOid::Unknown.as_u32());
641        }
642        if param_type_oids[idx] == PgOid::Unknown.as_u32() {
643            param_type_oids[idx] = oid;
644        }
645    }
646    let replaced = prepared.insert(
647        msg.statement,
648        PgPreparedStatement {
649            sql,
650            prepared_id: prepared_query.id,
651            param_type_oids,
652        },
653    );
654    if let Some(replaced) = replaced {
655        release_prepared_if_unreferenced(
656            prepared_registry,
657            replaced.prepared_id,
658            prepared,
659            portals,
660        );
661    }
662    write_frame(stream, &BackendMessage::ParseComplete).await
663}
664
665async fn handle_bind<S>(
666    stream: &mut S,
667    prepared_registry: &PreparedRegistry,
668    prepared: &HashMap<String, PgPreparedStatement>,
669    portals: &mut HashMap<String, PgPortal>,
670    msg: super::protocol::BindMessage,
671) -> Result<(), PgWireError>
672where
673    S: AsyncWrite + Unpin,
674{
675    let Some(stmt) = prepared.get(&msg.statement) else {
676        send_error(
677            stream,
678            "26000",
679            &format!("prepared statement {:?} does not exist", msg.statement),
680        )
681        .await?;
682        return Ok(());
683    };
684    let params = match bind_pg_params(stmt, &msg) {
685        Ok(params) => params,
686        Err(err) => {
687            send_error(stream, "22023", &err).await?;
688            return Ok(());
689        }
690    };
691    let replaced = portals.insert(
692        msg.portal,
693        PgPortal {
694            sql: stmt.sql.clone(),
695            prepared_id: stmt.prepared_id,
696            params,
697            result_format_codes: msg.result_format_codes,
698            row_description_sent: false,
699            result: None,
700            row_offset: 0,
701        },
702    );
703    if let Some(replaced) = replaced {
704        release_prepared_if_unreferenced(
705            prepared_registry,
706            replaced.prepared_id,
707            prepared,
708            portals,
709        );
710    }
711    write_frame(stream, &BackendMessage::BindComplete).await
712}
713
714async fn handle_describe<S>(
715    stream: &mut S,
716    runtime: &RedDBRuntime,
717    auth_context: Option<&PgAuthContext>,
718    prepared_registry: &PreparedRegistry,
719    prepared: &HashMap<String, PgPreparedStatement>,
720    portals: &mut HashMap<String, PgPortal>,
721    msg: super::protocol::DescribeMessage,
722) -> Result<(), PgWireError>
723where
724    S: AsyncWrite + Unpin,
725{
726    match msg.target {
727        DescribeTarget::Statement => {
728            let Some(stmt) = prepared.get(&msg.name) else {
729                send_error(
730                    stream,
731                    "26000",
732                    &format!("prepared statement {:?} does not exist", msg.name),
733                )
734                .await?;
735                return Ok(());
736            };
737            write_frame(
738                stream,
739                &BackendMessage::ParameterDescription(stmt.param_type_oids.clone()),
740            )
741            .await?;
742            if is_ask_query(&stmt.sql) {
743                emit_ask_row_description(stream).await
744            } else {
745                write_frame(stream, &BackendMessage::NoData).await
746            }
747        }
748        DescribeTarget::Portal => {
749            let Some(portal) = portals.get_mut(&msg.name) else {
750                send_error(
751                    stream,
752                    "34000",
753                    &format!("portal {:?} does not exist", msg.name),
754                )
755                .await?;
756                return Ok(());
757            };
758            if is_ask_query(&portal.sql) {
759                emit_ask_row_description(stream).await?;
760                portal.row_description_sent = true;
761                Ok(())
762            } else if is_row_returning_query(&portal.sql) {
763                let result = match execute_pg_query_result(
764                    runtime,
765                    auth_context,
766                    prepared_registry,
767                    portal.prepared_id,
768                    &portal.sql,
769                    &portal.params,
770                ) {
771                    Ok(result) => result,
772                    Err(err) => {
773                        let code = classify_sqlstate(&err);
774                        send_error(stream, code, &err).await?;
775                        return Ok(());
776                    }
777                };
778                emit_row_description_for_result(stream, &result).await?;
779                portal.row_description_sent = true;
780                portal.result = Some(result);
781                portal.row_offset = 0;
782                Ok(())
783            } else {
784                write_frame(stream, &BackendMessage::NoData).await
785            }
786        }
787    }
788}
789
790async fn handle_execute<S>(
791    stream: &mut S,
792    runtime: &RedDBRuntime,
793    auth_context: Option<&PgAuthContext>,
794    prepared_registry: &PreparedRegistry,
795    portals: &mut HashMap<String, PgPortal>,
796    msg: super::protocol::ExecuteMessage,
797) -> Result<(), PgWireError>
798where
799    S: AsyncWrite + Unpin,
800{
801    let Some(portal) = portals.get_mut(&msg.portal) else {
802        send_error(
803            stream,
804            "34000",
805            &format!("portal {:?} does not exist", msg.portal),
806        )
807        .await?;
808        return Ok(());
809    };
810    match prepared_registry.ddl_epoch(portal.prepared_id) {
811        Some(epoch) if epoch != runtime.ddl_epoch() => {
812            send_error(stream, "XX000", "prepared_needs_replan").await?;
813            return Ok(());
814        }
815        Some(_) => {}
816        None => {
817            send_error(stream, "26000", "prepared statement does not exist").await?;
818            return Ok(());
819        }
820    }
821    if portal.result.is_none() {
822        match execute_pg_query_result(
823            runtime,
824            auth_context,
825            prepared_registry,
826            portal.prepared_id,
827            &portal.sql,
828            &portal.params,
829        ) {
830            Ok(result) => {
831                portal.result = Some(result);
832                portal.row_offset = 0;
833            }
834            Err(err) => {
835                let code = classify_sqlstate(&err);
836                send_error(stream, code, &err).await?;
837                return Ok(());
838            }
839        }
840    }
841
842    let result = portal.result.as_ref().expect("portal result populated");
843    match emit_portal_execution(
844        stream,
845        result,
846        portal.row_description_sent,
847        portal.row_offset,
848        msg.max_rows,
849    )
850    .await
851    {
852        Ok(PortalExecution::Complete) => {
853            portal.result = None;
854            portal.row_offset = 0;
855            portal.row_description_sent = false;
856            Ok(())
857        }
858        Ok(PortalExecution::Suspended { next_row_offset }) => {
859            portal.row_offset = next_row_offset;
860            portal.row_description_sent = true;
861            Ok(())
862        }
863        Err(err) => {
864            let code = classify_sqlstate(&err);
865            send_error(stream, code, &err).await
866        }
867    }
868}
869
870async fn handle_close<S>(
871    stream: &mut S,
872    prepared_registry: &PreparedRegistry,
873    prepared: &mut HashMap<String, PgPreparedStatement>,
874    portals: &mut HashMap<String, PgPortal>,
875    msg: super::protocol::CloseMessage,
876) -> Result<(), PgWireError>
877where
878    S: AsyncWrite + Unpin,
879{
880    let prepared_id = match msg.target {
881        DescribeTarget::Statement => prepared.remove(&msg.name).map(|stmt| stmt.prepared_id),
882        DescribeTarget::Portal => portals.remove(&msg.name).map(|portal| portal.prepared_id),
883    };
884    if let Some(prepared_id) = prepared_id {
885        release_prepared_if_unreferenced(prepared_registry, prepared_id, prepared, portals);
886    }
887    write_frame(stream, &BackendMessage::CloseComplete).await
888}
889
890fn release_prepared_if_unreferenced(
891    prepared_registry: &PreparedRegistry,
892    prepared_id: PreparedId,
893    prepared: &HashMap<String, PgPreparedStatement>,
894    portals: &HashMap<String, PgPortal>,
895) {
896    let statement_references_id = prepared
897        .values()
898        .any(|statement| statement.prepared_id == prepared_id);
899    let portal_references_id = portals
900        .values()
901        .any(|portal| portal.prepared_id == prepared_id);
902    if !statement_references_id && !portal_references_id {
903        prepared_registry.release(prepared_id);
904    }
905}
906
907fn bind_pg_params(
908    stmt: &PgPreparedStatement,
909    msg: &super::protocol::BindMessage,
910) -> Result<Vec<ParamValue>, String> {
911    if !matches!(msg.param_format_codes.len(), 0 | 1)
912        && msg.param_format_codes.len() != msg.params.len()
913    {
914        return Err("Bind format count must be 0, 1, or match parameter count".to_string());
915    }
916    msg.params
917        .iter()
918        .enumerate()
919        .map(|(idx, param)| {
920            let oid = stmt
921                .param_type_oids
922                .get(idx)
923                .copied()
924                .unwrap_or(PgOid::Unknown.as_u32());
925            let format_code = match msg.param_format_codes.as_slice() {
926                [] => 0,
927                [format] => *format,
928                formats => formats[idx],
929            };
930            pg_param_to_param_value(oid, format_code, param.as_deref())
931        })
932        .collect()
933}
934
935fn execute_pg_query_result(
936    runtime: &RedDBRuntime,
937    auth_context: Option<&PgAuthContext>,
938    prepared_registry: &PreparedRegistry,
939    prepared_id: PreparedId,
940    sql: &str,
941    params: &[ParamValue],
942) -> Result<crate::runtime::OwnedQueryResult, String> {
943    let executor = QueryRequestExecutor::new(runtime, prepared_registry);
944    let bound = executor
945        .bind(QueryRequest::prepared(prepared_id, params.to_vec()))
946        .map_err(|err| err.to_string())?;
947    if pg_scalar_select_param_index(sql).is_some() {
948        let values = params.iter().cloned().map(Value::from).collect::<Vec<_>>();
949        if let Some(result) = try_execute_pg_scalar_select(sql, &values) {
950            return runtime
951                .own_query_result(result)
952                .map_err(|error| error.to_string());
953        }
954    }
955    if params.is_empty() {
956        return match translate_pg_catalog_query(runtime, sql) {
957            Ok(Some(result)) => runtime
958                .own_query_result(crate::runtime::RuntimeQueryResult {
959                    query: sql.to_string(),
960                    mode: reddb_rql::modes::QueryMode::Sql,
961                    statement: "select",
962                    engine: "pg-catalog",
963                    result,
964                    affected_rows: 0,
965                    statement_type: "select",
966                    bookmark: None,
967                    notice: None,
968                })
969                .map_err(|error| error.to_string()),
970            Ok(None) => run_runtime_blocking(|| {
971                execute_with_pg_auth_context(auth_context, || {
972                    // Utility statements use the full dispatcher after registry validation.
973                    // Their zero-parameter SQL can be reparsed without losing bindings;
974                    // data queries retain the cached prepared shape on repeated Execute.
975                    if matches!(
976                        bound.prepared_expr(),
977                        Some(
978                            QueryExpr::Function(_)
979                                | QueryExpr::Table(_)
980                                | QueryExpr::Graph(_)
981                                | QueryExpr::Path(_)
982                                | QueryExpr::Join(_)
983                                | QueryExpr::Vector(_)
984                                | QueryExpr::Hybrid(_)
985                                | QueryExpr::Insert(_)
986                                | QueryExpr::Update(_)
987                                | QueryExpr::Delete(_)
988                                | QueryExpr::QueueCommand(_)
989                                | QueryExpr::KvCommand(_)
990                                | QueryExpr::SearchCommand(_)
991                                | QueryExpr::Ask(_)
992                        )
993                    ) {
994                        executor.execute_bound(bound)
995                    } else {
996                        runtime
997                            .execute_query(sql)
998                            .and_then(|result| runtime.own_query_result(result))
999                    }
1000                })
1001            })
1002            .map_err(|err| err.to_string()),
1003            Err(err) => Err(err.to_string()),
1004        };
1005    }
1006
1007    run_runtime_blocking(|| {
1008        execute_with_pg_auth_context(auth_context, || executor.execute_bound(bound))
1009    })
1010    .map_err(|err| err.to_string())
1011}
1012
1013fn try_execute_pg_scalar_select(
1014    sql: &str,
1015    params: &[Value],
1016) -> Option<crate::runtime::RuntimeQueryResult> {
1017    let index = pg_scalar_select_param_index(sql)?;
1018    let value = params.get(index)?.clone();
1019    let mut result = UnifiedResult::with_columns(vec!["?column?".to_string()]);
1020    let mut record = UnifiedRecord::new();
1021    record.set("?column?", value);
1022    result.push(record);
1023    Some(crate::runtime::RuntimeQueryResult {
1024        query: sql.to_string(),
1025        mode: reddb_rql::modes::QueryMode::Sql,
1026        statement: "select",
1027        engine: "pg-wire",
1028        result,
1029        affected_rows: 0,
1030        statement_type: "select",
1031        bookmark: None,
1032        notice: None,
1033    })
1034}
1035
1036fn pg_scalar_select_param_index(sql: &str) -> Option<usize> {
1037    let trimmed = sql.trim().trim_end_matches(';').trim();
1038    let lower = trimmed.to_ascii_lowercase();
1039    let body = lower.strip_prefix("select ")?;
1040    let param = if let Some(inner) = body.strip_prefix("cast(") {
1041        let end = inner.find(" as ")?;
1042        &inner[..end]
1043    } else {
1044        body.split_whitespace().next()?
1045    };
1046    let digits = param.strip_prefix('$')?;
1047    let n = digits.parse::<usize>().ok()?;
1048    n.checked_sub(1)
1049}
1050
1051fn rewrite_pg_parameter_casts(sql: &str) -> String {
1052    let mut out = String::with_capacity(sql.len());
1053    let bytes = sql.as_bytes();
1054    let mut cursor = 0;
1055    let mut pos = 0;
1056    while pos < bytes.len() {
1057        if bytes[pos] != b'$' {
1058            pos += 1;
1059            continue;
1060        }
1061        let param_start = pos;
1062        pos += 1;
1063        let digits_start = pos;
1064        while pos < bytes.len() && bytes[pos].is_ascii_digit() {
1065            pos += 1;
1066        }
1067        if digits_start == pos {
1068            continue;
1069        }
1070        if pos + 2 <= bytes.len() && &bytes[pos..pos + 2] == b"::" {
1071            let param_end = pos;
1072            pos += 2;
1073            let type_start = pos;
1074            while pos < bytes.len()
1075                && (bytes[pos].is_ascii_alphanumeric() || matches!(bytes[pos], b'_' | b'.'))
1076            {
1077                pos += 1;
1078            }
1079            if type_start != pos {
1080                out.push_str(&sql[cursor..param_start]);
1081                out.push_str(&sql[param_start..param_end]);
1082                cursor = pos;
1083                continue;
1084            }
1085        }
1086    }
1087    out.push_str(&sql[cursor..]);
1088    out
1089}
1090
1091fn infer_pg_cast_param_type_oids(sql: &str) -> Vec<(usize, u32)> {
1092    let mut out = Vec::new();
1093    let bytes = sql.as_bytes();
1094    let mut pos = 0;
1095    while pos < bytes.len() {
1096        if bytes[pos] != b'$' {
1097            pos += 1;
1098            continue;
1099        }
1100        pos += 1;
1101        let digits_start = pos;
1102        while pos < bytes.len() && bytes[pos].is_ascii_digit() {
1103            pos += 1;
1104        }
1105        if digits_start == pos {
1106            continue;
1107        }
1108        let Some(param_index) = sql[digits_start..pos]
1109            .parse::<usize>()
1110            .ok()
1111            .and_then(|idx| idx.checked_sub(1))
1112        else {
1113            continue;
1114        };
1115        if pos + 2 > bytes.len() || &bytes[pos..pos + 2] != b"::" {
1116            continue;
1117        }
1118        pos += 2;
1119        let type_start = pos;
1120        while pos < bytes.len()
1121            && (bytes[pos].is_ascii_alphanumeric() || matches!(bytes[pos], b'_' | b'.'))
1122        {
1123            pos += 1;
1124        }
1125        if type_start == pos {
1126            continue;
1127        }
1128        if let Some(oid) = pg_cast_type_oid(&sql[type_start..pos]) {
1129            out.push((param_index, oid));
1130        }
1131    }
1132    out
1133}
1134
1135fn pg_cast_type_oid(ty: &str) -> Option<u32> {
1136    let lower = ty.to_ascii_lowercase();
1137    let short = lower.rsplit('.').next().unwrap_or(lower.as_str());
1138    let oid = match short {
1139        "bool" | "boolean" => PgOid::Bool,
1140        "int2" | "smallint" => PgOid::Int2,
1141        "int" | "int4" | "integer" => PgOid::Int4,
1142        "int8" | "bigint" => PgOid::Int8,
1143        "float4" | "real" => PgOid::Float4,
1144        "float8" | "double" | "doubleprecision" => PgOid::Float8,
1145        "numeric" | "decimal" => PgOid::Numeric,
1146        "bytea" => PgOid::Bytea,
1147        "json" => PgOid::Json,
1148        "jsonb" => PgOid::Jsonb,
1149        "text" => PgOid::Text,
1150        "varchar" | "character varying" => PgOid::Varchar,
1151        "uuid" => PgOid::Uuid,
1152        "timestamp" => PgOid::Timestamp,
1153        "timestamptz" | "timestampz" => PgOid::TimestampTz,
1154        "vector" => PgOid::Vector,
1155        _ => return None,
1156    };
1157    Some(oid.as_u32())
1158}
1159
1160fn is_row_returning_query(sql: &str) -> bool {
1161    let trimmed = sql.trim_start().to_ascii_lowercase();
1162    trimmed.starts_with("select")
1163        || trimmed.starts_with("with")
1164        || trimmed.starts_with("ask")
1165        || trimmed.starts_with("search")
1166        || trimmed.starts_with("vector")
1167        || trimmed.starts_with("hybrid")
1168}
1169
1170fn is_ask_query(sql: &str) -> bool {
1171    sql.trim_start().to_ascii_lowercase().starts_with("ask")
1172}
1173
1174/// Drive the PG-Wire authentication handshake after a `StartupMessage`.
1175///
1176/// Posture:
1177///   * Auth store enabled → request `AuthenticationCleartextPassword`,
1178///     read the client's `PasswordMessage`, and verify it against the same
1179///     auth store the native transports use. Success installs the resolved
1180///     identity; failure emits `28P01` and closes the socket.
1181///   * No auth store enabled + loopback bind → legacy trust: send
1182///     `AuthenticationOk` with no credentials required (`None` context).
1183///   * No auth store enabled + non-loopback bind → refuse with `28P01`;
1184///     an unconfigured listener must never grant unauthenticated access
1185///     off loopback.
1186///
1187/// TLS posture: a cleartext password may only cross an *unencrypted* link
1188/// on a loopback bind. On a non-loopback bind we refuse to even challenge
1189/// for a password unless the connection is already over TLS (`over_tls`),
1190/// so a client using `sslmode=disable` against a remote listener never
1191/// leaks its password in the clear. Loopback and TLS connections are
1192/// challenged as before.
1193///
1194/// Returns `Ok(Some(context))` once the connection is ready for queries
1195/// (`context` is `None` for trust connections), or `Ok(None)` when
1196/// authentication failed and the error frame has already been sent.
1197///
1198/// Failures give the same `28P01` "password authentication failed" for a
1199/// wrong password and a missing user, so there is no user-existence oracle.
1200async fn authenticate_startup<S>(
1201    stream: &mut S,
1202    runtime: &RedDBRuntime,
1203    config: &PgWireConfig,
1204    over_tls: bool,
1205    params: &super::protocol::StartupParams,
1206) -> Result<Option<Option<PgAuthContext>>, PgWireError>
1207where
1208    S: AsyncRead + AsyncWrite + Unpin,
1209{
1210    let auth_store = runtime.auth_store().filter(|store| store.is_enabled());
1211    let Some(auth_store) = auth_store else {
1212        // No credentials configured: trust survives for loopback only.
1213        if is_loopback_bind(&config.bind_addr) {
1214            send_startup_ready(stream, config, params).await?;
1215            return Ok(Some(None));
1216        }
1217        send_error(stream, "28P01", "password authentication failed").await?;
1218        return Ok(None);
1219    };
1220
1221    // Refuse to challenge for a cleartext password over an unencrypted
1222    // non-loopback link: it would traverse the network in the clear.
1223    // Clients must connect with `sslmode=require` (TLS) to authenticate
1224    // against a remote bind.
1225    if !over_tls && !is_loopback_bind(&config.bind_addr) {
1226        send_error(
1227            stream,
1228            "28000",
1229            "cleartext password authentication requires TLS on non-loopback binds",
1230        )
1231        .await?;
1232        return Ok(None);
1233    }
1234
1235    write_frame(stream, &BackendMessage::AuthenticationCleartextPassword).await?;
1236    let password = match read_frame(stream).await? {
1237        FrontendMessage::PasswordMessage(payload) => password_from_message(&payload),
1238        // Anything other than a PasswordMessage is a protocol violation
1239        // here; treat it as an auth failure rather than leaking detail.
1240        _ => None,
1241    };
1242    let username = params.get("user").unwrap_or("");
1243    let tenant = params.get("tenant");
1244    let Some(password) = password else {
1245        send_error(stream, "28P01", "password authentication failed").await?;
1246        return Ok(None);
1247    };
1248    // A missing user and a wrong password both surface as the same error
1249    // class below — no user-existence oracle leaks over the wire.
1250    let session = match auth_store.authenticate_in_tenant(tenant, username, &password) {
1251        Ok(session) => session,
1252        Err(_) => {
1253            send_error(stream, "28P01", "password authentication failed").await?;
1254            return Ok(None);
1255        }
1256    };
1257
1258    send_startup_ready(stream, config, params).await?;
1259    Ok(Some(Some(PgAuthContext {
1260        username: session.username,
1261        role: session.role,
1262        tenant: session.tenant_id,
1263    })))
1264}
1265
1266async fn send_startup_ready<S>(
1267    stream: &mut S,
1268    config: &PgWireConfig,
1269    params: &super::protocol::StartupParams,
1270) -> Result<(), PgWireError>
1271where
1272    S: AsyncWrite + Unpin,
1273{
1274    // Authentication has succeeded (or trust applies): signal AuthenticationOk.
1275    write_frame(stream, &BackendMessage::AuthenticationOk).await?;
1276
1277    // Standard ParameterStatus frames. Drivers gate capabilities on these.
1278    for (name, value) in [
1279        ("server_version", config.server_version.as_str()),
1280        ("server_encoding", "UTF8"),
1281        ("client_encoding", "UTF8"),
1282        ("DateStyle", "ISO, MDY"),
1283        ("TimeZone", "UTC"),
1284        ("integer_datetimes", "on"),
1285        ("standard_conforming_strings", "on"),
1286        (
1287            "application_name",
1288            params.get("application_name").unwrap_or(""),
1289        ),
1290    ] {
1291        write_frame(
1292            stream,
1293            &BackendMessage::ParameterStatus {
1294                name: name.to_string(),
1295                value: value.to_string(),
1296            },
1297        )
1298        .await?;
1299    }
1300
1301    // BackendKeyData: (pid, secret_key). CancelRequest is not honoured
1302    // yet, but the key is what a future cancel path authenticates on, so
1303    // it is random per connection rather than a constant a client could
1304    // learn once and reuse against every session.
1305    write_frame(
1306        stream,
1307        &BackendMessage::BackendKeyData {
1308            pid: std::process::id(),
1309            key: random_cancel_key(),
1310        },
1311    )
1312    .await?;
1313
1314    write_frame(
1315        stream,
1316        &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
1317    )
1318    .await?;
1319    Ok(())
1320}
1321
1322async fn handle_simple_query<S>(
1323    stream: &mut S,
1324    runtime: &RedDBRuntime,
1325    auth_context: Option<&PgAuthContext>,
1326    sql: &str,
1327) -> Result<(), PgWireError>
1328where
1329    S: AsyncWrite + Unpin,
1330{
1331    // Empty query convention: PG emits EmptyQueryResponse instead of a
1332    // CommandComplete. Some clients (psql `\;`) rely on this.
1333    if sql.trim().is_empty() {
1334        write_frame(stream, &BackendMessage::EmptyQueryResponse).await?;
1335        write_frame(
1336            stream,
1337            &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
1338        )
1339        .await?;
1340        return Ok(());
1341    }
1342
1343    if let Some(tag) = pg_session_compat_command_tag(sql) {
1344        write_frame(stream, &BackendMessage::CommandComplete(tag.to_string())).await?;
1345        write_frame(
1346            stream,
1347            &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
1348        )
1349        .await?;
1350        return Ok(());
1351    }
1352
1353    let query_result = match translate_pg_catalog_query(runtime, sql) {
1354        Ok(Some(result)) => runtime.own_query_result(crate::runtime::RuntimeQueryResult {
1355            query: sql.to_string(),
1356            mode: reddb_rql::modes::QueryMode::Sql,
1357            statement: "select",
1358            engine: "pg-catalog",
1359            result,
1360            affected_rows: 0,
1361            statement_type: "select",
1362            bookmark: None,
1363            notice: None,
1364        }),
1365        Ok(None) => run_runtime_blocking(|| {
1366            execute_with_pg_auth_context(auth_context, || {
1367                runtime
1368                    .execute_query(sql)
1369                    .and_then(|result| runtime.own_query_result(result))
1370            })
1371        }),
1372        Err(err) => Err(err),
1373    };
1374
1375    match query_result {
1376        Ok(result) => {
1377            emit_success_result(stream, &result).await?;
1378        }
1379        Err(err) => {
1380            // PG SQLSTATE class 42 covers syntax / binding errors; we use
1381            // 42P01 (undefined_table) and 42601 (syntax_error) when we can
1382            // detect; otherwise fall back to XX000 (internal error).
1383            let code = classify_sqlstate(&err.to_string());
1384            send_error(stream, code, &crate::api::client_facing_message(&err)).await?;
1385        }
1386    }
1387
1388    write_frame(
1389        stream,
1390        &BackendMessage::ReadyForQuery(TransactionStatus::Idle),
1391    )
1392    .await?;
1393    Ok(())
1394}
1395
1396fn pg_session_compat_command_tag(sql: &str) -> Option<&'static str> {
1397    let lower = sql.trim().trim_end_matches(';').to_ascii_lowercase();
1398    if lower.starts_with("set ") {
1399        return Some("SET");
1400    }
1401    None
1402}
1403
1404async fn emit_success_result<S>(
1405    stream: &mut S,
1406    result: &crate::runtime::RuntimeQueryResult,
1407) -> Result<(), PgWireError>
1408where
1409    S: AsyncWrite + Unpin,
1410{
1411    if result.statement == "ask" {
1412        emit_ask_result_row(stream, result).await?;
1413        write_frame(
1414            stream,
1415            &BackendMessage::CommandComplete("SELECT 1".to_string()),
1416        )
1417        .await?;
1418    } else if result_returns_rows(result) {
1419        emit_result_rows(stream, &result.result).await?;
1420        write_frame(
1421            stream,
1422            &BackendMessage::CommandComplete(format!("SELECT {}", result.result.records.len())),
1423        )
1424        .await?;
1425    } else {
1426        // DDL / DML / config statements: echo the runtime's
1427        // high-level statement tag back. PG format is
1428        // "<CMD> [<OID>] <COUNT>"; we keep the count where
1429        // applicable and fall back to the runtime's message.
1430        let tag = match result.statement_type {
1431            "insert" => format!("INSERT 0 {}", result.affected_rows),
1432            "update" => format!("UPDATE {}", result.affected_rows),
1433            "delete" => format!("DELETE {}", result.affected_rows),
1434            other => other.to_uppercase(),
1435        };
1436        write_frame(stream, &BackendMessage::CommandComplete(tag)).await?;
1437    }
1438    Ok(())
1439}
1440
1441async fn emit_success_result_without_row_description<S>(
1442    stream: &mut S,
1443    result: &crate::runtime::RuntimeQueryResult,
1444) -> Result<(), PgWireError>
1445where
1446    S: AsyncWrite + Unpin,
1447{
1448    if result.statement == "ask" {
1449        let row = ask_query_result_to_pg_wire_row(result)
1450            .ok_or_else(|| PgWireError::Protocol("ASK result missing row body".to_string()))?;
1451        write_frame(stream, &BackendMessage::DataRow(row.cells)).await?;
1452        write_frame(
1453            stream,
1454            &BackendMessage::CommandComplete("SELECT 1".to_string()),
1455        )
1456        .await?;
1457    } else if result_returns_rows(result) {
1458        emit_result_data_rows(stream, &result.result).await?;
1459        write_frame(
1460            stream,
1461            &BackendMessage::CommandComplete(format!("SELECT {}", result.result.records.len())),
1462        )
1463        .await?;
1464    } else {
1465        let tag = match result.statement_type {
1466            "insert" => format!("INSERT 0 {}", result.affected_rows),
1467            "update" => format!("UPDATE {}", result.affected_rows),
1468            "delete" => format!("DELETE {}", result.affected_rows),
1469            other => other.to_uppercase(),
1470        };
1471        write_frame(stream, &BackendMessage::CommandComplete(tag)).await?;
1472    }
1473    Ok(())
1474}
1475
1476enum PortalExecution {
1477    Complete,
1478    Suspended { next_row_offset: usize },
1479}
1480
1481async fn emit_portal_execution<S>(
1482    stream: &mut S,
1483    result: &crate::runtime::RuntimeQueryResult,
1484    row_description_sent: bool,
1485    row_offset: usize,
1486    max_rows: u32,
1487) -> Result<PortalExecution, String>
1488where
1489    S: AsyncWrite + Unpin,
1490{
1491    if result.statement == "ask" && row_description_sent {
1492        emit_success_result_without_row_description(stream, result)
1493            .await
1494            .map_err(|err| err.to_string())?;
1495        Ok(PortalExecution::Complete)
1496    } else if result.statement == "ask" {
1497        emit_success_result(stream, result)
1498            .await
1499            .map_err(|err| err.to_string())?;
1500        Ok(PortalExecution::Complete)
1501    } else if result_returns_rows(result) {
1502        emit_row_returning_portal_execution(
1503            stream,
1504            result,
1505            row_description_sent,
1506            row_offset,
1507            max_rows,
1508        )
1509        .await
1510    } else if row_description_sent {
1511        emit_success_result_without_row_description(stream, result)
1512            .await
1513            .map_err(|err| err.to_string())?;
1514        Ok(PortalExecution::Complete)
1515    } else {
1516        emit_success_result(stream, result)
1517            .await
1518            .map_err(|err| err.to_string())?;
1519        Ok(PortalExecution::Complete)
1520    }
1521}
1522
1523async fn emit_row_returning_portal_execution<S>(
1524    stream: &mut S,
1525    result: &crate::runtime::RuntimeQueryResult,
1526    row_description_sent: bool,
1527    row_offset: usize,
1528    max_rows: u32,
1529) -> Result<PortalExecution, String>
1530where
1531    S: AsyncWrite + Unpin,
1532{
1533    if !row_description_sent {
1534        emit_result_row_description(stream, &result.result)
1535            .await
1536            .map_err(|err| err.to_string())?;
1537    }
1538
1539    let total_rows = result.result.records.len();
1540    let remaining_rows = total_rows.saturating_sub(row_offset);
1541    let rows_to_emit = if max_rows == 0 {
1542        remaining_rows
1543    } else {
1544        remaining_rows.min(max_rows as usize)
1545    };
1546    emit_result_data_rows_range(stream, &result.result, row_offset, rows_to_emit)
1547        .await
1548        .map_err(|err| err.to_string())?;
1549
1550    let next_row_offset = row_offset + rows_to_emit;
1551    if next_row_offset < total_rows {
1552        write_frame(stream, &BackendMessage::PortalSuspended)
1553            .await
1554            .map_err(|err| err.to_string())?;
1555        Ok(PortalExecution::Suspended { next_row_offset })
1556    } else {
1557        write_frame(
1558            stream,
1559            &BackendMessage::CommandComplete(format!("SELECT {}", total_rows)),
1560        )
1561        .await
1562        .map_err(|err| err.to_string())?;
1563        Ok(PortalExecution::Complete)
1564    }
1565}
1566
1567async fn emit_row_description_for_result<S>(
1568    stream: &mut S,
1569    result: &crate::runtime::RuntimeQueryResult,
1570) -> Result<(), PgWireError>
1571where
1572    S: AsyncWrite + Unpin,
1573{
1574    if result.statement == "ask" {
1575        emit_ask_row_description(stream).await
1576    } else if result_returns_rows(result) {
1577        emit_result_row_description(stream, &result.result).await
1578    } else {
1579        write_frame(stream, &BackendMessage::NoData).await
1580    }
1581}
1582
1583fn result_returns_rows(result: &crate::runtime::RuntimeQueryResult) -> bool {
1584    result.statement_type == "select"
1585}
1586
1587async fn emit_result_rows<S>(
1588    stream: &mut S,
1589    result: &crate::storage::query::unified::UnifiedResult,
1590) -> Result<(), PgWireError>
1591where
1592    S: AsyncWrite + Unpin,
1593{
1594    emit_result_row_description(stream, result).await?;
1595    emit_result_data_rows(stream, result).await
1596}
1597
1598async fn emit_result_row_description<S>(
1599    stream: &mut S,
1600    result: &crate::storage::query::unified::UnifiedResult,
1601) -> Result<(), PgWireError>
1602where
1603    S: AsyncWrite + Unpin,
1604{
1605    // RowDescription: derived from the first record's column ordering.
1606    // When `result.columns` is non-empty we honour that order; otherwise
1607    // we synthesise one from the record's field order.
1608    let columns: Vec<String> = if !result.columns.is_empty() {
1609        result.columns.clone()
1610    } else if let Some(first) = result.records.first() {
1611        record_field_names(first)
1612    } else {
1613        Vec::new()
1614    };
1615
1616    // Peek at the first record for per-column type OIDs. When there's no
1617    // data row we fall back to TEXT for every column — clients render
1618    // empty result sets happily.
1619    let type_oids: Vec<PgOid> = columns
1620        .iter()
1621        .map(|col| {
1622            result
1623                .records
1624                .first()
1625                .and_then(|r| record_get(r, col))
1626                .map(PgOid::from_value)
1627                .unwrap_or(PgOid::Text)
1628        })
1629        .collect();
1630
1631    let descriptors: Vec<ColumnDescriptor> = columns
1632        .iter()
1633        .zip(type_oids.iter())
1634        .map(|(name, oid)| ColumnDescriptor {
1635            name: name.clone(),
1636            table_oid: 0,
1637            column_attr: 0,
1638            type_oid: oid.as_u32(),
1639            type_size: -1,
1640            type_mod: -1,
1641            format: 0,
1642        })
1643        .collect();
1644
1645    write_frame(stream, &BackendMessage::RowDescription(descriptors)).await
1646}
1647
1648async fn emit_result_data_rows<S>(
1649    stream: &mut S,
1650    result: &crate::storage::query::unified::UnifiedResult,
1651) -> Result<(), PgWireError>
1652where
1653    S: AsyncWrite + Unpin,
1654{
1655    emit_result_data_rows_range(stream, result, 0, result.records.len()).await
1656}
1657
1658async fn emit_result_data_rows_range<S>(
1659    stream: &mut S,
1660    result: &crate::storage::query::unified::UnifiedResult,
1661    row_offset: usize,
1662    row_count: usize,
1663) -> Result<(), PgWireError>
1664where
1665    S: AsyncWrite + Unpin,
1666{
1667    let columns: Vec<String> = if !result.columns.is_empty() {
1668        result.columns.clone()
1669    } else if let Some(first) = result.records.first() {
1670        record_field_names(first)
1671    } else {
1672        Vec::new()
1673    };
1674    for record in result.records.iter().skip(row_offset).take(row_count) {
1675        let fields: Vec<Option<Vec<u8>>> = columns
1676            .iter()
1677            .map(|col| record_get(record, col).and_then(value_to_pg_wire_bytes))
1678            .collect();
1679        write_frame(stream, &BackendMessage::DataRow(fields)).await?;
1680    }
1681
1682    Ok(())
1683}
1684
1685async fn emit_ask_result_row<S>(
1686    stream: &mut S,
1687    result: &crate::runtime::RuntimeQueryResult,
1688) -> Result<(), PgWireError>
1689where
1690    S: AsyncWrite + Unpin,
1691{
1692    let row = ask_query_result_to_pg_wire_row(result)
1693        .ok_or_else(|| PgWireError::Protocol("ASK result missing row body".to_string()))?;
1694
1695    emit_ask_row_description(stream).await?;
1696    write_frame(stream, &BackendMessage::DataRow(row.cells)).await?;
1697    Ok(())
1698}
1699
1700async fn emit_ask_row_description<S>(stream: &mut S) -> Result<(), PgWireError>
1701where
1702    S: AsyncWrite + Unpin,
1703{
1704    let descriptors: Vec<ColumnDescriptor> = crate::runtime::ai::pg_wire_ask_row_encoder::columns()
1705        .iter()
1706        .map(|col| ColumnDescriptor {
1707            name: col.name.to_string(),
1708            table_oid: 0,
1709            column_attr: 0,
1710            type_oid: col.oid.as_u32(),
1711            type_size: -1,
1712            type_mod: -1,
1713            format: 0,
1714        })
1715        .collect();
1716    write_frame(stream, &BackendMessage::RowDescription(descriptors)).await
1717}
1718
1719fn ask_query_result_to_pg_wire_row(
1720    result: &crate::runtime::RuntimeQueryResult,
1721) -> Option<crate::runtime::ai::pg_wire_ask_row_encoder::AskRow> {
1722    if result.statement != "ask" {
1723        return None;
1724    }
1725    let ask = crate::presentation::query_result::ask_result_from_unified_result(&result.result)?;
1726
1727    Some(crate::runtime::ai::pg_wire_ask_row_encoder::encode(&ask))
1728}
1729
1730/// Best-effort field lookup on a `UnifiedRecord`. The record API lives in
1731/// `storage::query::unified` and today uses `HashMap<String, Value>` under
1732/// the hood — we use `get` if it exists, else fall back to serialised map.
1733fn record_get<'a>(record: &'a UnifiedRecord, key: &str) -> Option<&'a Value> {
1734    record.get(key)
1735}
1736
1737/// Extract column names in iteration order from a single record. When
1738/// the caller didn't supply an explicit `columns` projection we use the
1739/// first record's field ordering as the canonical tuple shape.
1740///
1741/// HashMap iteration order is non-deterministic — for Phase 3.1 we
1742/// accept the shuffle since PG clients receive the ordered header via
1743/// RowDescription and match cells positionally. A stable ordering
1744/// would require keeping an insertion-order index alongside `values`.
1745fn record_field_names(record: &UnifiedRecord) -> Vec<String> {
1746    // `column_names()` merges the columnar scan side-channel with
1747    // the HashMap so scan rows (which populate only columnar) still
1748    // surface their field names in PG wire output.
1749    record
1750        .column_names()
1751        .into_iter()
1752        .map(|k| k.to_string())
1753        .collect()
1754}
1755
1756async fn send_error<S>(stream: &mut S, code: &str, message: &str) -> Result<(), PgWireError>
1757where
1758    S: AsyncWrite + Unpin,
1759{
1760    write_frame(
1761        stream,
1762        &BackendMessage::ErrorResponse {
1763            severity: "ERROR".to_string(),
1764            code: code.to_string(),
1765            message: message.to_string(),
1766        },
1767    )
1768    .await
1769}
1770
1771/// Heuristically map a runtime error message onto a PG SQLSTATE. Full
1772/// coverage would map every `RedDBError` variant; this is enough for the
1773/// common psql / JDBC paths.
1774fn classify_sqlstate(msg: &str) -> &'static str {
1775    let lower = msg.to_ascii_lowercase();
1776    if lower.contains("not found") || lower.contains("does not exist") {
1777        // 42P01 undefined_table; close enough for collection-not-found.
1778        "42P01"
1779    } else if lower.contains("parse") || lower.contains("expected") || lower.contains("syntax") {
1780        "42601"
1781    } else if lower.contains("already exists") {
1782        "42P07"
1783    } else if lower.contains("permission") || lower.contains("auth") {
1784        "28000"
1785    } else {
1786        "XX000"
1787    }
1788}
1789
1790#[cfg(test)]
1791mod tests {
1792    use super::*;
1793    use crate::api::RedDBOptions;
1794    use crate::auth::{AuthConfig, AuthStore};
1795    use crate::runtime::RuntimeQueryResult;
1796    use crate::storage::query::unified::UnifiedResult;
1797    use reddb_rql::modes::QueryMode;
1798    use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
1799
1800    #[tokio::test]
1801    async fn response_writer_releases_execution_slot_before_socket_write() {
1802        let runtime = RedDBRuntime::with_options(RedDBOptions::in_memory()).unwrap();
1803        runtime.configure_execution_admission(1, 0, 100);
1804        let admission = runtime.execution_admission();
1805        let permit = admission.acquire().await.expect("execution permit");
1806        assert_eq!(admission.active(), 1);
1807
1808        let (mut server_io, mut client_io) = tokio::io::duplex(64);
1809        let mut writer = ExecutionResponseWriter::new(&mut server_io, Some(permit));
1810        writer.write_all(b"ok").await.unwrap();
1811        assert_eq!(admission.active(), 0);
1812
1813        let mut received = [0_u8; 2];
1814        client_io.read_exact(&mut received).await.unwrap();
1815        assert_eq!(&received, b"ok");
1816    }
1817
1818    #[tokio::test]
1819    async fn enabled_auth_store_rejects_bad_pgwire_password_with_28p01() {
1820        let runtime = Arc::new(RedDBRuntime::with_options(RedDBOptions::in_memory()).unwrap());
1821        let auth = Arc::new(AuthStore::new(AuthConfig {
1822            enabled: true,
1823            require_auth: true,
1824            ..AuthConfig::default()
1825        }));
1826        auth.create_user("alice", "secret", Role::Write).unwrap();
1827        runtime.set_auth_store(auth);
1828
1829        let config = Arc::new(PgWireConfig::default());
1830        let (server_io, mut client_io) = tokio::io::duplex(64 * 1024);
1831        let server = tokio::spawn(async move {
1832            handle_connection(server_io, runtime, config, None)
1833                .await
1834                .unwrap();
1835        });
1836
1837        write_startup_params(&mut client_io, &[("user", "alice")]).await;
1838        // Server must challenge for a cleartext password first.
1839        let (tag, body) = read_backend_frame(&mut client_io).await;
1840        assert_eq!(tag, b'R');
1841        assert_eq!(i32::from_be_bytes(body[..4].try_into().unwrap()), 3);
1842
1843        write_frontend_frame(&mut client_io, b'p', cstring_body("wrong")).await;
1844        let (tag, body) = read_backend_frame(&mut client_io).await;
1845        assert_eq!(tag, b'E');
1846        assert_eq!(error_response_field(&body, b'C'), Some("28P01"));
1847
1848        // The connection is closed after an auth failure.
1849        let mut eof = [0u8; 1];
1850        assert_eq!(client_io.read(&mut eof).await.unwrap(), 0);
1851        server.await.unwrap();
1852    }
1853
1854    #[tokio::test]
1855    async fn pgwire_password_auth_installs_statement_identity() {
1856        let runtime = Arc::new(RedDBRuntime::with_options(RedDBOptions::in_memory()).unwrap());
1857        let auth = Arc::new(AuthStore::new(AuthConfig {
1858            enabled: true,
1859            require_auth: true,
1860            ..AuthConfig::default()
1861        }));
1862        auth.create_user_in_tenant(Some("acme"), "alice", "secret", Role::Write)
1863            .unwrap();
1864        runtime.set_auth_store(auth);
1865
1866        let config = Arc::new(PgWireConfig::default());
1867        let (server_io, mut client_io) = tokio::io::duplex(64 * 1024);
1868        let server = tokio::spawn(async move {
1869            handle_connection(server_io, runtime, config, None)
1870                .await
1871                .unwrap();
1872        });
1873
1874        write_startup_params(&mut client_io, &[("user", "alice"), ("tenant", "acme")]).await;
1875        let (tag, body) = read_backend_frame(&mut client_io).await;
1876        assert_eq!(tag, b'R');
1877        assert_eq!(i32::from_be_bytes(body[..4].try_into().unwrap()), 3);
1878        write_frontend_frame(&mut client_io, b'p', cstring_body("secret")).await;
1879        read_until_ready(&mut client_io).await;
1880
1881        // The authenticated identity must drive CURRENT_USER/ROLE/TENANT so
1882        // the simple-query path yields real rows (T,D,C,Z), transport-
1883        // invariant with the native wire.
1884        write_frontend_frame(
1885            &mut client_io,
1886            b'Q',
1887            cstring_body("SELECT CURRENT_USER(), CURRENT_ROLE(), CURRENT_TENANT()"),
1888        )
1889        .await;
1890        let frames = read_until_ready(&mut client_io).await;
1891        assert_eq!(
1892            frames.iter().map(|(tag, _)| *tag).collect::<Vec<_>>(),
1893            b"TDCZ"
1894        );
1895        let cells = decode_data_row(&frames[1].1);
1896        assert_eq!(cells[0].as_deref(), Some(b"alice".as_slice()));
1897        assert_eq!(cells[1].as_deref(), Some(b"write".as_slice()));
1898        assert_eq!(cells[2].as_deref(), Some(b"acme".as_slice()));
1899
1900        write_frontend_frame(&mut client_io, b'X', Vec::new()).await;
1901        server.await.unwrap();
1902    }
1903
1904    #[tokio::test]
1905    async fn extended_parse_bind_execute_returns_rows() {
1906        let runtime = Arc::new(RedDBRuntime::with_options(RedDBOptions::in_memory()).unwrap());
1907        let config = Arc::new(PgWireConfig::default());
1908        let (server_io, mut client_io) = tokio::io::duplex(64 * 1024);
1909        let server = tokio::spawn(async move {
1910            handle_connection(server_io, runtime, config, None)
1911                .await
1912                .unwrap();
1913        });
1914
1915        write_startup(&mut client_io).await;
1916        read_until_ready(&mut client_io).await;
1917
1918        write_frontend_frame(
1919            &mut client_io,
1920            b'P',
1921            parse_body("", "SELECT $1::int", &[PgOid::Int4.as_u32()]),
1922        )
1923        .await;
1924        write_frontend_frame(
1925            &mut client_io,
1926            b'B',
1927            bind_body("", "", &[0], &[Some(b"42".as_slice())], &[]),
1928        )
1929        .await;
1930        write_frontend_frame(&mut client_io, b'D', describe_body(b'P', "")).await;
1931        write_frontend_frame(&mut client_io, b'E', execute_body("", 0)).await;
1932        write_frontend_frame(&mut client_io, b'S', Vec::new()).await;
1933
1934        let frames = read_until_ready(&mut client_io).await;
1935        assert_eq!(
1936            frames.iter().map(|(tag, _)| *tag).collect::<Vec<_>>(),
1937            b"12TDCZ"
1938        );
1939        let columns = decode_row_description(&frames[2].1);
1940        assert_eq!(columns.len(), 1);
1941        let cells = decode_data_row(&frames[3].1);
1942        assert_eq!(cells.len(), 1);
1943        assert_eq!(cells[0].as_deref(), Some(b"42".as_slice()));
1944        assert_eq!(decode_command_complete(&frames[4].1), "SELECT 1");
1945
1946        write_frontend_frame(&mut client_io, b'X', Vec::new()).await;
1947        server.await.unwrap();
1948    }
1949
1950    #[test]
1951    fn prepared_pg_utility_statements_preserve_ddl_and_transactions() {
1952        let runtime = RedDBRuntime::in_memory().expect("runtime");
1953        let registry = PreparedRegistry::new();
1954        let execute = |sql: &str, params: Vec<ParamValue>| {
1955            let prepared = registry
1956                .prepare(&runtime, sql)
1957                .expect("prepare utility or DML");
1958            execute_pg_query_result(&runtime, None, &registry, prepared.id, sql, &params)
1959                .unwrap_or_else(|error| panic!("{sql}: {error}"))
1960        };
1961        execute("CREATE TABLE pg_utilities (id INT)", Vec::new());
1962        execute("BEGIN", Vec::new());
1963        execute(
1964            "INSERT INTO pg_utilities (id) VALUES ($1)",
1965            vec![ParamValue::Int64(1)],
1966        );
1967        execute("SAVEPOINT before_second", Vec::new());
1968        execute(
1969            "INSERT INTO pg_utilities (id) VALUES ($1)",
1970            vec![ParamValue::Int64(2)],
1971        );
1972        execute("ROLLBACK TO SAVEPOINT before_second", Vec::new());
1973        execute("COMMIT", Vec::new());
1974        let result = execute("SELECT id FROM pg_utilities", Vec::new());
1975        assert_eq!(result.result.records.len(), 1);
1976        assert_eq!(result.result.records[0].get("id"), Some(&Value::Integer(1)));
1977        execute("DROP TABLE pg_utilities", Vec::new());
1978        assert!(runtime
1979            .db()
1980            .store()
1981            .get_collection("pg_utilities")
1982            .is_none());
1983    }
1984
1985    #[tokio::test]
1986    async fn extended_repeated_execute_parses_prepared_shape_once() {
1987        let runtime = RedDBRuntime::with_options(RedDBOptions::in_memory()).unwrap();
1988        runtime
1989            .execute_query("CREATE TABLE pg_parse_once (id INT)")
1990            .unwrap();
1991        runtime
1992            .execute_query("INSERT INTO pg_parse_once (id) VALUES (42)")
1993            .unwrap();
1994        let prepared_registry = PreparedRegistry::new();
1995        let mut prepared = HashMap::new();
1996        let mut portals = HashMap::new();
1997        let (mut server_io, _client_io) = tokio::io::duplex(64 * 1024);
1998
1999        handle_parse(
2000            &mut server_io,
2001            &runtime,
2002            &prepared_registry,
2003            &mut prepared,
2004            &portals,
2005            super::super::protocol::ParseMessage {
2006                statement: "parse_once_stmt".to_string(),
2007                query: "SELECT id FROM pg_parse_once WHERE id = $1".to_string(),
2008                param_type_oids: vec![PgOid::Int4.as_u32()],
2009            },
2010        )
2011        .await
2012        .unwrap();
2013        assert_eq!(prepared_registry.parse_count(), 1);
2014
2015        handle_bind(
2016            &mut server_io,
2017            &prepared_registry,
2018            &prepared,
2019            &mut portals,
2020            super::super::protocol::BindMessage {
2021                portal: "parse_once_portal".to_string(),
2022                statement: "parse_once_stmt".to_string(),
2023                param_format_codes: vec![0],
2024                params: vec![Some(b"42".to_vec())],
2025                result_format_codes: Vec::new(),
2026            },
2027        )
2028        .await
2029        .unwrap();
2030
2031        for _ in 0..2 {
2032            handle_execute(
2033                &mut server_io,
2034                &runtime,
2035                None,
2036                &prepared_registry,
2037                &mut portals,
2038                super::super::protocol::ExecuteMessage {
2039                    portal: "parse_once_portal".to_string(),
2040                    max_rows: 0,
2041                },
2042            )
2043            .await
2044            .unwrap();
2045        }
2046        assert_eq!(prepared_registry.parse_count(), 1);
2047    }
2048
2049    #[tokio::test]
2050    async fn non_loopback_plaintext_refuses_cleartext_password_without_tls() {
2051        // Posture (issue #1653): a cleartext password may not be challenged
2052        // over an unencrypted non-loopback link. The server must refuse
2053        // *before* sending AuthenticationCleartextPassword, so a driver
2054        // using sslmode=disable against a remote bind never leaks its
2055        // password in the clear. `tls_acceptor = None` + no SSLRequest is
2056        // exactly that plaintext, no-TLS path.
2057        let runtime = Arc::new(RedDBRuntime::with_options(RedDBOptions::in_memory()).unwrap());
2058        let auth = Arc::new(AuthStore::new(AuthConfig {
2059            enabled: true,
2060            require_auth: true,
2061            ..AuthConfig::default()
2062        }));
2063        auth.create_user("alice", "secret", Role::Write).unwrap();
2064        runtime.set_auth_store(auth);
2065
2066        // Non-loopback bind: the wildcard address is not a loopback host.
2067        let config = Arc::new(PgWireConfig {
2068            bind_addr: "0.0.0.0:5432".to_string(),
2069            ..PgWireConfig::default()
2070        });
2071        let (server_io, mut client_io) = tokio::io::duplex(64 * 1024);
2072        let server = tokio::spawn(async move {
2073            handle_connection(server_io, runtime, config, None)
2074                .await
2075                .unwrap();
2076        });
2077
2078        write_startup_params(&mut client_io, &[("user", "alice")]).await;
2079        // The refusal arrives directly — never an 'R' cleartext challenge.
2080        let (tag, body) = read_backend_frame(&mut client_io).await;
2081        assert_eq!(tag, b'E', "expected an ErrorResponse, not a challenge");
2082        assert_eq!(error_response_field(&body, b'C'), Some("28000"));
2083
2084        // The connection is closed after the refusal.
2085        let mut eof = [0u8; 1];
2086        assert_eq!(client_io.read(&mut eof).await.unwrap(), 0);
2087        server.await.unwrap();
2088    }
2089
2090    #[test]
2091    fn infer_pg_cast_param_type_oids_from_parameter_casts() {
2092        assert_eq!(
2093            infer_pg_cast_param_type_oids("INSERT INTO t (id, name) VALUES ($1::int, $2::text)"),
2094            vec![(0, PgOid::Int4.as_u32()), (1, PgOid::Text.as_u32())]
2095        );
2096        assert_eq!(
2097            infer_pg_cast_param_type_oids("SEARCH SIMILAR [1.0] COLLECTION v LIMIT $1::int8"),
2098            vec![(0, PgOid::Int8.as_u32())]
2099        );
2100    }
2101
2102    #[test]
2103    fn pg_session_compat_accepts_driver_setup_set_commands() {
2104        assert_eq!(
2105            pg_session_compat_command_tag("SET extra_float_digits = 3"),
2106            Some("SET")
2107        );
2108        assert_eq!(
2109            pg_session_compat_command_tag("SET application_name = 'pgjdbc'"),
2110            Some("SET")
2111        );
2112        assert_eq!(pg_session_compat_command_tag("SELECT 1"), None);
2113    }
2114
2115    #[tokio::test]
2116    async fn ask_success_result_uses_canonical_pg_wire_row_shape() {
2117        let mut result = UnifiedResult::with_columns(vec![
2118            "answer".into(),
2119            "provider".into(),
2120            "model".into(),
2121            "prompt_tokens".into(),
2122            "completion_tokens".into(),
2123            "sources_count".into(),
2124            "sources_flat".into(),
2125            "citations".into(),
2126            "validation".into(),
2127        ]);
2128        let mut record = UnifiedRecord::new();
2129        record.set("answer", Value::text("Deploy failed [^1]."));
2130        record.set("provider", Value::text("openai"));
2131        record.set("model", Value::text("gpt-4o-mini"));
2132        record.set("prompt_tokens", Value::Integer(11));
2133        record.set("completion_tokens", Value::Integer(7));
2134        record.set(
2135            "sources_flat",
2136            Value::Json(
2137                br#"[{"urn":"urn:reddb:row:deployments:1","kind":"row","collection":"deployments","id":"1"}]"#
2138                    .to_vec(),
2139            ),
2140        );
2141        record.set(
2142            "citations",
2143            Value::Json(br#"[{"marker":1,"urn":"urn:reddb:row:deployments:1"}]"#.to_vec()),
2144        );
2145        record.set(
2146            "validation",
2147            Value::Json(br#"{"ok":true,"warnings":[],"errors":[]}"#.to_vec()),
2148        );
2149        result.push(record);
2150
2151        let qr = RuntimeQueryResult {
2152            query: "ASK 'why did deploy fail?'".to_string(),
2153            mode: QueryMode::Sql,
2154            statement: "ask",
2155            engine: "runtime-ai",
2156            result,
2157            affected_rows: 0,
2158            statement_type: "select",
2159            bookmark: None,
2160            notice: None,
2161        };
2162
2163        let mut out = Vec::new();
2164        emit_success_result(&mut out, &qr).await.unwrap();
2165        let frames = decode_frames(&out);
2166
2167        assert_eq!(
2168            frames.iter().map(|(tag, _)| *tag).collect::<Vec<_>>(),
2169            b"TDC"
2170        );
2171
2172        let columns = decode_row_description(frames[0].1);
2173        assert_eq!(
2174            columns,
2175            vec![
2176                ("answer".to_string(), PgOid::Text.as_u32()),
2177                ("cache_hit".to_string(), PgOid::Bool.as_u32()),
2178                ("citations".to_string(), PgOid::Jsonb.as_u32()),
2179                ("completion_tokens".to_string(), PgOid::Int8.as_u32()),
2180                ("cost_usd".to_string(), PgOid::Numeric.as_u32()),
2181                ("mode".to_string(), PgOid::Text.as_u32()),
2182                ("model".to_string(), PgOid::Text.as_u32()),
2183                ("prompt_tokens".to_string(), PgOid::Int8.as_u32()),
2184                ("provider".to_string(), PgOid::Text.as_u32()),
2185                ("retry_count".to_string(), PgOid::Int8.as_u32()),
2186                ("sources_flat".to_string(), PgOid::Jsonb.as_u32()),
2187                ("validation".to_string(), PgOid::Jsonb.as_u32()),
2188            ]
2189        );
2190
2191        let cells = decode_data_row(frames[1].1);
2192        assert_eq!(cells.len(), 12);
2193        assert_eq!(cells[0].as_deref(), Some(b"Deploy failed [^1].".as_slice()));
2194        assert_eq!(cells[1].as_deref(), Some(b"f".as_slice()));
2195        assert_eq!(cells[4].as_deref(), Some(b"0".as_slice()));
2196        assert_eq!(cells[5].as_deref(), Some(b"strict".as_slice()));
2197        assert_eq!(cells[9].as_deref(), Some(b"0".as_slice()));
2198        assert!(std::str::from_utf8(cells[10].as_deref().unwrap())
2199            .unwrap()
2200            .contains(r#""payload""#));
2201        assert_eq!(decode_command_complete(frames[2].1), "SELECT 1");
2202    }
2203
2204    fn decode_frames(bytes: &[u8]) -> Vec<(u8, &[u8])> {
2205        let mut pos = 0;
2206        let mut frames = Vec::new();
2207        while pos < bytes.len() {
2208            let tag = bytes[pos];
2209            let len = u32::from_be_bytes([
2210                bytes[pos + 1],
2211                bytes[pos + 2],
2212                bytes[pos + 3],
2213                bytes[pos + 4],
2214            ]) as usize;
2215            let body_start = pos + 5;
2216            let body_end = pos + 1 + len;
2217            frames.push((tag, &bytes[body_start..body_end]));
2218            pos = body_end;
2219        }
2220        frames
2221    }
2222
2223    fn decode_row_description(body: &[u8]) -> Vec<(String, u32)> {
2224        let count = i16::from_be_bytes([body[0], body[1]]) as usize;
2225        let mut pos = 2;
2226        let mut columns = Vec::with_capacity(count);
2227        for _ in 0..count {
2228            let end = body[pos..].iter().position(|&b| b == 0).unwrap() + pos;
2229            let name = std::str::from_utf8(&body[pos..end]).unwrap().to_string();
2230            pos = end + 1;
2231            pos += 4; // table oid
2232            pos += 2; // column attr
2233            let oid = u32::from_be_bytes([body[pos], body[pos + 1], body[pos + 2], body[pos + 3]]);
2234            pos += 4;
2235            pos += 2; // type size
2236            pos += 4; // type mod
2237            pos += 2; // format
2238            columns.push((name, oid));
2239        }
2240        columns
2241    }
2242
2243    fn decode_data_row(body: &[u8]) -> Vec<Option<Vec<u8>>> {
2244        let count = i16::from_be_bytes([body[0], body[1]]) as usize;
2245        let mut pos = 2;
2246        let mut cells = Vec::with_capacity(count);
2247        for _ in 0..count {
2248            let len = i32::from_be_bytes([body[pos], body[pos + 1], body[pos + 2], body[pos + 3]]);
2249            pos += 4;
2250            if len < 0 {
2251                cells.push(None);
2252            } else {
2253                let len = len as usize;
2254                cells.push(Some(body[pos..pos + len].to_vec()));
2255                pos += len;
2256            }
2257        }
2258        cells
2259    }
2260
2261    fn decode_command_complete(body: &[u8]) -> &str {
2262        let nul = body.iter().position(|&b| b == 0).unwrap_or(body.len());
2263        std::str::from_utf8(&body[..nul]).unwrap()
2264    }
2265
2266    async fn write_startup<W: AsyncWrite + Unpin>(stream: &mut W) {
2267        write_startup_params(stream, &[("user", "reddb")]).await;
2268    }
2269
2270    async fn write_startup_params<W: AsyncWrite + Unpin>(stream: &mut W, params: &[(&str, &str)]) {
2271        let mut payload = Vec::new();
2272        payload.extend_from_slice(&crate::wire::postgres::protocol::PG_PROTOCOL_V3.to_be_bytes());
2273        for (name, value) in params {
2274            push_pg_cstring(&mut payload, name);
2275            push_pg_cstring(&mut payload, value);
2276        }
2277        payload.push(0);
2278        let len = (payload.len() + 4) as u32;
2279        stream.write_all(&len.to_be_bytes()).await.unwrap();
2280        stream.write_all(&payload).await.unwrap();
2281    }
2282
2283    fn cstring_body(value: &str) -> Vec<u8> {
2284        let mut out = Vec::new();
2285        push_pg_cstring(&mut out, value);
2286        out
2287    }
2288
2289    fn error_response_field(body: &[u8], field_tag: u8) -> Option<&str> {
2290        let mut pos = 0;
2291        while pos < body.len() && body[pos] != 0 {
2292            let tag = body[pos];
2293            pos += 1;
2294            let end = body[pos..].iter().position(|&b| b == 0)? + pos;
2295            if tag == field_tag {
2296                return std::str::from_utf8(&body[pos..end]).ok();
2297            }
2298            pos = end + 1;
2299        }
2300        None
2301    }
2302
2303    async fn write_frontend_frame<W: AsyncWrite + Unpin>(
2304        stream: &mut W,
2305        tag: u8,
2306        payload: Vec<u8>,
2307    ) {
2308        stream.write_all(&[tag]).await.unwrap();
2309        stream
2310            .write_all(&((payload.len() + 4) as u32).to_be_bytes())
2311            .await
2312            .unwrap();
2313        stream.write_all(&payload).await.unwrap();
2314    }
2315
2316    async fn read_backend_frame<R: AsyncRead + Unpin>(stream: &mut R) -> (u8, Vec<u8>) {
2317        let mut tag = [0u8; 1];
2318        stream.read_exact(&mut tag).await.unwrap();
2319        let mut len = [0u8; 4];
2320        stream.read_exact(&mut len).await.unwrap();
2321        let len = u32::from_be_bytes(len) as usize;
2322        let mut body = vec![0u8; len - 4];
2323        stream.read_exact(&mut body).await.unwrap();
2324        (tag[0], body)
2325    }
2326
2327    async fn read_until_ready<R: AsyncRead + Unpin>(stream: &mut R) -> Vec<(u8, Vec<u8>)> {
2328        let mut frames = Vec::new();
2329        loop {
2330            let frame = read_backend_frame(stream).await;
2331            let done = frame.0 == b'Z';
2332            frames.push(frame);
2333            if done {
2334                return frames;
2335            }
2336        }
2337    }
2338
2339    fn parse_body(statement: &str, query: &str, oids: &[u32]) -> Vec<u8> {
2340        let mut out = Vec::new();
2341        push_pg_cstring(&mut out, statement);
2342        push_pg_cstring(&mut out, query);
2343        out.extend_from_slice(&(oids.len() as i16).to_be_bytes());
2344        for oid in oids {
2345            out.extend_from_slice(&oid.to_be_bytes());
2346        }
2347        out
2348    }
2349
2350    fn bind_body(
2351        portal: &str,
2352        statement: &str,
2353        formats: &[i16],
2354        params: &[Option<&[u8]>],
2355        result_formats: &[i16],
2356    ) -> Vec<u8> {
2357        let mut out = Vec::new();
2358        push_pg_cstring(&mut out, portal);
2359        push_pg_cstring(&mut out, statement);
2360        out.extend_from_slice(&(formats.len() as i16).to_be_bytes());
2361        for format in formats {
2362            out.extend_from_slice(&format.to_be_bytes());
2363        }
2364        out.extend_from_slice(&(params.len() as i16).to_be_bytes());
2365        for param in params {
2366            match param {
2367                Some(bytes) => {
2368                    out.extend_from_slice(&(bytes.len() as i32).to_be_bytes());
2369                    out.extend_from_slice(bytes);
2370                }
2371                None => out.extend_from_slice(&(-1i32).to_be_bytes()),
2372            }
2373        }
2374        out.extend_from_slice(&(result_formats.len() as i16).to_be_bytes());
2375        for format in result_formats {
2376            out.extend_from_slice(&format.to_be_bytes());
2377        }
2378        out
2379    }
2380
2381    fn describe_body(target: u8, name: &str) -> Vec<u8> {
2382        let mut out = vec![target];
2383        push_pg_cstring(&mut out, name);
2384        out
2385    }
2386
2387    fn execute_body(portal: &str, max_rows: u32) -> Vec<u8> {
2388        let mut out = Vec::new();
2389        push_pg_cstring(&mut out, portal);
2390        out.extend_from_slice(&max_rows.to_be_bytes());
2391        out
2392    }
2393
2394    fn push_pg_cstring(out: &mut Vec<u8>, value: &str) {
2395        out.extend_from_slice(value.as_bytes());
2396        out.push(0);
2397    }
2398}
2399
2400/// A fresh 32-bit cancel key from the OS CSPRNG. Falls back to a
2401/// process/time-derived value only if the CSPRNG is unavailable — the
2402/// key guards nothing yet, so refusing the connection would be the
2403/// wrong trade.
2404fn random_cancel_key() -> u32 {
2405    let mut bytes = [0u8; 4];
2406    if crate::crypto::os_random::fill_bytes(&mut bytes).is_ok() {
2407        return u32::from_le_bytes(bytes);
2408    }
2409    let nanos = std::time::SystemTime::now()
2410        .duration_since(std::time::UNIX_EPOCH)
2411        .map(|d| d.subsec_nanos())
2412        .unwrap_or(0);
2413    std::process::id() ^ nanos
2414}