//! 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::ServiceID,
    connection::{
        self,
        handshake::{
            self, Handshake, HandshakeClient, HandshakeClientList,
            HandshakeServer,
        },
        Connection, IDRecv,
    },
    enc::{
        self,
        asym::{self, KeyID, PrivKey, PubKey},
        hkdf::{Hkdf, HkdfKind},
        sym::{CipherKind, CipherRecv},
    },
    Error,
};
use ::std::vec::Vec;

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

/// Client information needed to fully establish the conenction
#[derive(Debug)]
pub(crate) struct ClientConnectInfo {
    /// The service ID that we are connecting to
    pub service_id: ServiceID,
    /// The service ID that we are connecting to
    pub service_connection_id: IDRecv,
    /// Parsed handshake packet
    pub handshake: Handshake,
    /// Connection
    pub connection: Connection,
}
/// Intermediate actions to be taken while parsing the handshake
#[derive(Debug)]
pub(crate) 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),
}

/// Track the total number of threads and our index
/// 65K cpus should be enough for anybody
#[derive(Debug, Clone, Copy)]
pub(crate) struct ThreadTracker {
    pub total: u16,
    /// Note: starts from 1
    pub id: u16,
}

/// Tracking of handhsakes and conenctions
/// Note that we have multiple Handshake trackers, pinned to different cores
/// Each of them will handle a subset of all handshakes.
/// Each handshake is routed to a different tracker by checking
/// core = (udp_src_sender_port % total_threads) - 1
pub(crate) struct HandshakeTracker {
    thread_id: ThreadTracker,
    key_exchanges: Vec<(asym::KeyKind, asym::KeyExchangeKind)>,
    ciphers: Vec<CipherKind>,
    /// ephemeral keys used server side in key exchange
    keys_srv: Vec<HandshakeServer>,
    /// ephemeral keys used client side in key exchange
    hshake_cli: HandshakeClientList,
}

impl HandshakeTracker {
    pub(crate) fn new(thread_id: ThreadTracker) -> Self {
        Self {
            thread_id,
            ciphers: Vec::new(),
            key_exchanges: Vec::new(),
            keys_srv: Vec::new(),
            hshake_cli: HandshakeClientList::new(),
        }
    }
    pub(crate) fn new_client(
        &mut self,
        priv_key: PrivKey,
        pub_key: PubKey,
        service_id: ServiceID,
        service_conn_id: IDRecv,
        connection: Connection,
    ) -> Result<(KeyID, &mut HandshakeClient), ()> {
        self.hshake_cli.add(
            priv_key,
            pub_key,
            service_id,
            service_conn_id,
            connection,
        )
    }
    pub(crate) fn timeout_client(
        &mut self,
        key_id: KeyID,
    ) -> Option<[IDRecv; 2]> {
        if let Some(hshake) = self.hshake_cli.remove(key_id) {
            Some([hshake.connection.id_recv, hshake.service_conn_id])
        } else {
            None
        }
    }
    pub(crate) fn recv_handshake(
        &mut self,
        mut handshake: Handshake,
        handshake_raw: &mut [u8],
    ) -> Result<HandshakeAction, Error> {
        use connection::handshake::{dirsync::DirSync, HandshakeData};
        match handshake.data {
            HandshakeData::DirSync(ref mut ds) => match ds {
                DirSync::Req(ref mut req) => {
                    let ephemeral_key = {
                        if let Some(h_k) =
                            self.keys_srv.iter().find(|k| k.id == req.key_id)
                        {
                            // 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();
                    {
                        if None
                            == self.key_exchanges.iter().find(|&x| {
                                *x == (ephemeral_key.kind(), req.exchange)
                            })
                        {
                            return Err(
                                enc::Error::UnsupportedKeyExchange.into()
                            );
                        }
                    }
                    {
                        if None
                            == self.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 = Hkdf::new(HkdfKind::Sha3, 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 = match self.hshake_cli.get(resp.client_key_id) {
                        Some(hshake) => hshake,
                        None => {
                            ::tracing::debug!(
                                "No such client key id: {:?}",
                                resp.client_key_id
                            );
                            return Err(handshake::Error::UnknownKeyID.into());
                        }
                    };
                    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());
                        }
                    }
                    let hshake =
                        self.hshake_cli.remove(resp.client_key_id).unwrap();
                    return Ok(HandshakeAction::ClientConnect(
                        ClientConnectInfo {
                            service_id: hshake.service_id,
                            service_connection_id: hshake.service_conn_id,
                            handshake,
                            connection: hshake.connection,
                        },
                    ));
                }
            },
        }
    }
}