nusadb 0.1.0

//! Low-level Nusa Wire Protocol codec (docs/wire-protocol.md, PROTOCOL_VERSION 1.1).
//!
//! A frame is `[type:u8][len:u32][payload]`, `len` the **total** size including the 5-byte header,
//! big-endian throughout. Strings are length-prefixed UTF-8 (`[len:u32][bytes]`, **not**
//! null-terminated). This module is pure encode/decode — no I/O policy, no SQL semantics.

use std::io::{Read, Write};

use crate::error::{Error, Result};

/// `"NUSA"` — every `Startup` payload begins with this magic (§2).
pub const PROTOCOL_MAGIC: u32 = 0x4E55_5341;
/// Wire-protocol major version. A mismatch is rejected by the server.
pub const PROTOCOL_MAJOR: u16 = 1;
/// Wire-protocol minor version. Requesting `1` opts into the typed `RowDescriptionTyped` (§9.2); a
/// 1.0 server ignores it and answers with the untyped form, which the driver still handles.
pub const PROTOCOL_MINOR: u16 = 1;
/// The frame header is `type:u8` + `len:u32`.
pub const HEADER_LEN: usize = 5;
/// Frames larger than this are rejected before allocation (§17) — bounds memory vs a crafted length.
pub const MAX_FRAME_LEN: u32 = 256 * 1024 * 1024;
/// The only SASL mechanism offered/used.
pub const SCRAM_MECHANISM: &str = "SCRAM-SHA-256";

/// Backend (server→client) message type bytes (§16.2).
#[allow(
    dead_code,
    reason = "complete protocol taxonomy; COPY bytes are reserved for COPY support"
)]
pub mod backend {
    /// Authentication (`R`): the first `u32` of the payload is the auth sub-code.
    pub const AUTH: u8 = b'R';
    /// `BackendKeyData` (`K`): `[pid:u32][secret:u32]`.
    pub const BACKEND_KEY: u8 = b'K';
    /// `ReadyForQuery` (`Z`): `[status:u8]`.
    pub const READY: u8 = b'Z';
    /// `CommandComplete` (`C`): `[tag:Str]`.
    pub const COMMAND_COMPLETE: u8 = b'C';
    /// `Error` (`E`): `[code:Str][message:Str]`.
    pub const ERROR: u8 = b'E';
    /// `RowDescription` (`T`): `[count:u16][name:Str]*` (1.0, untyped).
    pub const ROW_DESCRIPTION: u8 = b'T';
    /// `RowDescriptionTyped` (`y`): `[count:u16]([name:Str][tag:u8])*` (1.1).
    pub const ROW_DESCRIPTION_TYPED: u8 = b'y';
    /// `DataRow` (`D`): a `Fields` list.
    pub const DATA_ROW: u8 = b'D';
    /// `ParseComplete` (`1`).
    pub const PARSE_COMPLETE: u8 = b'1';
    /// `BindComplete` (`2`).
    pub const BIND_COMPLETE: u8 = b'2';
    /// `CloseComplete` (`3`).
    pub const CLOSE_COMPLETE: u8 = b'3';
    /// `ParameterDescription` (`t`): `[count:u16]`.
    pub const PARAMETER_DESCRIPTION: u8 = b't';
    /// `NoData` (`n`).
    pub const NO_DATA: u8 = b'n';
    /// `PortalSuspended` (`z`).
    pub const PORTAL_SUSPENDED: u8 = b'z';
    /// `CopyInResponse` (`G`).
    pub const COPY_IN: u8 = b'G';
    /// `CopyOutResponse` (`H`).
    pub const COPY_OUT: u8 = b'H';
    /// `CopyData` (`d`).
    pub const COPY_DATA: u8 = b'd';
    /// `CopyDone` (`c`).
    pub const COPY_DONE: u8 = b'c';
    /// `ParameterStatus` (`S`) — a `name=value` server runtime parameter, sent during the startup
    /// handshake. The driver records nothing from it, so it is consumed and skipped.
    pub const PARAMETER_STATUS: u8 = b'S';
    /// `NotificationResponse` (`A`) — an asynchronous `LISTEN`/`NOTIFY` message: `[pid:u32][channel:str]
    /// [payload:str]`. Pushed by the server while the connection is idle, so it can arrive between
    /// query responses; the driver buffers it for `notifications()`/`poll_notification`.
    pub const NOTIFICATION_RESPONSE: u8 = b'A';
}

/// Frontend (client→server) message type bytes (§16.1).
#[allow(
    dead_code,
    reason = "complete protocol taxonomy; CLOSE is reserved for explicit statement close"
)]
pub mod frontend {
    /// `Startup` (`S`).
    pub const STARTUP: u8 = b'S';
    /// `Query` (`Q`): a simple query.
    pub const QUERY: u8 = b'Q';
    /// `Parse` (`P`).
    pub const PARSE: u8 = b'P';
    /// `Bind` (`B`).
    pub const BIND: u8 = b'B';
    /// `Describe` (`D`).
    pub const DESCRIBE: u8 = b'D';
    /// `Execute` (`E`).
    pub const EXECUTE: u8 = b'E';
    /// `Sync` (`Y`).
    pub const SYNC: u8 = b'Y';
    /// `Close` (`C`).
    pub const CLOSE: u8 = b'C';
    /// `SASLInitialResponse` (`p`).
    pub const SASL_INITIAL: u8 = b'p';
    /// `SASLResponse` (`r`).
    pub const SASL_RESPONSE: u8 = b'r';
    /// `CancelRequest` (`K`) — sent in place of `Startup` on a fresh connection.
    pub const CANCEL: u8 = b'K';
    /// `Terminate` (`X`).
    pub const TERMINATE: u8 = b'X';
    /// `CopyData` (`d`) — a chunk of `COPY ... FROM STDIN` data (§12.1).
    pub const COPY_DATA: u8 = b'd';
    /// `CopyDone` (`c`) — end of the client's COPY data stream.
    pub const COPY_DONE: u8 = b'c';
    /// `CopyFail` (`f`): `[message:Str]` — abort the in-progress COPY.
    pub const COPY_FAIL: u8 = b'f';
}

/// SASL auth sub-codes (the first `u32` of an `R` payload, §7.2).
pub mod auth {
    /// Authentication succeeded.
    pub const OK: u32 = 0;
    /// Offer SASL mechanisms.
    pub const SASL: u32 = 10;
    /// SASL challenge (server-first).
    pub const SASL_CONTINUE: u32 = 11;
    /// SASL final (server verifier).
    pub const SASL_FINAL: u32 = 12;
}

/// A column's declared type, decoded from the 1-byte `RowDescriptionTyped` tag (§9.2). An
/// unrecognised tag (and `0x00`) maps to [`TypeTag::Unknown`], treated as `TEXT`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum TypeTag {
    /// `0x00` — the server could not resolve the type; treat as `TEXT`.
    Unknown,
    /// `0x01` `BOOL`.
    Bool,
    /// `0x02` `INT` (64-bit).
    Int,
    /// `0x03` `FLOAT` (IEEE-754 f64).
    Float,
    /// `0x04` `NUMERIC` (arbitrary precision; canonical text).
    Numeric,
    /// `0x05` `TEXT`.
    Text,
    /// `0x06` `BYTES`.
    Bytes,
    /// `0x07` `DATE`.
    Date,
    /// `0x08` `TIME`.
    Time,
    /// `0x09` `TIMETZ`.
    TimeTz,
    /// `0x0A` `TIMESTAMP`.
    Timestamp,
    /// `0x0B` `TIMESTAMPTZ`.
    TimestampTz,
    /// `0x0C` `INTERVAL`.
    Interval,
    /// `0x0D` `UUID`.
    Uuid,
    /// `0x0E` `JSON`.
    Json,
    /// `0x0F` `ARRAY`.
    Array,
    /// `0x10` `VECTOR`.
    Vector,
}

impl TypeTag {
    /// Map a wire tag byte to a [`TypeTag`]; unknown bytes become [`TypeTag::Unknown`].
    #[must_use]
    pub const fn from_byte(tag: u8) -> Self {
        match tag {
            0x01 => Self::Bool,
            0x02 => Self::Int,
            0x03 => Self::Float,
            0x04 => Self::Numeric,
            0x05 => Self::Text,
            0x06 => Self::Bytes,
            0x07 => Self::Date,
            0x08 => Self::Time,
            0x09 => Self::TimeTz,
            0x0A => Self::Timestamp,
            0x0B => Self::TimestampTz,
            0x0C => Self::Interval,
            0x0D => Self::Uuid,
            0x0E => Self::Json,
            0x0F => Self::Array,
            0x10 => Self::Vector,
            _ => Self::Unknown,
        }
    }

    /// The canonical uppercase name (`"INT"`, `"TEXT"`, …), as a JDBC-style `getColumnTypeName`.
    #[must_use]
    pub const fn name(self) -> &'static str {
        match self {
            Self::Unknown => "UNKNOWN",
            Self::Bool => "BOOL",
            Self::Int => "INT",
            Self::Float => "FLOAT",
            Self::Numeric => "NUMERIC",
            Self::Text => "TEXT",
            Self::Bytes => "BYTES",
            Self::Date => "DATE",
            Self::Time => "TIME",
            Self::TimeTz => "TIMETZ",
            Self::Timestamp => "TIMESTAMP",
            Self::TimestampTz => "TIMESTAMPTZ",
            Self::Interval => "INTERVAL",
            Self::Uuid => "UUID",
            Self::Json => "JSON",
            Self::Array => "ARRAY",
            Self::Vector => "VECTOR",
        }
    }
}

/// A decoded frame: its type byte plus the raw payload (the bytes after the 5-byte header).
#[derive(Debug, Clone)]
pub struct Frame {
    /// The message type byte.
    pub kind: u8,
    /// The payload (`len - 5` bytes).
    pub payload: Vec<u8>,
}

/// Builds a frame payload from primitive fields, then wraps it with the type byte + length header.
#[derive(Default)]
pub struct Writer {
    buf: Vec<u8>,
}

impl Writer {
    /// A fresh, empty payload writer.
    #[must_use]
    pub fn new() -> Self {
        Self {
            buf: Vec::with_capacity(32),
        }
    }

    /// Append a `u8`.
    pub fn u8(&mut self, v: u8) -> &mut Self {
        self.buf.push(v);
        self
    }

    /// Append a big-endian `u16`.
    pub fn u16(&mut self, v: u16) -> &mut Self {
        self.buf.extend_from_slice(&v.to_be_bytes());
        self
    }

    /// Append a big-endian `u32`.
    pub fn u32(&mut self, v: u32) -> &mut Self {
        self.buf.extend_from_slice(&v.to_be_bytes());
        self
    }

    /// Append a `Str` (`[len:u32][utf8 bytes]`, no terminator).
    pub fn str(&mut self, s: &str) -> &mut Self {
        self.u32(s.len() as u32);
        self.buf.extend_from_slice(s.as_bytes());
        self
    }

    /// Append raw bytes verbatim (no length prefix).
    pub fn raw(&mut self, bytes: &[u8]) -> &mut Self {
        self.buf.extend_from_slice(bytes);
        self
    }

    /// Append a `Fields` list (§4.3): `[count:u16]` then per field `[present:u8]` + optional
    /// `[len:u32][bytes]`. `None` is SQL `NULL`.
    pub fn fields(&mut self, values: &[Option<Vec<u8>>]) -> &mut Self {
        self.u16(values.len() as u16);
        for v in values {
            match v {
                None => {
                    self.u8(0);
                }
                Some(bytes) => {
                    self.u8(1);
                    self.u32(bytes.len() as u32);
                    self.buf.extend_from_slice(bytes);
                }
            }
        }
        self
    }

    /// Finish: prepend the `type` byte + total-length header and return the full frame bytes.
    #[must_use]
    pub fn frame(&self, kind: u8) -> Vec<u8> {
        let total = (self.buf.len() + HEADER_LEN) as u32;
        let mut out = Vec::with_capacity(self.buf.len() + HEADER_LEN);
        out.push(kind);
        out.extend_from_slice(&total.to_be_bytes());
        out.extend_from_slice(&self.buf);
        out
    }
}

/// Reads primitive fields out of a frame payload, bounds-checking every access (§4.5).
pub struct Reader<'a> {
    buf: &'a [u8],
    pos: usize,
}

impl<'a> Reader<'a> {
    /// Wrap a payload slice.
    #[must_use]
    pub fn new(buf: &'a [u8]) -> Self {
        Self { buf, pos: 0 }
    }

    /// Bytes not yet consumed.
    #[must_use]
    pub fn remaining(&self) -> usize {
        self.buf.len().saturating_sub(self.pos)
    }

    fn take(&mut self, n: usize) -> Result<&'a [u8]> {
        if self.pos + n > self.buf.len() {
            return Err(Error::Protocol(format!(
                "truncated payload: need {n} bytes, {} remain",
                self.remaining()
            )));
        }
        let slice = &self.buf[self.pos..self.pos + n];
        self.pos += n;
        Ok(slice)
    }

    /// Read a `u8`.
    pub fn u8(&mut self) -> Result<u8> {
        Ok(self.take(1)?[0])
    }

    /// Read a big-endian `u16`.
    pub fn u16(&mut self) -> Result<u16> {
        let b = self.take(2)?;
        Ok(u16::from_be_bytes([b[0], b[1]]))
    }

    /// Read a big-endian `u32`.
    pub fn u32(&mut self) -> Result<u32> {
        let b = self.take(4)?;
        Ok(u32::from_be_bytes([b[0], b[1], b[2], b[3]]))
    }

    /// Read a `Str` (`[len:u32][utf8]`), validating UTF-8.
    pub fn str(&mut self) -> Result<String> {
        let n = self.u32()? as usize;
        let bytes = self.take(n)?;
        String::from_utf8(bytes.to_vec())
            .map_err(|_| Error::Protocol("invalid UTF-8 in string field".to_owned()))
    }

    /// The rest of the payload as raw bytes (e.g. SASL server-first/final, frame-delimited).
    pub fn rest(&mut self) -> Result<&'a [u8]> {
        let n = self.remaining();
        self.take(n)
    }

    /// Read a `Fields` list (`[count:u16]` then per field present-byte + optional `[len:u32][bytes]`).
    /// Each entry is `None` (SQL NULL) or `Some(raw field bytes)`.
    pub fn fields(&mut self) -> Result<Vec<Option<Vec<u8>>>> {
        let count = self.u16()? as usize;
        // Each present field is >= 1 byte; bound the count before allocating (§4.5).
        if count > self.remaining() + 1 {
            return Err(Error::Protocol(format!(
                "field count {count} exceeds remaining payload"
            )));
        }
        let mut out = Vec::with_capacity(count);
        for _ in 0..count {
            if self.u8()? == 0 {
                out.push(None);
            } else {
                let n = self.u32()? as usize;
                out.push(Some(self.take(n)?.to_vec()));
            }
        }
        Ok(out)
    }
}

/// Read exactly one frame from `r` (blocking): the 5-byte header, then `len - 5` payload bytes.
/// Rejects `len < 5` (`MalformedFrame`) and `len > MAX_FRAME_LEN` (`FrameTooLarge`) before allocating.
pub fn read_frame<R: Read>(r: &mut R) -> Result<Frame> {
    let mut header = [0u8; HEADER_LEN];
    r.read_exact(&mut header)?;
    let kind = header[0];
    let total = u32::from_be_bytes([header[1], header[2], header[3], header[4]]);
    if total < HEADER_LEN as u32 {
        return Err(Error::Protocol(format!("malformed frame: len {total} < 5")));
    }
    if total > MAX_FRAME_LEN {
        return Err(Error::Protocol(format!(
            "frame too large: {total} > {MAX_FRAME_LEN}"
        )));
    }
    let payload_len = (total - HEADER_LEN as u32) as usize;
    let mut payload = vec![0u8; payload_len];
    r.read_exact(&mut payload)?;
    Ok(Frame { kind, payload })
}

/// Write a frame's bytes to `w` and flush.
pub fn write_frame<W: Write>(w: &mut W, frame: &[u8]) -> Result<()> {
    w.write_all(frame)?;
    w.flush()?;
    Ok(())
}