pgwire 0.39.0

//! APIs for building postgresql compatible servers.

use std::any::{Any, TypeId};
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::{Arc, RwLock};

use bytes::Bytes;
use futures::channel::oneshot;
use futures::lock::Mutex;
pub use postgres_types::Type;
#[cfg(any(feature = "_ring", feature = "_aws-lc-rs"))]
use rustls_pki_types::CertificateDer;

use crate::error::PgWireError;
use crate::messages::ProtocolVersion;
use crate::messages::response::TransactionStatus;
use crate::messages::startup::SecretKey;

/// Authentication handlers for the startup phase.
pub mod auth;
/// Cancel request handling.
pub mod cancel;
#[cfg(feature = "client-api")]
/// Client-side PostgreSQL protocol implementation.
pub mod client;
/// COPY protocol handlers.
pub mod copy;
/// Portal management for the extended query protocol.
pub mod portal;
/// Simple and extended query handlers.
pub mod query;
/// Query result encoding and formatting.
pub mod results;
/// Prepared statement management.
pub mod stmt;
/// Portal and prepared statement storage backends.
pub mod store;
/// Transaction status tracking.
pub mod transaction;

/// Default name used when no name is provided for prepared statements or portals.
pub const DEFAULT_NAME: &str = "POSTGRESQL_DEFAULT_NAME";

/// States of a PostgreSQL connection lifecycle.
#[derive(Debug, Clone, Copy, Default)]
pub enum PgWireConnectionState {
    /// Waiting for an SSL request from the client.
    #[default]
    AwaitingSslRequest,
    /// Waiting for a startup message from the client.
    AwaitingStartup,
    /// Authentication handshake in progress.
    AuthenticationInProgress,
    /// Connection is idle and ready for queries.
    ReadyForQuery,
    /// A query is currently being executed.
    QueryInProgress,
    /// A COPY operation is in progress.
    CopyInProgress(bool),
    /// Waiting for a Sync message from the client.
    AwaitingSync,
}

/// Per-connection typed extension store, keyed by `TypeId`.
///
/// Allows hooks and handlers to store arbitrary per-session state that is
/// automatically cleaned up when the connection closes. Uses interior
/// mutability so it can be accessed from `&dyn ClientInfo`.
#[derive(Default)]
pub struct SessionExtensions {
    map: RwLock<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
}

impl std::fmt::Debug for SessionExtensions {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let count = self.map.read().unwrap().len();
        f.debug_struct("SessionExtensions")
            .field("entries", &count)
            .finish()
    }
}

impl SessionExtensions {
    /// Create a new empty extension store.
    pub fn new() -> Self {
        Self {
            map: RwLock::new(HashMap::new()),
        }
    }

    /// Get an extension by type, or insert a default value if not present.
    pub fn get_or_insert_with<T: Send + Sync + 'static>(&self, f: impl FnOnce() -> T) -> Arc<T> {
        let type_id = TypeId::of::<T>();

        // Fast path: read lock
        {
            let map = self.map.read().unwrap();
            if let Some(val) = map.get(&type_id) {
                return val.clone().downcast::<T>().unwrap();
            }
        }

        // Slow path: write lock
        let mut map = self.map.write().unwrap();
        let val = map.entry(type_id).or_insert_with(|| Arc::new(f()));
        val.clone().downcast::<T>().unwrap()
    }

    /// Get an extension by type, returning `None` if not present.
    pub fn get<T: Send + Sync + 'static>(&self) -> Option<Arc<T>> {
        let type_id = TypeId::of::<T>();
        let map = self.map.read().unwrap();
        map.get(&type_id)
            .map(|val| val.clone().downcast::<T>().unwrap())
    }

    /// Insert an extension by type, replacing any existing value.
    pub fn insert<T: Send + Sync + 'static>(&self, val: T) -> Option<Arc<T>> {
        let type_id = TypeId::of::<T>();
        let mut map = self.map.write().unwrap();
        map.insert(type_id, Arc::new(val))
            .map(|old| old.downcast::<T>().unwrap())
    }
}

// TODO: add oauth scope and issuer
/// Describe a client information holder
pub trait ClientInfo {
    /// Returns the client's socket address.
    fn socket_addr(&self) -> SocketAddr;

    /// Returns whether the connection is secured (TLS).
    fn is_secure(&self) -> bool;

    /// Returns the protocol version negotiated with the client.
    fn protocol_version(&self) -> ProtocolVersion;

    /// Sets the protocol version for this connection.
    fn set_protocol_version(&mut self, version: ProtocolVersion);

    /// Returns the process ID and secret key used for cancel requests.
    fn pid_and_secret_key(&self) -> (i32, SecretKey);

    /// Sets the process ID and secret key for cancel requests.
    fn set_pid_and_secret_key(&mut self, pid: i32, secret_key: SecretKey);

    /// Returns the current connection state.
    fn state(&self) -> PgWireConnectionState;

    /// Transitions the connection to a new state.
    fn set_state(&mut self, new_state: PgWireConnectionState);

    /// Returns the current transaction status.
    fn transaction_status(&self) -> TransactionStatus;

    /// Updates the transaction status.
    fn set_transaction_status(&mut self, new_status: TransactionStatus);

    /// Returns the connection metadata key-value map.
    fn metadata(&self) -> &HashMap<String, String>;

    /// Returns a mutable reference to the connection metadata.
    fn metadata_mut(&mut self) -> &mut HashMap<String, String>;

    /// Returns the per-session extension store.
    fn session_extensions(&self) -> &SessionExtensions;

    /// Returns the TLS SNI server name, if available.
    #[cfg(any(feature = "_ring", feature = "_aws-lc-rs"))]
    fn sni_server_name(&self) -> Option<&str>;

    /// Returns the client TLS certificates, if available.
    #[cfg(any(feature = "_ring", feature = "_aws-lc-rs"))]
    fn client_certificates<'a>(&self) -> Option<&[CertificateDer<'a>]>;
}

/// Client Portal Store
pub trait ClientPortalStore {
    /// The underlying portal store type.
    type PortalStore;

    /// Returns a reference to the portal store.
    fn portal_store(&self) -> &Self::PortalStore;
}

/// Metadata key for the connected user name.
pub const METADATA_USER: &str = "user";
/// Metadata key for the target database name.
pub const METADATA_DATABASE: &str = "database";
/// Metadata key for the client encoding setting.
pub const METADATA_CLIENT_ENCODING: &str = "client_encoding";
/// Metadata key for the application name.
pub const METADATA_APPLICATION_NAME: &str = "application_name";

/// Default client implementation holding all per-connection state.
#[non_exhaustive]
#[derive(Debug)]
pub struct DefaultClient<S> {
    /// The client's socket address.
    pub socket_addr: SocketAddr,
    /// Whether the connection is secured with TLS.
    pub is_secure: bool,
    /// The negotiated protocol version.
    pub protocol_version: ProtocolVersion,
    /// Process ID and secret key for cancel requests.
    pub pid_secret_key: (i32, SecretKey),
    /// Current connection state.
    pub state: PgWireConnectionState,
    /// Current transaction status.
    pub transaction_status: TransactionStatus,
    /// Connection metadata key-value pairs.
    pub metadata: HashMap<String, String>,
    /// The TLS SNI server name, if using TLS.
    #[cfg(any(feature = "_ring", feature = "_aws-lc-rs"))]
    pub sni_server_name: Option<String>,
    /// In-memory portal and prepared statement store.
    pub portal_store: store::MemPortalStore<S>,
    /// Per-session typed extension store.
    pub session_extensions: SessionExtensions,
}

impl<S> ClientInfo for DefaultClient<S> {
    fn socket_addr(&self) -> SocketAddr {
        self.socket_addr
    }

    fn is_secure(&self) -> bool {
        self.is_secure
    }

    fn pid_and_secret_key(&self) -> (i32, SecretKey) {
        self.pid_secret_key.clone()
    }

    fn set_pid_and_secret_key(&mut self, pid: i32, secret_key: SecretKey) {
        self.pid_secret_key = (pid, secret_key);
    }

    fn protocol_version(&self) -> ProtocolVersion {
        self.protocol_version
    }

    fn set_protocol_version(&mut self, version: ProtocolVersion) {
        self.protocol_version = version;
    }

    fn state(&self) -> PgWireConnectionState {
        self.state
    }

    fn set_state(&mut self, new_state: PgWireConnectionState) {
        self.state = new_state;
    }

    fn metadata(&self) -> &HashMap<String, String> {
        &self.metadata
    }

    fn metadata_mut(&mut self) -> &mut HashMap<String, String> {
        &mut self.metadata
    }

    fn session_extensions(&self) -> &SessionExtensions {
        &self.session_extensions
    }

    fn transaction_status(&self) -> TransactionStatus {
        self.transaction_status
    }

    fn set_transaction_status(&mut self, new_status: TransactionStatus) {
        self.transaction_status = new_status
    }

    #[cfg(any(feature = "_ring", feature = "_aws-lc-rs"))]
    fn sni_server_name(&self) -> Option<&str> {
        self.sni_server_name.as_deref()
    }

    #[cfg(any(feature = "_ring", feature = "_aws-lc-rs"))]
    fn client_certificates<'a>(&self) -> Option<&[CertificateDer<'a>]> {
        None
    }
}

impl<S> DefaultClient<S> {
    /// Create a new `DefaultClient` with the given socket address and TLS status.
    pub fn new(socket_addr: SocketAddr, is_secure: bool) -> DefaultClient<S> {
        DefaultClient {
            socket_addr,
            is_secure,
            protocol_version: ProtocolVersion::default(),
            pid_secret_key: (0, SecretKey::default()),
            state: PgWireConnectionState::default(),
            transaction_status: TransactionStatus::Idle,
            metadata: HashMap::new(),
            #[cfg(any(feature = "_ring", feature = "_aws-lc-rs"))]
            sni_server_name: None,
            portal_store: store::MemPortalStore::new(),
            session_extensions: SessionExtensions::new(),
        }
    }
}

impl<S> ClientPortalStore for DefaultClient<S> {
    type PortalStore = store::MemPortalStore<S>;

    fn portal_store(&self) -> &Self::PortalStore {
        &self.portal_store
    }
}

/// Generator for process ID and secret key pairs used to identify connections
/// for cancel requests.
///
/// Implement this trait to customize how PID and secret key values are
/// produced. The default implementation [`RandomPidSecretKeyGenerator`]
/// generates random values.
pub trait PidSecretKeyGenerator: Send + Sync {
    fn generate(&self, client: &dyn ClientInfo) -> (i32, SecretKey);
}

/// Default implementation of [`PidSecretKeyGenerator`] that produces random
/// values using the `rand` crate.
///
/// For protocol 3.0, generates a 4-byte `SecretKey::I32`. For protocol 3.2,
/// generates a 32-byte `SecretKey::Bytes`.
#[derive(Debug, Default)]
pub struct RandomPidSecretKeyGenerator;

impl PidSecretKeyGenerator for RandomPidSecretKeyGenerator {
    fn generate(&self, client: &dyn ClientInfo) -> (i32, SecretKey) {
        let pid = (rand::random::<u32>() >> 1) as i32;
        let secret_key = match client.protocol_version() {
            ProtocolVersion::PROTOCOL3_0 => SecretKey::I32(rand::random::<i32>()),
            ProtocolVersion::PROTOCOL3_2 => {
                let mut bytes = vec![0u8; 32];
                rand::fill(&mut bytes);
                SecretKey::Bytes(bytes.into())
            }
        };
        (pid, secret_key)
    }
}

/// Per-connection cancel handle.
///
/// Stores a replaceable `oneshot::Sender` that is swapped out before each
/// query via [`ConnectionHandle::start_query`]. When a cancel request arrives,
/// [`ConnectionHandle::cancel`] fires the current sender.
pub struct ConnectionHandle {
    cancel_tx: Mutex<Option<oneshot::Sender<()>>>,
}

impl std::fmt::Debug for ConnectionHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ConnectionHandle").finish()
    }
}

impl ConnectionHandle {
    pub fn new() -> Self {
        Self {
            cancel_tx: Mutex::new(None),
        }
    }

    /// Install a fresh cancel sender and return the receiver.
    ///
    /// Call this before each query. The previous sender (if any) is dropped,
    /// causing its paired receiver to resolve with `Err(Canceled)`.
    pub async fn start_query(&self) -> oneshot::Receiver<()> {
        let (tx, rx) = oneshot::channel();
        *self.cancel_tx.lock().await = Some(tx);
        rx
    }

    /// Fire the current cancel sender.
    ///
    /// Returns `true` if there was an active query to cancel, `false` if no
    /// query was in progress.
    pub async fn cancel(&self) -> bool {
        if let Some(tx) = self.cancel_tx.lock().await.take() {
            let _ = tx.send(());
            true
        } else {
            false
        }
    }
}

impl Default for ConnectionHandle {
    fn default() -> Self {
        Self::new()
    }
}

/// RAII guard that unregisters a connection from the [`ConnectionManager`] on
/// drop.
pub struct ConnectionGuard {
    pid: i32,
    secret_key: Bytes,
    manager: Arc<ConnectionManager>,
}

impl std::fmt::Debug for ConnectionGuard {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ConnectionGuard")
            .field("pid", &self.pid)
            .finish()
    }
}

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        self.manager
            .inner
            .write()
            .unwrap()
            .remove(&(self.pid, self.secret_key.clone()));
    }
}

/// A registry mapping `(pid, secret_key)` to per-connection cancel handles.
///
/// Create one `ConnectionManager` at server startup, share it as
/// `Arc<ConnectionManager>` across all handler implementations.
///
/// ## Lifecycle
///
/// 1. In your `StartupHandler`, call [`register`](ConnectionManager::register)
///    to obtain an [`Arc<ConnectionHandle>`] and a [`ConnectionGuard`].
/// 2. Store both in [`SessionExtensions`] so they live as long as the
///    connection.
/// 3. Before each query, call
///    [`handle.start_query()`](ConnectionHandle::start_query) to get a
///    `oneshot::Receiver`. `select!` on it alongside your query.
/// 4. A cancel request fires the receiver, interrupting the current query.
/// 5. When the connection ends (task dropped), the `ConnectionGuard` drops and
///    unregisters from this manager.
#[derive(Debug)]
pub struct ConnectionManager {
    inner: RwLock<HashMap<(i32, Bytes), Arc<ConnectionHandle>>>,
}

impl ConnectionManager {
    pub fn new() -> Self {
        Self {
            inner: RwLock::new(HashMap::new()),
        }
    }

    /// Register a new connection.
    ///
    /// Returns a shared [`ConnectionHandle`] for query-time cancel signaling
    /// and a [`ConnectionGuard`] that unregisters on drop.
    ///
    /// If the `(pid, secret_key)` pair is already registered, returns the
    /// existing handle and a new guard.
    pub fn register(
        self: &Arc<Self>,
        pid: i32,
        secret_key: SecretKey,
    ) -> (Arc<ConnectionHandle>, ConnectionGuard) {
        let key = (pid, secret_key.to_bytes());
        let mut map = self.inner.write().unwrap();
        let handle = if let Some(existing) = map.get(&key) {
            existing.clone()
        } else {
            let handle = Arc::new(ConnectionHandle::new());
            map.insert(key.clone(), handle.clone());
            handle
        };
        (
            handle,
            ConnectionGuard {
                pid,
                secret_key: key.1,
                manager: Arc::clone(self),
            },
        )
    }

    /// Cancel the current query on a connection identified by `(pid,
    /// secret_key)`.
    ///
    /// Returns `true` if the connection was found (regardless of whether a
    /// query was active), `false` if the connection was not registered.
    pub async fn cancel(&self, pid: i32, secret_key: &SecretKey) -> bool {
        let key = (pid, secret_key.to_bytes());
        let handle = self.inner.read().unwrap().get(&key).cloned();
        if let Some(handle) = handle {
            handle.cancel().await;
            true
        } else {
            false
        }
    }
}

impl Default for ConnectionManager {
    fn default() -> Self {
        Self::new()
    }
}

/// A centralized handler for all errors
///
/// This handler captures all errors produces by authentication, query and
/// copy. You can do logging, filtering or masking the error before it sent to
/// client.
pub trait ErrorHandler: Send + Sync {
    fn on_error<C>(&self, _client: &C, _error: &mut PgWireError)
    where
        C: ClientInfo,
    {
    }
}

/// A noop implementation for `ErrorHandler`.
#[derive(Debug)]
pub struct NoopHandler;

impl ErrorHandler for NoopHandler {}

/// A collection of all handler traits required to serve a PostgreSQL connection.
///
/// Implement this trait to provide handlers for each phase of the wire protocol.
pub trait PgWireServerHandlers: 'static {
    /// Returns the handler for simple (text) queries.
    fn simple_query_handler(&self) -> Arc<impl query::SimpleQueryHandler> {
        Arc::new(NoopHandler)
    }

    /// Returns the handler for extended query protocol messages.
    fn extended_query_handler(&self) -> Arc<impl query::ExtendedQueryHandler> {
        Arc::new(NoopHandler)
    }

    /// Returns the handler for the startup/authentication phase.
    fn startup_handler(&self) -> Arc<impl auth::StartupHandler> {
        Arc::new(NoopHandler)
    }

    /// Returns the handler for COPY sub-protocol messages.
    fn copy_handler(&self) -> Arc<impl copy::CopyHandler> {
        Arc::new(NoopHandler)
    }

    /// Returns the error handler for processing errors before sending to the client.
    fn error_handler(&self) -> Arc<impl ErrorHandler> {
        Arc::new(NoopHandler)
    }

    /// Returns the handler for cancel requests.
    fn cancel_handler(&self) -> Arc<impl cancel::CancelHandler> {
        Arc::new(NoopHandler)
    }
}

impl<T> PgWireServerHandlers for Arc<T>
where
    T: PgWireServerHandlers,
{
    fn simple_query_handler(&self) -> Arc<impl query::SimpleQueryHandler> {
        (**self).simple_query_handler()
    }

    fn extended_query_handler(&self) -> Arc<impl query::ExtendedQueryHandler> {
        (**self).extended_query_handler()
    }

    fn startup_handler(&self) -> Arc<impl auth::StartupHandler> {
        (**self).startup_handler()
    }

    fn copy_handler(&self) -> Arc<impl copy::CopyHandler> {
        (**self).copy_handler()
    }

    fn error_handler(&self) -> Arc<impl ErrorHandler> {
        (**self).error_handler()
    }

    fn cancel_handler(&self) -> Arc<impl cancel::CancelHandler> {
        (**self).cancel_handler()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::RwLock;

    #[test]
    fn session_extensions_get_or_insert_with() {
        let ext = SessionExtensions::new();
        let val: Arc<RwLock<Vec<i32>>> = ext.get_or_insert_with(|| RwLock::new(vec![1, 2, 3]));
        assert_eq!(*val.read().unwrap(), vec![1, 2, 3]);

        // Second call returns same instance, ignores closure
        let val2: Arc<RwLock<Vec<i32>>> = ext.get_or_insert_with(|| RwLock::new(vec![99]));
        assert_eq!(*val2.read().unwrap(), vec![1, 2, 3]);
        assert!(Arc::ptr_eq(&val, &val2));
    }

    #[test]
    fn session_extensions_different_types_are_independent() {
        let ext = SessionExtensions::new();
        ext.insert(42u32);
        ext.insert("hello".to_string());

        assert_eq!(*ext.get::<u32>().unwrap(), 42);
        assert_eq!(*ext.get::<String>().unwrap(), "hello");
    }

    #[test]
    fn session_extensions_get_returns_none_when_missing() {
        let ext = SessionExtensions::new();
        assert!(ext.get::<u64>().is_none());
    }

    #[test]
    fn session_extensions_insert_replaces_previous() {
        let ext = SessionExtensions::new();
        ext.insert(1u32);
        let old = ext.insert(2u32);
        assert_eq!(*old.unwrap(), 1);
        assert_eq!(*ext.get::<u32>().unwrap(), 2);
    }

    #[tokio::test]
    async fn connection_manager_register_and_cancel() {
        let manager = Arc::new(ConnectionManager::new());
        let key = SecretKey::I32(12345);

        let (handle, _guard) = manager.register(1, key.clone());

        let mut rx = handle.start_query().await;
        assert!(manager.cancel(1, &key).await);
        assert_eq!(rx.try_recv(), Ok(Some(())));
    }

    #[tokio::test]
    async fn connection_manager_cancel_unknown_connection() {
        let manager = Arc::new(ConnectionManager::new());
        assert!(!manager.cancel(999, &SecretKey::I32(0)).await);
    }

    #[tokio::test]
    async fn connection_manager_guard_unregisters() {
        let manager = Arc::new(ConnectionManager::new());
        let key = SecretKey::I32(99);

        {
            let (_handle, _guard) = manager.register(42, key.clone());
        }

        assert!(!manager.cancel(42, &key).await);
    }

    #[tokio::test]
    async fn connection_handle_start_query_replaces_previous() {
        let handle = ConnectionHandle::new();

        let mut rx1 = handle.start_query().await;
        let mut rx2 = handle.start_query().await;

        assert!(rx1.try_recv().is_err());
        assert!(handle.cancel().await);
        assert_eq!(rx2.try_recv(), Ok(Some(())));
    }

    #[tokio::test]
    async fn connection_handle_cancel_no_active_query() {
        let handle = ConnectionHandle::new();
        assert!(!handle.cancel().await);
    }
}