dosh/src/auth.rs

use crate::config::{ServerConfig, expand_tilde};
use crate::crypto;
use crate::protocol;
use anyhow::{Context, Result};
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use serde::{Deserialize, Serialize};
use std::fs;
use std::io::Write;
#[cfg(unix)]
use std::os::unix::fs::OpenOptionsExt;
use std::time::{SystemTime, UNIX_EPOCH};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BootstrapResponse {
    pub protocol_version: u8,
    pub server_id: [u8; 32],
    pub server_key_epoch: u64,
    pub issued_at: u64,
    pub expires_at: u64,
    pub user: String,
    pub session: String,
    pub mode: String,
    pub terminal_size: (u16, u16),
    pub client_nonce: [u8; 12],
    pub attach_token: [u8; 32],
    pub attach_ticket: Vec<u8>,
    pub attach_ticket_psk: [u8; 32],
    pub session_key: [u8; 32],
    pub session_key_id: [u8; 16],
    pub udp_host: String,
    pub udp_port: u16,
    pub aead_algorithm: String,
}

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

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AttachTicketPlain {
    pub server_id: [u8; 32],
    pub server_key_epoch: u64,
    pub user: String,
    pub session: String,
    pub mode: String,
    pub issued_at: u64,
    pub expires_at: u64,
    pub psk: [u8; 32],
}

pub fn now_secs() -> Result<u64> {
    Ok(SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .context("system clock before UNIX epoch")?
        .as_secs())
}

pub fn load_or_create_server_secret(config: &ServerConfig) -> Result<[u8; 32]> {
    let path = expand_tilde(&config.secret_path);
    if path.exists() {
        let raw = fs::read(&path).with_context(|| format!("read {}", path.display()))?;
        let decoded = if raw.len() == 32 {
            raw
        } else {
            URL_SAFE_NO_PAD
                .decode(String::from_utf8_lossy(&raw).trim())
                .context("decode server secret")?
        };
        anyhow::ensure!(
            decoded.len() == 32,
            "server secret at {} must be 32 bytes, got {}",
            path.display(),
            decoded.len()
        );
        let mut out = [0u8; 32];
        out.copy_from_slice(&decoded[..32]);
        return Ok(out);
    }
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
    }
    let secret = crypto::random_32();
    let mut options = fs::OpenOptions::new();
    options.create_new(true).write(true);
    #[cfg(unix)]
    options.mode(0o600);
    let mut file = options
        .open(&path)
        .with_context(|| format!("create {}", path.display()))?;
    file.write_all(URL_SAFE_NO_PAD.encode(secret).as_bytes())?;
    file.write_all(b"\n")?;
    Ok(secret)
}

#[allow(clippy::too_many_arguments)]
pub fn build_bootstrap(
    config: &ServerConfig,
    secret: &[u8; 32],
    user: String,
    session: String,
    mode: String,
    terminal_size: (u16, u16),
    client_nonce: [u8; 12],
    udp_host: String,
) -> Result<BootstrapResponse> {
    let issued_at = now_secs()?;
    let expires_at = issued_at + config.auth_ttl_secs;
    let ticket_expires = issued_at + config.attach_ticket_ttl_secs;
    let server_id = crypto::sha256(secret);
    let session_key = crypto::hkdf32(
        secret,
        &client_nonce,
        format!("dosh/session/{user}/{session}/{issued_at}").as_bytes(),
    )?;
    let session_key_id = protocol::session_key_id(&session_key);
    let attach_token = attach_token(
        secret,
        &user,
        &session,
        &mode,
        terminal_size,
        &client_nonce,
        issued_at,
        expires_at,
        &session_key_id,
    );
    let attach_ticket_psk = crypto::random_32();
    let attach_ticket = build_attach_ticket(
        secret,
        server_id,
        1,
        user.clone(),
        session.clone(),
        mode.clone(),
        issued_at,
        ticket_expires,
        &attach_ticket_psk,
    )?;

    Ok(BootstrapResponse {
        protocol_version: 1,
        server_id,
        server_key_epoch: 1,
        issued_at,
        expires_at,
        user,
        session,
        mode,
        terminal_size,
        client_nonce,
        attach_token,
        attach_ticket,
        attach_ticket_psk,
        session_key,
        session_key_id,
        udp_host,
        udp_port: config.port,
        aead_algorithm: "chacha20poly1305".to_string(),
    })
}

#[allow(clippy::too_many_arguments)]
pub fn attach_token(
    secret: &[u8; 32],
    user: &str,
    session: &str,
    mode: &str,
    terminal_size: (u16, u16),
    client_nonce: &[u8; 12],
    issued_at: u64,
    expires_at: u64,
    session_key_id: &[u8; 16],
) -> [u8; 32] {
    crypto::hmac_sha256(
        secret,
        &[
            user.as_bytes(),
            session.as_bytes(),
            mode.as_bytes(),
            &terminal_size.0.to_be_bytes(),
            &terminal_size.1.to_be_bytes(),
            client_nonce,
            &issued_at.to_be_bytes(),
            &expires_at.to_be_bytes(),
            session_key_id,
        ],
    )
}

pub fn verify_bootstrap(resp: &BootstrapResponse, secret: &[u8; 32]) -> Result<bool> {
    if now_secs()? > resp.expires_at {
        return Ok(false);
    }
    let expected = attach_token(
        secret,
        &resp.user,
        &resp.session,
        &resp.mode,
        resp.terminal_size,
        &resp.client_nonce,
        resp.issued_at,
        resp.expires_at,
        &resp.session_key_id,
    );
    Ok(expected == resp.attach_token)
}

#[allow(clippy::too_many_arguments)]
pub fn build_attach_ticket(
    secret: &[u8; 32],
    server_id: [u8; 32],
    server_key_epoch: u64,
    user: String,
    session: String,
    mode: String,
    issued_at: u64,
    expires_at: u64,
    psk: &[u8; 32],
) -> Result<Vec<u8>> {
    let payload = AttachTicketPlain {
        server_id,
        server_key_epoch,
        user,
        session,
        mode,
        issued_at,
        expires_at,
        psk: *psk,
    };
    let key = attach_ticket_key(secret)?;
    let nonce = crypto::random_12();
    let encoded = bincode::serialize(&payload)?;
    let ciphertext = crypto::seal(&key, &nonce, b"dosh-sealed-attach-ticket-v1", &encoded)?;
    Ok(bincode::serialize(&SealedAttachTicket {
        nonce,
        ciphertext,
    })?)
}

pub fn verify_attach_ticket(
    secret: &[u8; 32],
    ticket_bytes: &[u8],
    psk: &[u8; 32],
    session: &str,
    mode: &str,
) -> Result<Option<AttachTicketPlain>> {
    let sealed: SealedAttachTicket = bincode::deserialize(ticket_bytes)?;
    let key = attach_ticket_key(secret)?;
    let plain = crypto::open(
        &key,
        &sealed.nonce,
        b"dosh-sealed-attach-ticket-v1",
        &sealed.ciphertext,
    )?;
    let ticket: AttachTicketPlain = bincode::deserialize(&plain)?;
    if now_secs()? > ticket.expires_at {
        return Ok(None);
    }
    if ticket.session != session || ticket.mode != mode {
        return Ok(None);
    }
    if ticket.psk != *psk {
        return Ok(None);
    }
    Ok(Some(ticket))
}

pub fn open_attach_ticket(secret: &[u8; 32], ticket_bytes: &[u8]) -> Result<AttachTicketPlain> {
    let sealed: SealedAttachTicket = bincode::deserialize(ticket_bytes)?;
    let key = attach_ticket_key(secret)?;
    let plain = crypto::open(
        &key,
        &sealed.nonce,
        b"dosh-sealed-attach-ticket-v1",
        &sealed.ciphertext,
    )?;
    Ok(bincode::deserialize(&plain)?)
}

fn attach_ticket_key(secret: &[u8; 32]) -> Result<[u8; 32]> {
    crypto::hkdf32(secret, b"dosh-ticket-key-salt-v1", b"dosh/attach-ticket/v1")
}

pub fn encode_bootstrap(resp: &BootstrapResponse) -> Result<String> {
    Ok(URL_SAFE_NO_PAD.encode(bincode::serialize(resp)?))
}

pub fn decode_bootstrap(raw: &str) -> Result<BootstrapResponse> {
    let bytes = URL_SAFE_NO_PAD.decode(raw.trim())?;
    Ok(bincode::deserialize(&bytes)?)
}