//! Inner Fenrir tracking
//! This is meant to be **async-free** so that others might use it
//! without the tokio runtime

pub(crate) mod worker;

use crate::{
    auth,
    connection::{
        self,
        handshake::{self, Handshake, HandshakeClient, HandshakeServer},
        Connection,
    },
    enc::{
        self, asym,
        hkdf::HkdfSha3,
        sym::{CipherKind, CipherRecv, CipherSend, HeadLen, TagLen},
    },
    Error,
};
use ::arc_swap::{ArcSwap, ArcSwapAny, ArcSwapOption};
use ::std::{sync::Arc, vec::Vec};

/// Information needed to reply after the key exchange
#[derive(Debug, Clone)]
pub struct AuthNeededInfo {
    /// Parsed handshake packet
    pub handshake: Handshake,
    /// hkdf generated from the handshake
    pub hkdf: HkdfSha3,
    /// cipher to be used in both directions
    pub cipher: CipherKind,
}

/// Client information needed to fully establish the conenction
#[derive(Debug)]
pub struct ClientConnectInfo {
    /// Parsed handshake packet
    pub handshake: Handshake,
    /// Connection
    pub connection: Arc<Connection>,
}
/// Intermediate actions to be taken while parsing the handshake
#[derive(Debug)]
pub enum HandshakeAction {
    /// Parsing finished, all ok, nothing to do
    None,
    /// Packet parsed, now go perform authentication
    AuthNeeded(AuthNeededInfo),
    /// the client can fully establish a connection with this info
    ClientConnect(ClientConnectInfo),
}

/// Async free but thread safe tracking of handhsakes and conenctions
pub struct HandshakeTracker {
    key_exchanges: ArcSwapAny<Arc<Vec<(asym::Key, asym::KeyExchange)>>>,
    ciphers: ArcSwapAny<Arc<Vec<CipherKind>>>,
    /// ephemeral keys used server side in key exchange
    keys_srv: ArcSwapAny<Arc<Vec<HandshakeServer>>>,
    /// ephemeral keys used client side in key exchange
    hshake_cli: ArcSwapAny<Arc<Vec<HandshakeClient>>>,
}
#[allow(unsafe_code)]
unsafe impl Send for HandshakeTracker {}
#[allow(unsafe_code)]
unsafe impl Sync for HandshakeTracker {}

impl HandshakeTracker {
    pub fn new() -> Self {
        Self {
            ciphers: ArcSwapAny::new(Arc::new(Vec::new())),
            key_exchanges: ArcSwapAny::new(Arc::new(Vec::new())),
            keys_srv: ArcSwapAny::new(Arc::new(Vec::new())),
            hshake_cli: ArcSwapAny::new(Arc::new(Vec::new())),
        }
    }
    pub(crate) fn recv_handshake(
        &self,
        mut handshake: Handshake,
        handshake_raw: &mut [u8],
    ) -> Result<HandshakeAction, Error> {
        use connection::handshake::{
            dirsync::{self, DirSync},
            HandshakeData,
        };
        match handshake.data {
            HandshakeData::DirSync(ref mut ds) => match ds {
                DirSync::Req(ref mut req) => {
                    let ephemeral_key = {
                        // Keep this block short to avoid contention
                        // on self.keys_srv
                        let keys = self.keys_srv.load();
                        if let Some(h_k) =
                            keys.iter().find(|k| k.id == req.key_id)
                        {
                            use enc::asym::PrivKey;
                            // Directory synchronized can only use keys
                            // for key exchange, not signing keys
                            if let PrivKey::Exchange(k) = &h_k.key {
                                Some(k.clone())
                            } else {
                                None
                            }
                        } else {
                            None
                        }
                    };
                    if ephemeral_key.is_none() {
                        ::tracing::debug!(
                            "No such server key id: {:?}",
                            req.key_id
                        );
                        return Err(handshake::Error::UnknownKeyID.into());
                    }
                    let ephemeral_key = ephemeral_key.unwrap();
                    {
                        let exchanges = self.key_exchanges.load();
                        if None
                            == exchanges.iter().find(|&x| {
                                *x == (ephemeral_key.kind(), req.exchange)
                            })
                        {
                            return Err(
                                enc::Error::UnsupportedKeyExchange.into()
                            );
                        }
                    }
                    {
                        let ciphers = self.ciphers.load();
                        if None == ciphers.iter().find(|&x| *x == req.cipher) {
                            return Err(enc::Error::UnsupportedCipher.into());
                        }
                    }
                    let shared_key = match ephemeral_key
                        .key_exchange(req.exchange, req.exchange_key)
                    {
                        Ok(shared_key) => shared_key,
                        Err(e) => return Err(handshake::Error::Key(e).into()),
                    };
                    let hkdf = HkdfSha3::new(b"fenrir", shared_key);
                    let secret_recv = hkdf.get_secret(b"to_server");
                    let cipher_recv = CipherRecv::new(req.cipher, secret_recv);
                    use crate::enc::sym::AAD;
                    let aad = AAD(&mut []); // no aad for now
                    match cipher_recv.decrypt(
                        aad,
                        &mut handshake_raw[req.encrypted_offset()..],
                    ) {
                        Ok(cleartext) => {
                            req.data.deserialize_as_cleartext(cleartext)
                        }
                        Err(e) => {
                            return Err(handshake::Error::Key(e).into());
                        }
                    }

                    let cipher = req.cipher;

                    return Ok(HandshakeAction::AuthNeeded(AuthNeededInfo {
                        handshake,
                        hkdf,
                        cipher,
                    }));
                }
                DirSync::Resp(resp) => {
                    let hshake = {
                        // Keep this block short to avoid contention
                        // on self.hshake_cli
                        let hshake_cli_lock = self.hshake_cli.load();
                        match hshake_cli_lock
                            .iter()
                            .find(|h| h.id == resp.client_key_id)
                        {
                            Some(h) => Some(h.clone()),
                            None => None,
                        }
                    };
                    if hshake.is_none() {
                        ::tracing::debug!(
                            "No such client key id: {:?}",
                            resp.client_key_id
                        );
                        return Err(handshake::Error::UnknownKeyID.into());
                    }
                    let hshake = hshake.unwrap();
                    let cipher_recv = &hshake.connection.cipher_recv;
                    use crate::enc::sym::AAD;
                    // no aad for now
                    let aad = AAD(&mut []);
                    let mut raw_data = &mut handshake_raw[resp
                        .encrypted_offset()
                        ..(resp.encrypted_offset() + resp.encrypted_length())];
                    match cipher_recv.decrypt(aad, &mut raw_data) {
                        Ok(cleartext) => {
                            resp.data.deserialize_as_cleartext(&cleartext)
                        }
                        Err(e) => {
                            return Err(handshake::Error::Key(e).into());
                        }
                    }
                    return Ok(HandshakeAction::ClientConnect(
                        ClientConnectInfo {
                            handshake,
                            connection: hshake.connection,
                        },
                    ));
                }
            },
        }
    }
}