//! Per-connection dispatch and the six wire operations.

use crate::bind::TransportBindValues;
use crate::crypto::{b32, ct_eq, hmac_sha256, message_id, verify};
use crate::proto::{
    valid_username, ProtocolError, Reader, AUTH_FAILED, AUTH_REQUIRED, BAD_VERSION, CERT_LEN,
    ENVELOPE_MAGIC, ENVELOPE_MIN, ENVELOPE_VERSION, ID_LEN, KEY_LEN, LABEL_AUTH, LABEL_MAC,
    LABEL_REGISTER, LABEL_ROTATE, MAC_LEN, MALFORMED, MAX_CHAIN, NOT_PERMITTED, OK, OP_AUTH,
    OP_DELETE, OP_FETCH, OP_REGISTER, OP_RESOLVE, OP_SEND, QUOTA_EXCEEDED, RATE_LIMITED, TOKEN_LEN,
    TOO_LARGE, UNKNOWN_USER,
};
use crate::ratelimit::RateLimiter;
use crate::store::Store;

pub struct ServerConfig {
    pub max_envelope: usize,
    pub fetch_budget: usize,
    pub main_quota: i64,
    pub requests_quota: i64,
    pub max_tokens: u16,
    pub invite_token: Option<Vec<u8>>,
    pub conn_limiter: RateLimiter,
    pub send_limiter: RateLimiter,
    pub token_limiter: RateLimiter,
}

/// A parse failure (-> MALFORMED) or a storage failure (-> INTERNAL_ERROR).
pub enum HandlerError {
    Protocol(ProtocolError),
    Store(rusqlite::Error),
}

impl From<ProtocolError> for HandlerError {
    fn from(e: ProtocolError) -> Self {
        HandlerError::Protocol(e)
    }
}

impl From<rusqlite::Error> for HandlerError {
    fn from(e: rusqlite::Error) -> Self {
        HandlerError::Store(e)
    }
}

type OpResult = Result<(u8, Vec<u8>), HandlerError>;

/// Wire op byte to log name; unknown bytes are logged as dispatched too.
fn op_name(op: u8) -> &'static str {
    match op {
        OP_AUTH => "AUTH",
        OP_RESOLVE => "RESOLVE",
        OP_SEND => "SEND",
        OP_FETCH => "FETCH",
        OP_DELETE => "DELETE",
        OP_REGISTER => "REGISTER",
        _ => "UNKNOWN",
    }
}

pub struct Session<'a> {
    config: &'a ServerConfig,
    store: Store,
    peer_ip: String,
    bind: TransportBindValues,
    username: Option<String>,
}

impl<'a> Session<'a> {
    pub fn new(
        config: &'a ServerConfig,
        store: Store,
        peer_ip: String,
        bind: TransportBindValues,
    ) -> Self {
        Session {
            config,
            store,
            peer_ip,
            bind,
            username: None,
        }
    }

    /// Dispatches one frame, always producing a status to send back, never
    /// panicking or propagating errors to the caller: a bad frame or a
    /// storage error both become a response, and the caller decides
    /// separately whether to keep the connection open.
    ///
    /// Timing lives in this wrapper because both carriers dispatch through
    /// it, so the metrics flag covers TCP and RNS from one place.
    pub fn dispatch(&mut self, op: u8, body: &[u8]) -> (u8, Vec<u8>) {
        let started = std::time::Instant::now();
        let (status, payload) = self.dispatch_op(op, body);
        log::debug!(
            "op {} from {} in {:?}: status {}, {}B request, {}B response",
            op_name(op),
            self.peer_ip,
            started.elapsed(),
            status,
            body.len(),
            payload.len()
        );
        (status, payload)
    }

    fn dispatch_op(&mut self, op: u8, body: &[u8]) -> (u8, Vec<u8>) {
        if matches!(op, OP_FETCH | OP_DELETE) && self.username.is_none() {
            return (AUTH_REQUIRED, Vec::new());
        }
        let mut r = Reader::new(body);
        let result = match op {
            OP_AUTH => self.op_auth(&mut r),
            OP_RESOLVE => self.op_resolve(&mut r),
            OP_SEND => self.op_send(&mut r),
            OP_FETCH => self.op_fetch(&mut r),
            OP_DELETE => self.op_delete(&mut r),
            OP_REGISTER => self.op_register(&mut r),
            _ => return (MALFORMED, Vec::new()),
        };
        match result {
            Ok(response) => response,
            Err(HandlerError::Protocol(e)) => {
                log::info!("bad body from {}: {}", self.peer_ip, e);
                (MALFORMED, Vec::new())
            }
            Err(HandlerError::Store(e)) => {
                log::error!("storage error from {}: {}", self.peer_ip, e);
                (crate::proto::INTERNAL_ERROR, Vec::new())
            }
        }
    }

    fn read_str(r: &mut Reader) -> Result<String, ProtocolError> {
        let len = r.u8()? as usize;
        let bytes = r.take(len)?;
        std::str::from_utf8(bytes)
            .map(str::to_string)
            .map_err(|_| ProtocolError::new("invalid utf-8"))
    }

    fn op_auth(&mut self, r: &mut Reader) -> OpResult {
        let username = Self::read_str(r)?;
        let identity = r.take(KEY_LEN)?.to_vec();
        let signature = r.take(64)?.to_vec();
        let sync = r.u8()?;
        let count = r.u16()? as usize;
        let mut tokens = Vec::with_capacity(count);
        for _ in 0..count {
            tokens.push(r.take(TOKEN_LEN)?.to_vec());
        }
        r.done()?;

        // sync = 0 leaves the stored set untouched, so it carries no tokens.
        if sync > 1 || (sync == 0 && count != 0) {
            return Ok((MALFORMED, Vec::new()));
        }

        let bound = self.store.identity_of(&username)?;
        // A wrong username and a wrong signature are both AUTH_FAILED: telling
        // them apart would turn this into an account-existence oracle.
        if bound.as_deref() != Some(identity.as_slice()) {
            return Ok((AUTH_FAILED, Vec::new()));
        }
        let mut msg = LABEL_AUTH.to_vec();
        msg.extend_from_slice(&self.bind.h);
        if !verify(&identity, &signature, &msg) {
            return Ok((AUTH_FAILED, Vec::new()));
        }

        if sync == 1 {
            if count > self.config.max_tokens as usize {
                // Leaves the session unauthenticated, per SPEC.md sec 4.
                return Ok((TOO_LARGE, Vec::new()));
            }
            self.store.set_tokens(&username, &tokens)?;
        }

        let accepted = self.store.token_count(&username)?;
        self.username = Some(username);
        Ok((OK, accepted.to_be_bytes().to_vec()))
    }

    fn op_resolve(&mut self, r: &mut Reader) -> OpResult {
        let username = Self::read_str(r)?;
        r.done()?;
        let identity = match self.store.identity_of(&username)? {
            Some(id) => id,
            None => return Ok((UNKNOWN_USER, Vec::new())),
        };
        let chain = self.store.chain(&username)?;
        let mut out = identity;
        out.push(chain.len() as u8);
        for cert in chain {
            out.extend_from_slice(&cert);
        }
        Ok((OK, out))
    }

    fn op_send(&mut self, r: &mut Reader) -> OpResult {
        let mac_len = r.u8()? as usize;
        if mac_len != 0 && mac_len != MAC_LEN {
            return Ok((MALFORMED, Vec::new()));
        }
        let mac = r.take(mac_len)?.to_vec();
        let envelope = r.rest().to_vec();

        if !self.config.send_limiter.allow(&self.peer_ip) {
            return Ok((RATE_LIMITED, Vec::new()));
        }
        if envelope.len() > self.config.max_envelope {
            return Ok((TOO_LARGE, Vec::new()));
        }
        if envelope.len() < ENVELOPE_MIN || &envelope[..4] != ENVELOPE_MAGIC {
            return Ok((MALFORMED, Vec::new()));
        }
        if envelope[4] != ENVELOPE_VERSION {
            return Ok((BAD_VERSION, Vec::new()));
        }
        let recipient = &envelope[5..37];
        let username = match self.store.username_for_key(recipient)? {
            Some(u) => u,
            None => return Ok((UNKNOWN_USER, Vec::new())),
        };

        // The ciphertext is never inspected; the server cannot read it.
        let mid = message_id(&envelope);

        // A matching accept token (SPEC.md sec 5.8) puts the envelope in the
        // mailbox's main tier; anything else lands in the smaller, short-lived
        // requests tier. A failed match is never reported to the sender.
        let matched_token = if mac.len() == MAC_LEN {
            let mut msg = LABEL_MAC.to_vec();
            msg.extend_from_slice(&mid);
            self.store
                .tokens_of(&username)?
                .into_iter()
                .find(|t| ct_eq(&hmac_sha256(t, &msg), &mac))
        } else {
            None
        };

        if let Some(token) = &matched_token {
            if !self.config.token_limiter.allow(&b32(token)) {
                return Ok((RATE_LIMITED, Vec::new()));
            }
        }

        let unsolicited = matched_token.is_none();
        let tier_quota = if unsolicited {
            self.config.requests_quota
        } else {
            self.config.main_quota
        };
        let keys = self.store.keys_of(&username)?;
        let used = self.store.mailbox_bytes(&keys, unsolicited)?;
        if used + envelope.len() as i64 > tier_quota {
            return Ok((QUOTA_EXCEEDED, Vec::new()));
        }

        self.store
            .store_message(&mid, recipient, &envelope, unsolicited)?;
        Ok((OK, mid.to_vec()))
    }

    fn op_fetch(&mut self, r: &mut Reader) -> OpResult {
        let after_received_at = r.i64()?;
        let after_id = r.take(ID_LEN)?.to_vec();
        r.done()?;
        let username = self.username.as_ref().expect("AUTH_REQUIRED gate above");
        let keys = self.store.keys_of(username)?;
        let records = self.store.pending(
            &keys,
            after_received_at,
            &after_id,
            self.config.fetch_budget,
        )?;

        let mut out = Vec::new();
        out.extend_from_slice(&(records.len() as u16).to_be_bytes());
        for (mid, received_at, unsolicited, envelope) in records {
            out.extend_from_slice(&mid);
            out.extend_from_slice(&received_at.to_be_bytes());
            out.push(if unsolicited { 1 } else { 0 });
            out.extend_from_slice(&(envelope.len() as u32).to_be_bytes());
            out.extend_from_slice(&envelope);
        }
        Ok((OK, out))
    }

    fn op_delete(&mut self, r: &mut Reader) -> OpResult {
        let count = r.u16()? as usize;
        let mut ids = Vec::with_capacity(count);
        for _ in 0..count {
            ids.push(r.take(ID_LEN)?.to_vec());
        }
        r.done()?;
        let username = self.username.as_ref().expect("AUTH_REQUIRED gate above");

        if ids.is_empty() {
            return Ok((OK, 0u16.to_be_bytes().to_vec()));
        }
        // Scoped to the caller's own keys, so ids cannot be used to probe or
        // delete another mailbox.
        let keys = self.store.keys_of(username)?;
        let removed = self.store.delete(&keys, &ids)?;
        Ok((OK, (removed as u16).to_be_bytes().to_vec()))
    }

    fn op_register(&mut self, r: &mut Reader) -> OpResult {
        let username = Self::read_str(r)?;
        let identity = r.take(KEY_LEN)?.to_vec();
        let signature = r.take(64)?.to_vec();
        let token_len = r.u8()? as usize;
        let token = r.take(token_len)?.to_vec();
        let cert_len = r.u8()? as usize;
        let cert = r.take(cert_len)?.to_vec();
        r.done()?;

        if !valid_username(&username) {
            return Ok((MALFORMED, Vec::new()));
        }
        // identity is exactly KEY_LEN bytes by construction (Reader::take
        // enforces it); no separate curve-point validity check is needed.

        // Proof of possession, required on every registration and rotation
        // (SPEC.md sec 6.1). Binding server_static stops the attestation from
        // being replayed against another server.
        let mut pop_msg = LABEL_REGISTER.to_vec();
        pop_msg.extend_from_slice(&self.bind.server_static);
        pop_msg.extend_from_slice(username.as_bytes());
        pop_msg.extend_from_slice(&identity);
        if !verify(&identity, &signature, &pop_msg) {
            return Ok((AUTH_FAILED, Vec::new()));
        }

        if let Some(expected) = &self.config.invite_token {
            if &token != expected {
                return Ok((NOT_PERMITTED, Vec::new()));
            }
        }

        if cert.is_empty() {
            if !self.store.register(&username, &identity)? {
                return Ok((NOT_PERMITTED, Vec::new()));
            }
            return Ok((OK, Vec::new()));
        }

        if cert.len() != CERT_LEN {
            return Ok((MALFORMED, Vec::new()));
        }
        let old_pub = &cert[..32];
        let new_pub = &cert[32..64];
        let when = &cert[64..72];
        let sig_old = &cert[72..136];
        let sig_new = &cert[136..200];
        if new_pub != identity.as_slice() {
            return Ok((MALFORMED, Vec::new()));
        }
        let bound = match self.store.identity_of(&username)? {
            Some(b) => b,
            None => return Ok((UNKNOWN_USER, Vec::new())),
        };
        // Only the currently bound key may hand the username on.
        if bound != old_pub {
            return Ok((NOT_PERMITTED, Vec::new()));
        }
        let mut msg = LABEL_ROTATE.to_vec();
        msg.extend_from_slice(username.as_bytes());
        msg.extend_from_slice(old_pub);
        msg.extend_from_slice(new_pub);
        msg.extend_from_slice(when);
        if !verify(old_pub, sig_old, &msg) || !verify(new_pub, sig_new, &msg) {
            return Ok((AUTH_FAILED, Vec::new()));
        }
        let chain = self.store.chain(&username)?;
        if chain.len() >= MAX_CHAIN {
            return Ok((NOT_PERMITTED, Vec::new()));
        }
        if !self.store.rotate(&username, new_pub, &cert, chain.len())? {
            return Ok((NOT_PERMITTED, Vec::new()));
        }
        Ok((OK, Vec::new()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::crypto::sha256;
    use ed25519_dalek::{Signer, SigningKey};

    fn test_config() -> ServerConfig {
        ServerConfig {
            max_envelope: 1024,
            fetch_budget: 512,
            main_quota: 1 << 20,
            requests_quota: 1 << 20,
            max_tokens: 16,
            invite_token: None,
            conn_limiter: RateLimiter::new(0),
            send_limiter: RateLimiter::new(0),
            token_limiter: RateLimiter::new(0),
        }
    }

    fn temp_store(name: &str) -> Store {
        let path =
            std::env::temp_dir().join(format!("bunshin-session-{name}-{}.db", std::process::id()));
        let _ = std::fs::remove_file(&path);
        Store::open(path.to_str().unwrap()).unwrap()
    }

    fn str_field(s: &str) -> Vec<u8> {
        let mut out = vec![s.len() as u8];
        out.extend_from_slice(s.as_bytes());
        out
    }

    fn register_body(username: &str, identity: &[u8], signature: &[u8]) -> Vec<u8> {
        let mut body = str_field(username);
        body.extend_from_slice(identity);
        body.extend_from_slice(signature);
        body.push(0); // invite token, none configured
        body.push(0); // cert, plain registration
        body
    }

    fn auth_body(username: &str, identity: &[u8], signature: &[u8]) -> Vec<u8> {
        let mut body = str_field(username);
        body.extend_from_slice(identity);
        body.extend_from_slice(signature);
        body.push(0); // sync = 0, stored tokens untouched
        body.extend_from_slice(&0u16.to_be_bytes());
        body
    }

    #[test]
    fn rns_bind_signatures_accepted_and_tcp_bind_signatures_rejected() {
        let config = test_config();
        let key = SigningKey::from_bytes(&[3u8; 32]);
        let identity = key.verifying_key().as_bytes().to_vec();

        let bind = TransportBindValues::rns(&[0x11; 16], &[0x22; 16]);

        let mut pop = LABEL_REGISTER.to_vec();
        pop.extend_from_slice(&bind.server_static);
        pop.extend_from_slice(b"alice");
        pop.extend_from_slice(&identity);
        let good_register = register_body("alice", &identity, &key.sign(&pop).to_bytes());

        // A TCP client would sign over the real X25519 static key instead.
        let tcp_static = crate::crypto::derive_public(&[9u8; 32]);
        let mut tcp_pop = LABEL_REGISTER.to_vec();
        tcp_pop.extend_from_slice(&tcp_static);
        tcp_pop.extend_from_slice(b"alice");
        tcp_pop.extend_from_slice(&identity);
        let tcp_register = register_body("alice", &identity, &key.sign(&tcp_pop).to_bytes());

        let mut auth_msg = LABEL_AUTH.to_vec();
        auth_msg.extend_from_slice(&bind.h);
        let good_auth = auth_body("alice", &identity, &key.sign(&auth_msg).to_bytes());

        // A Noise client would sign over the handshake hash instead.
        let handshake_hash = sha256(&[b"handshake"]);
        let mut tcp_auth = LABEL_AUTH.to_vec();
        tcp_auth.extend_from_slice(&handshake_hash);
        let tcp_auth = auth_body("alice", &identity, &key.sign(&tcp_auth).to_bytes());

        let mut session = Session::new(&config, temp_store("rns-bind"), "rns".into(), bind);
        assert_eq!(session.dispatch(OP_REGISTER, &good_register).0, OK);
        assert_eq!(session.dispatch(OP_AUTH, &good_auth).0, OK);

        let mut session = Session::new(
            &config,
            temp_store("rns-bind-reject"),
            "rns".into(),
            TransportBindValues::rns(&[0x11; 16], &[0x22; 16]),
        );
        assert_eq!(session.dispatch(OP_REGISTER, &good_register).0, OK);
        assert_eq!(session.dispatch(OP_AUTH, &tcp_auth).0, AUTH_FAILED);

        let mut session = Session::new(
            &config,
            temp_store("rns-bind-tcp-register"),
            "rns".into(),
            TransportBindValues::rns(&[0x11; 16], &[0x22; 16]),
        );
        assert_eq!(session.dispatch(OP_REGISTER, &tcp_register).0, AUTH_FAILED);
    }

    #[test]
    fn tcp_bind_signatures_accepted_and_rns_bind_signatures_rejected() {
        let config = test_config();
        let key = SigningKey::from_bytes(&[4u8; 32]);
        let identity = key.verifying_key().as_bytes().to_vec();

        let server_static = crate::crypto::derive_public(&[9u8; 32]);
        let handshake_hash = sha256(&[b"handshake"]);
        let bind = TransportBindValues::tcp(&handshake_hash, &server_static).unwrap();

        let mut pop = LABEL_REGISTER.to_vec();
        pop.extend_from_slice(&server_static);
        pop.extend_from_slice(b"alice");
        pop.extend_from_slice(&identity);
        let good_register = register_body("alice", &identity, &key.sign(&pop).to_bytes());

        let mut auth_msg = LABEL_AUTH.to_vec();
        auth_msg.extend_from_slice(&handshake_hash);
        let good_auth = auth_body("alice", &identity, &key.sign(&auth_msg).to_bytes());

        // An RNS client would sign over the derived bind values instead.
        let rns_bind = TransportBindValues::rns(&[0x11; 16], &[0x22; 16]);
        let mut rns_pop = LABEL_REGISTER.to_vec();
        rns_pop.extend_from_slice(&rns_bind.server_static);
        rns_pop.extend_from_slice(b"alice");
        rns_pop.extend_from_slice(&identity);
        let rns_register = register_body("alice", &identity, &key.sign(&rns_pop).to_bytes());

        let mut rns_auth = LABEL_AUTH.to_vec();
        rns_auth.extend_from_slice(&rns_bind.h);
        let rns_auth = auth_body("alice", &identity, &key.sign(&rns_auth).to_bytes());

        let mut session = Session::new(&config, temp_store("tcp-bind"), "tcp".into(), bind);
        assert_eq!(session.dispatch(OP_REGISTER, &good_register).0, OK);
        assert_eq!(session.dispatch(OP_AUTH, &good_auth).0, OK);

        let mut session = Session::new(
            &config,
            temp_store("tcp-bind-reject"),
            "tcp".into(),
            TransportBindValues::tcp(&handshake_hash, &server_static).unwrap(),
        );
        assert_eq!(session.dispatch(OP_REGISTER, &good_register).0, OK);
        assert_eq!(session.dispatch(OP_AUTH, &rns_auth).0, AUTH_FAILED);

        let mut session = Session::new(
            &config,
            temp_store("tcp-bind-rns-register"),
            "tcp".into(),
            TransportBindValues::tcp(&handshake_hash, &server_static).unwrap(),
        );
        assert_eq!(session.dispatch(OP_REGISTER, &rns_register).0, AUTH_FAILED);
    }
}