dosh/src/protocol.rs

use crate::auth::BootstrapResponse;
use crate::crypto;
use crate::native::{EnvVar, NativeAuthOk, NativeClientHello, NativeServerHello, NativeUserAuth};
use anyhow::{Context, Result, bail};
use serde::{Deserialize, Serialize};

pub const MAGIC: &[u8; 4] = b"DOSH";
pub const VERSION: u8 = 2;
pub const HEADER_LEN: usize = 58;

/// Stable, user-facing reason string the server puts in an `AttachReject` when a
/// native handshake arrives carrying a `protocol_version` it cannot speak. The
/// client recognizes this prefix and surfaces a clear "upgrade dosh" message
/// instead of the generic transport rejection or, worse, a silent timeout.
pub const VERSION_MISMATCH_REASON: &str = "protocol version mismatch — upgrade dosh";
const HEADER_AAD_LEN: usize = HEADER_LEN - 2;
pub const CLIENT_TO_SERVER: u32 = 1;
pub const SERVER_TO_CLIENT: u32 = 2;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum PacketKind {
    BootstrapAttachRequest = 1,
    TicketAttachRequest = 2,
    AttachOk = 3,
    AttachReject = 4,
    ResumeRequest = 5,
    ResumeOk = 6,
    Input = 7,
    Resize = 8,
    Frame = 9,
    Ack = 10,
    Ping = 11,
    Pong = 12,
    Detach = 13,
    NativeClientHello = 14,
    NativeServerHello = 15,
    NativeUserAuth = 16,
    NativeAuthOk = 17,
    StreamOpen = 18,
    StreamOpenOk = 19,
    StreamOpenReject = 20,
    StreamData = 21,
    StreamWindowAdjust = 22,
    StreamEof = 23,
    StreamClose = 24,
    NativeAuthCheckOk = 25,
    Rekey = 26,
    RekeyAck = 27,
}

impl TryFrom<u8> for PacketKind {
    type Error = anyhow::Error;

    fn try_from(value: u8) -> Result<Self> {
        Ok(match value {
            1 => Self::BootstrapAttachRequest,
            2 => Self::TicketAttachRequest,
            3 => Self::AttachOk,
            4 => Self::AttachReject,
            5 => Self::ResumeRequest,
            6 => Self::ResumeOk,
            7 => Self::Input,
            8 => Self::Resize,
            9 => Self::Frame,
            10 => Self::Ack,
            11 => Self::Ping,
            12 => Self::Pong,
            13 => Self::Detach,
            14 => Self::NativeClientHello,
            15 => Self::NativeServerHello,
            16 => Self::NativeUserAuth,
            17 => Self::NativeAuthOk,
            18 => Self::StreamOpen,
            19 => Self::StreamOpenOk,
            20 => Self::StreamOpenReject,
            21 => Self::StreamData,
            22 => Self::StreamWindowAdjust,
            23 => Self::StreamEof,
            24 => Self::StreamClose,
            25 => Self::NativeAuthCheckOk,
            26 => Self::Rekey,
            27 => Self::RekeyAck,
            _ => bail!("unknown packet kind {value}"),
        })
    }
}

#[derive(Debug, Clone)]
pub struct Header {
    pub kind: PacketKind,
    pub flags: u16,
    pub conn_id: [u8; 16],
    pub seq: u64,
    pub ack: u64,
    pub session_key_id: [u8; 16],
    pub body_len: u16,
}

impl Header {
    pub fn aad(&self) -> [u8; HEADER_LEN] {
        let mut out = [0u8; HEADER_LEN];
        out[..4].copy_from_slice(MAGIC);
        out[4] = VERSION;
        out[5] = self.kind as u8;
        out[6..8].copy_from_slice(&self.flags.to_be_bytes());
        out[8..24].copy_from_slice(&self.conn_id);
        out[24..32].copy_from_slice(&self.seq.to_be_bytes());
        out[32..40].copy_from_slice(&self.ack.to_be_bytes());
        out[40..56].copy_from_slice(&self.session_key_id);
        out[56..58].copy_from_slice(&self.body_len.to_be_bytes());
        out
    }

    pub fn parse(input: &[u8]) -> Result<Self> {
        if input.len() < HEADER_LEN {
            bail!("packet too short");
        }
        if &input[..4] != MAGIC {
            bail!("bad magic");
        }
        if input[4] != VERSION {
            bail!(
                "{} (peer wire protocol v{}, this build speaks v{VERSION})",
                VERSION_MISMATCH_REASON,
                input[4]
            );
        }
        let kind = PacketKind::try_from(input[5])?;
        let flags = u16::from_be_bytes(input[6..8].try_into().unwrap());
        let mut conn_id = [0u8; 16];
        conn_id.copy_from_slice(&input[8..24]);
        let seq = u64::from_be_bytes(input[24..32].try_into().unwrap());
        let ack = u64::from_be_bytes(input[32..40].try_into().unwrap());
        let mut session_key_id = [0u8; 16];
        session_key_id.copy_from_slice(&input[40..56]);
        let body_len = u16::from_be_bytes(input[56..58].try_into().unwrap());
        Ok(Self {
            kind,
            flags,
            conn_id,
            seq,
            ack,
            session_key_id,
            body_len,
        })
    }
}

/// Cheaply inspect a datagram that carries our [`MAGIC`] but whose wire
/// [`VERSION`] byte differs from this build's. Returns the peer's advertised
/// wire version when (and only when) the packet is a Dosh packet we cannot
/// otherwise decode because of a version skew, so the receiver can answer with a
/// clear, named version-mismatch reject instead of dropping it (a silent
/// timeout for the peer). Returns `None` for our own version, foreign magic, or
/// runt packets.
pub fn peek_foreign_wire_version(input: &[u8]) -> Option<u8> {
    if input.len() < 5 || &input[..4] != MAGIC || input[4] == VERSION {
        return None;
    }
    Some(input[4])
}

#[derive(Debug, Clone)]
pub struct Packet {
    pub header: Header,
    pub body: Vec<u8>,
}

pub fn encode_plain(
    kind: PacketKind,
    conn_id: [u8; 16],
    seq: u64,
    ack: u64,
    body: &[u8],
) -> Result<Vec<u8>> {
    if body.len() > u16::MAX as usize {
        bail!("packet body too large");
    }
    let header = Header {
        kind,
        flags: 0,
        conn_id,
        seq,
        ack,
        session_key_id: [0u8; 16],
        body_len: body.len() as u16,
    };
    let mut out = Vec::with_capacity(HEADER_LEN + body.len());
    out.extend_from_slice(&header.aad());
    out.extend_from_slice(body);
    Ok(out)
}

pub fn encode_encrypted(
    kind: PacketKind,
    conn_id: [u8; 16],
    seq: u64,
    ack: u64,
    key: &[u8; 32],
    direction: u32,
    plaintext: &[u8],
) -> Result<Vec<u8>> {
    let nonce = crypto::nonce_from(direction, seq);
    let header = Header {
        kind,
        flags: 1,
        conn_id,
        seq,
        ack,
        session_key_id: session_key_id(key),
        body_len: 0,
    };
    let aad_without_len = header.aad();
    let ciphertext = crypto::seal(key, &nonce, &aad_without_len[..HEADER_AAD_LEN], plaintext)?;
    if ciphertext.len() > u16::MAX as usize {
        bail!("packet body too large");
    }
    let header = Header {
        body_len: ciphertext.len() as u16,
        ..header
    };
    let mut out = Vec::with_capacity(HEADER_LEN + ciphertext.len());
    out.extend_from_slice(&header.aad());
    out.extend_from_slice(&ciphertext);
    Ok(out)
}

pub fn decode(input: &[u8]) -> Result<Packet> {
    let header = Header::parse(input)?;
    let end = HEADER_LEN + header.body_len as usize;
    if input.len() < end {
        bail!("truncated packet body");
    }
    Ok(Packet {
        header,
        body: input[HEADER_LEN..end].to_vec(),
    })
}

pub fn decrypt_body(packet: &Packet, key: &[u8; 32], direction: u32) -> Result<Vec<u8>> {
    if packet.header.flags & 1 == 0 {
        return Ok(packet.body.clone());
    }
    let nonce = crypto::nonce_from(direction, packet.header.seq);
    let aad = packet.header.aad();
    let expected_key_id = session_key_id(key);
    if packet.header.session_key_id != expected_key_id {
        bail!("stale or wrong session key id");
    }
    crypto::open(key, &nonce, &aad[..HEADER_AAD_LEN], &packet.body)
}

pub fn session_key_id(key: &[u8; 32]) -> [u8; 16] {
    let digest = crypto::sha256(key);
    let mut out = [0u8; 16];
    out.copy_from_slice(&digest[..16]);
    out
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BootstrapAttachRequest {
    pub bootstrap: BootstrapResponse,
    pub cols: u16,
    pub rows: u16,
    #[serde(default)]
    pub requested_env: Vec<EnvVar>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TicketAttachEnvelope {
    pub ticket: Vec<u8>,
    pub client_nonce: [u8; 12],
    pub ciphertext: Vec<u8>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TicketAttachBody {
    pub session: String,
    pub mode: String,
    pub cols: u16,
    pub rows: u16,
    #[serde(default)]
    pub requested_env: Vec<EnvVar>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TicketAttachOkEnvelope {
    pub server_nonce: [u8; 12],
    pub ciphertext: Vec<u8>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NativeClientHelloBody {
    pub hello: NativeClientHello,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NativeServerHelloBody {
    pub hello: NativeServerHello,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NativeUserAuthBody {
    pub auth: NativeUserAuth,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NativeAuthOkBody {
    pub ok: NativeAuthOk,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NativeAuthCheckOkBody {
    pub requested_user: String,
    pub native_auth_enabled: bool,
    pub allow_tcp_forwarding: bool,
    pub allow_remote_forwarding: bool,
    pub allow_agent_forwarding: bool,
    pub policy_flags: Vec<String>,
    #[serde(default)]
    pub server_version: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AttachOk {
    pub client_id: [u8; 16],
    pub session: String,
    pub mode: String,
    pub session_key: [u8; 32],
    pub session_key_id: [u8; 16],
    pub initial_seq: u64,
    pub snapshot: Vec<u8>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AttachReject {
    pub reason: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResumeRequest {
    pub session: String,
    pub last_rendered_seq: u64,
    pub cols: u16,
    pub rows: u16,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Input {
    pub bytes: Vec<u8>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Resize {
    pub cols: u16,
    pub rows: u16,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamOpen {
    pub stream_id: u64,
    pub target_host: String,
    pub target_port: u16,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamOpenOk {
    pub stream_id: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamOpenReject {
    pub stream_id: u64,
    pub reason: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamData {
    pub stream_id: u64,
    pub bytes: Vec<u8>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamWindowAdjust {
    pub stream_id: u64,
    pub bytes: u32,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamEof {
    pub stream_id: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamClose {
    pub stream_id: u64,
}

/// Server→client transport rekey, sealed under the *current* session key.
///
/// Carries the fresh server-generated `rekey_material` and the new `epoch`; both
/// peers feed these into [`crate::native::derive_rekey_session_key`] to compute
/// the next traffic key. `new_session_key_id` is the id the next epoch's packets
/// will carry, so the client can recognize and switch atomically. The client
/// replies with a [`PacketKind::RekeyAck`] encrypted under the *new* key.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Rekey {
    pub epoch: u64,
    pub rekey_material: [u8; 32],
    pub new_session_key_id: [u8; 16],
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Frame {
    pub session: String,
    pub output_seq: u64,
    pub bytes: Vec<u8>,
    pub snapshot: bool,
    pub closed: bool,
}

pub fn to_body<T: Serialize>(value: &T) -> Result<Vec<u8>> {
    bincode::serialize(value).context("serialize protocol body")
}

pub fn from_body<T: for<'de> Deserialize<'de>>(body: &[u8]) -> Result<T> {
    bincode::deserialize(body).context("deserialize protocol body")
}

#[derive(Debug, Clone)]
pub struct ReplayWindow {
    highest: u64,
    seen: u128,
    width: u32,
}

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

impl ReplayWindow {
    pub fn new(width: u32) -> Self {
        assert!((1..=128).contains(&width));
        Self {
            highest: 0,
            seen: 0,
            width,
        }
    }

    pub fn accept(&mut self, seq: u64) -> bool {
        if seq == 0 {
            return false;
        }
        if self.highest == 0 {
            self.highest = seq;
            self.seen = 1;
            return true;
        }
        if seq > self.highest {
            let shift = (seq - self.highest).min(128) as u32;
            self.seen = if shift >= self.width {
                1
            } else {
                ((self.seen << shift) | 1) & self.mask()
            };
            self.highest = seq;
            return true;
        }
        let offset = self.highest - seq;
        if offset >= self.width as u64 {
            return false;
        }
        let bit = 1u128 << offset;
        if self.seen & bit != 0 {
            false
        } else {
            self.seen |= bit;
            true
        }
    }

    fn mask(&self) -> u128 {
        if self.width == 128 {
            u128::MAX
        } else {
            (1u128 << self.width) - 1
        }
    }
}