livevue-rs/src/pg_replication/proto.rs

//! Raw PostgreSQL wire-protocol connection for logical replication.
//!
//! `tokio-postgres` 0.7 does not expose a `copy_both_simple` API suitable for
//! the replication stream.  This module implements the minimum subset of the
//! PostgreSQL wire protocol needed to:
//!
//! 1. Connect with `replication=database` startup parameter.
//! 2. Perform MD5 or cleartext password authentication.
//! 3. Send simple queries (for `CREATE_REPLICATION_SLOT` and `START_REPLICATION`).
//! 4. Enter `COPY BOTH` mode and stream `CopyData` frames.
//! 5. Send `StandbyStatusUpdate` CopyData frames back to the server.
//!
//! `postgres-protocol` is used for building most frontend messages and for
//! MD5 hash computation.  The backend parser is bypassed for `CopyBothResponse`
//! (tag `'W'` / 0x57) which is not present in that crate's public API.

use super::standby_status_payload;
use anyhow::{bail, Context as _};
use bytes::{Buf, Bytes, BytesMut};
use postgres_protocol::authentication::md5_hash;
use postgres_protocol::message::frontend;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

// ---------------------------------------------------------------------------
// RawPgConn
// ---------------------------------------------------------------------------

/// A raw PostgreSQL connection suitable for replication streaming.
pub struct RawPgConn {
    stream: TcpStream,
    /// Accumulates bytes read from the server until a complete frame arrives.
    read_buf: BytesMut,
}

impl RawPgConn {
    // -----------------------------------------------------------------------
    // Constructor
    // -----------------------------------------------------------------------

    /// Connect to a PostgreSQL server in **logical replication** mode.
    ///
    /// Sends a startup message with `replication=database` so the backend
    /// enters walsender mode.  Handles MD5 and cleartext password auth.
    pub async fn connect(
        host: &str,
        port: u16,
        user: &str,
        password: &str,
        database: &str,
    ) -> anyhow::Result<Self> {
        let addr = format!("{host}:{port}");
        let stream = TcpStream::connect(&addr)
            .await
            .with_context(|| format!("TCP connect to {addr}"))?;

        let mut conn = RawPgConn {
            stream,
            read_buf: BytesMut::with_capacity(4096),
        };

        // --- Startup message -----------------------------------------------
        // Parameters include `replication=database` which instructs PostgreSQL
        // to accept replication commands on this connection.
        let mut startup_buf = BytesMut::new();
        frontend::startup_message(
            [
                ("user", user),
                ("database", database),
                ("replication", "database"),
                ("application_name", "livevue-rs"),
            ]
            .iter()
            .copied(),
            &mut startup_buf,
        )?;
        conn.stream.write_all(&startup_buf).await?;
        conn.stream.flush().await?;

        // --- Authentication loop -------------------------------------------
        loop {
            let (tag, body) = conn.read_frame().await?;
            match tag {
                b'R' => {
                    // Authentication request
                    conn.handle_auth_request(&body, user, password).await?;
                }
                b'K' => {
                    // BackendKeyData — ignore
                }
                b'S' => {
                    // ParameterStatus — ignore
                }
                b'Z' => {
                    // ReadyForQuery — handshake complete
                    break;
                }
                b'E' => {
                    let msg = parse_error_message(&body);
                    bail!("PostgreSQL startup error: {msg}");
                }
                b'N' => {
                    // NoticeResponse — ignore
                }
                other => {
                    tracing::debug!("startup: unexpected tag 0x{:02x}", other);
                }
            }
        }

        Ok(conn)
    }

    // -----------------------------------------------------------------------
    // Simple query (used for CREATE_REPLICATION_SLOT etc.)
    // -----------------------------------------------------------------------

    /// Send a simple query and drain responses until `ReadyForQuery`.
    ///
    /// Returns `Err` if the server sends an `ErrorResponse`.
    pub async fn simple_query(&mut self, query: &str) -> anyhow::Result<()> {
        let mut buf = BytesMut::new();
        frontend::query(query, &mut buf)?;
        self.stream.write_all(&buf).await?;
        self.stream.flush().await?;

        loop {
            let (tag, body) = self.read_frame().await?;
            match tag {
                b'Z' => break, // ReadyForQuery
                b'E' => {
                    let msg = parse_error_message(&body);
                    bail!("query error: {msg}");
                }
                _ => {} // CommandComplete, RowDescription, DataRow, NoticeResponse …
            }
        }
        Ok(())
    }

    // -----------------------------------------------------------------------
    // Enter COPY BOTH mode for replication streaming
    // -----------------------------------------------------------------------

    /// Send `START_REPLICATION` and absorb the `CopyBothResponse`.
    ///
    /// After this call the connection is in COPY BOTH mode.  Use
    /// [`read_copy_data`] / [`send_standby_status`] to exchange messages.
    pub async fn start_replication(&mut self, query: &str) -> anyhow::Result<()> {
        let mut buf = BytesMut::new();
        frontend::query(query, &mut buf)?;
        self.stream.write_all(&buf).await?;
        self.stream.flush().await?;

        // Expect CopyBothResponse ('W') followed by the replication stream.
        // postgres-protocol doesn't parse 'W', so we handle it manually.
        loop {
            let (tag, _body) = self.read_frame().await?;
            match tag {
                b'W' => {
                    // CopyBothResponse — body contains column-format info we
                    // don't need for the replication stream.  Just break out.
                    break;
                }
                b'E' => {
                    let msg = parse_error_message(&_body);
                    bail!("START_REPLICATION error: {msg}");
                }
                b'N' => {} // NoticeResponse — skip
                other => {
                    tracing::debug!("start_replication: unexpected tag 0x{:02x}", other);
                }
            }
        }
        Ok(())
    }

    // -----------------------------------------------------------------------
    // Replication stream I/O
    // -----------------------------------------------------------------------

    /// Read one `CopyData` payload from the replication stream.
    ///
    /// Returns:
    /// - `Ok(Some(bytes))` — a CopyData payload (XLogData or PrimaryKeepalive).
    /// - `Ok(None)` — the server sent `CopyDone` or `ErrorResponse` (stream end).
    /// - `Err(_)` — I/O or protocol error.
    pub async fn read_copy_data(&mut self) -> anyhow::Result<Option<Bytes>> {
        let (tag, body) = self.read_frame().await?;
        match tag {
            b'd' => Ok(Some(body)),
            b'c' => Ok(None), // CopyDone
            b'E' => {
                let msg = parse_error_message(&body);
                bail!("replication error: {msg}");
            }
            other => {
                tracing::debug!("copy stream: unexpected tag 0x{:02x}", other);
                Ok(Some(Bytes::new()))
            }
        }
    }

    /// Send a `StandbyStatusUpdate` CopyData frame to acknowledge `lsn`.
    pub async fn send_standby_status(&mut self, lsn: u64) -> anyhow::Result<()> {
        let payload = standby_status_payload(lsn);
        let mut buf = BytesMut::new();
        frontend::CopyData::new(&payload[..])?.write(&mut buf);
        self.stream.write_all(&buf).await?;
        self.stream.flush().await?;
        Ok(())
    }

    // -----------------------------------------------------------------------
    // Low-level frame I/O
    // -----------------------------------------------------------------------

    /// Read one complete PostgreSQL message frame.
    ///
    /// Frame format: `tag (u8) | length (i32, BE, includes itself) | body`.
    /// Returns `(tag, body_bytes)`.
    async fn read_frame(&mut self) -> anyhow::Result<(u8, Bytes)> {
        // Ensure we have at least 5 bytes (1 tag + 4 length).
        while self.read_buf.len() < 5 {
            let n = self.stream.read_buf(&mut self.read_buf).await?;
            if n == 0 {
                bail!("server closed connection unexpectedly");
            }
        }

        let tag = self.read_buf[0];
        // Length field is i32 big-endian, includes the 4-length bytes but NOT
        // the tag byte.  Total frame length = 1 (tag) + length.
        let body_len = (i32::from_be_bytes(
            self.read_buf[1..5].try_into().expect("4 byte slice"),
        ) as usize)
            .saturating_sub(4); // subtract the 4 length bytes themselves

        let total = 5 + body_len; // tag(1) + length(4) + body
        while self.read_buf.len() < total {
            let n = self.stream.read_buf(&mut self.read_buf).await?;
            if n == 0 {
                bail!("server closed connection while reading message body");
            }
        }

        self.read_buf.advance(5); // consume tag + length
        let body = self.read_buf.split_to(body_len).freeze();
        Ok((tag, body))
    }

    // -----------------------------------------------------------------------
    // Authentication
    // -----------------------------------------------------------------------

    async fn handle_auth_request(
        &mut self,
        body: &[u8],
        user: &str,
        password: &str,
    ) -> anyhow::Result<()> {
        if body.len() < 4 {
            bail!("AuthenticationRequest body too short");
        }
        let auth_type = i32::from_be_bytes(body[0..4].try_into().expect("4 bytes"));

        match auth_type {
            0 => {
                // AuthenticationOk — nothing to do
            }
            3 => {
                // AuthenticationCleartextPassword
                let mut buf = BytesMut::new();
                frontend::password_message(password.as_bytes(), &mut buf)?;
                self.stream.write_all(&buf).await?;
                self.stream.flush().await?;
            }
            5 => {
                // AuthenticationMD5Password — salt is bytes 4..8
                if body.len() < 8 {
                    bail!("MD5 auth: body too short for salt");
                }
                let salt: [u8; 4] = body[4..8].try_into().expect("4 byte salt");
                let hashed = md5_hash(user.as_bytes(), password.as_bytes(), salt);
                let mut buf = BytesMut::new();
                frontend::password_message(hashed.as_bytes(), &mut buf)?;
                self.stream.write_all(&buf).await?;
                self.stream.flush().await?;
            }
            10 => {
                // AuthenticationSASL (SCRAM-SHA-256) — not implemented here.
                // For Docker dev setups, use `md5` or `trust` in pg_hba.conf.
                bail!(
                    "SCRAM-SHA-256 authentication is not supported by the \
                     replication module.  Configure PostgreSQL to use `md5` \
                     or `trust` for this user in pg_hba.conf."
                );
            }
            other => {
                bail!("unsupported authentication type: {other}");
            }
        }

        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Extract a human-readable message from an `ErrorResponse` body.
///
/// Fields are null-terminated key-value pairs.  We extract the 'M' (message)
/// and 'D' (detail) fields.
fn parse_error_message(body: &[u8]) -> String {
    let mut parts = Vec::new();
    let mut i = 0;
    while i < body.len() {
        let field_type = body[i];
        i += 1;
        if field_type == 0 {
            break;
        }
        // Find the null terminator of the field value.
        let start = i;
        while i < body.len() && body[i] != 0 {
            i += 1;
        }
        let value = String::from_utf8_lossy(&body[start..i]).into_owned();
        i += 1; // skip null terminator
        if matches!(field_type, b'M' | b'D' | b'C') {
            parts.push(value);
        }
    }
    if parts.is_empty() {
        "unknown PostgreSQL error".into()
    } else {
        parts.join(" — ")
    }
}