libFenrir/src/connection/mod.rs

//! Connection handling and send/receive queues

pub mod handshake;
pub mod packet;
pub mod socket;
pub mod stream;

use ::core::num::Wrapping;
use ::std::{
    collections::{BTreeMap, HashMap},
    vec::Vec,
};

pub use crate::connection::{handshake::Handshake, packet::Packet};

use crate::{
    dnssec,
    enc::{
        asym::PubKey,
        hkdf::Hkdf,
        sym::{self, CipherRecv, CipherSend},
        Random,
    },
    inner::{worker, ThreadTracker},
};

/// Fenrir Connection ID
///
/// 0 is special as it represents the handshake  
/// Connection IDs are to be considered u64 little endian
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum ID {
    /// Connection id 0 represent the handshake
    Handshake,
    /// Non-zero id can represent any connection
    ID(::core::num::NonZeroU64),
}

impl ID {
    /// Set the conenction id to handshake
    pub fn new_handshake() -> Self {
        Self::Handshake
    }
    /// New id from u64. PLZ NON ZERO
    pub(crate) fn new_u64(raw: u64) -> Self {
        #[allow(unsafe_code)]
        unsafe {
            ID::ID(::core::num::NonZeroU64::new_unchecked(raw))
        }
    }
    pub(crate) fn as_u64(&self) -> u64 {
        match self {
            ID::Handshake => 0,
            ID::ID(id) => id.get(),
        }
    }
    /// New random service ID
    pub fn new_rand(rand: &Random) -> Self {
        let mut raw = [0; 8];
        let mut num = 0;
        while num == 0 {
            rand.fill(&mut raw);
            num = u64::from_le_bytes(raw);
        }
        #[allow(unsafe_code)]
        unsafe {
            ID::ID(::core::num::NonZeroU64::new_unchecked(num))
        }
    }
    /// Quick check to know if this is an handshake
    pub fn is_handshake(&self) -> bool {
        *self == ID::Handshake
    }
    /// length if the connection ID in bytes
    pub const fn len() -> usize {
        8
    }
    /// write the ID to a buffer
    pub fn serialize(&self, out: &mut [u8]) {
        match self {
            ID::Handshake => out[..8].copy_from_slice(&[0; 8]),
            ID::ID(id) => out[..8].copy_from_slice(&id.get().to_le_bytes()),
        }
    }
}

impl From<u64> for ID {
    fn from(raw: u64) -> Self {
        if raw == 0 {
            ID::Handshake
        } else {
            #[allow(unsafe_code)]
            unsafe {
                ID::ID(::core::num::NonZeroU64::new_unchecked(raw))
            }
        }
    }
}

impl From<[u8; 8]> for ID {
    fn from(raw: [u8; 8]) -> Self {
        let raw_u64 = u64::from_le_bytes(raw);
        raw_u64.into()
    }
}

/// strong typedef for receiving connection id
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct IDRecv(pub ID);
/// strong typedef for sending connection id
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct IDSend(pub ID);

/// Version of the fenrir protocol in use
#[derive(::num_derive::FromPrimitive, Debug, Copy, Clone, PartialEq)]
#[repr(u8)]
pub enum ProtocolVersion {
    /// First Fenrir Protocol Version
    V0 = 0,
}
impl ProtocolVersion {
    /// actual length of the protocol version field
    pub const fn len() -> usize {
        1
    }
    /// Serialize into raw bytes
    pub fn serialize(&self, out: &mut u8) {
        *out = *self as u8;
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd)]
pub(crate) struct UserConnTracker(Wrapping<usize>);
impl UserConnTracker {
    fn advance(&mut self) -> Self {
        let old = self.0;
        self.0 = self.0 + Wrapping(1);
        UserConnTracker(old)
    }
}

/// Connection to an Authentication Server
#[derive(Debug)]
pub struct AuthSrvConn(pub(crate) Conn);
/// Connection to a service
#[derive(Debug)]
pub struct ServiceConn(pub(crate) Conn);

/// The connection, as seen from a user of libFenrir
#[derive(Debug)]
pub struct Conn {
    pub(crate) queue: ::async_channel::Sender<worker::Work>,
    pub(crate) conn: UserConnTracker,
}

impl Conn {
    /// Queue some data to be sent in this connection
    // TODO: send_and_wait, that wait for recipient ACK
    pub async fn send(&mut self, stream: stream::ID, data: Vec<u8>) {
        use crate::inner::worker::Work;
        let _ = self
            .queue
            .send(Work::UserSend((self.conn, stream, data)))
            .await;
    }
}

/// A single connection and its data
#[derive(Debug)]
pub(crate) struct Connection {
    /// Receiving Conn ID
    pub(crate) id_recv: IDRecv,
    /// Sending Conn ID
    pub(crate) id_send: IDSend,
    /// The main hkdf used for all secrets in this connection
    hkdf: Hkdf,
    /// Cipher for decrypting data
    pub(crate) cipher_recv: CipherRecv,
    /// Cipher for encrypting data
    pub(crate) cipher_send: CipherSend,
    /// send queue for each Stream
    send_queue: BTreeMap<stream::ID, Vec<Vec<u8>>>,
}

/// Role: track the connection direction
///
/// The Role is used to select the correct secrets, and track the direction
/// of the connection
/// * Server: Conn is Incoming
/// * Client: Conn is Outgoing
#[derive(Debug, Copy, Clone)]
#[repr(u8)]
pub enum Role {
    /// Server: we receive the connection
    Server = 0,
    /// Client: we initate the connection
    Client,
}

impl Connection {
    pub(crate) fn new(
        hkdf: Hkdf,
        cipher: sym::Kind,
        role: Role,
        rand: &Random,
    ) -> Self {
        let (secret_recv, secret_send) = match role {
            Role::Server => {
                (hkdf.get_secret(b"to_server"), hkdf.get_secret(b"to_client"))
            }
            Role::Client => {
                (hkdf.get_secret(b"to_client"), hkdf.get_secret(b"to_server"))
            }
        };
        let cipher_recv = CipherRecv::new(cipher, secret_recv);
        let cipher_send = CipherSend::new(cipher, secret_send, rand);

        Self {
            id_recv: IDRecv(ID::Handshake),
            id_send: IDSend(ID::Handshake),
            hkdf,
            cipher_recv,
            cipher_send,
            send_queue: BTreeMap::new(),
        }
    }
    pub(crate) fn send(&mut self, stream: stream::ID, data: Vec<u8>) {
        let stream = match self.send_queue.get_mut(&stream) {
            None => return,
            Some(stream) => stream,
        };
        stream.push(data);
    }
}

pub(crate) struct ConnList {
    thread_id: ThreadTracker,
    connections: Vec<Option<Connection>>,
    user_tracker: BTreeMap<UserConnTracker, usize>,
    last_tracked: UserConnTracker,
    /// Bitmap to track which connection ids are used or free
    ids_used: Vec<::bitmaps::Bitmap<1024>>,
}

impl ConnList {
    pub(crate) fn new(thread_id: ThreadTracker) -> Self {
        let mut bitmap_id = ::bitmaps::Bitmap::<1024>::new();
        if thread_id.id == 0 {
            // make sure we don't count the Handshake ID
            bitmap_id.set(0, true);
        }
        const INITIAL_CAP: usize = 128;
        let mut ret = Self {
            thread_id,
            connections: Vec::with_capacity(INITIAL_CAP),
            user_tracker: BTreeMap::new(),
            last_tracked: UserConnTracker(Wrapping(0)),
            ids_used: vec![bitmap_id],
        };
        ret.connections.resize_with(INITIAL_CAP, || None);
        ret
    }
    pub fn get_mut(
        &mut self,
        tracker: UserConnTracker,
    ) -> Option<&mut Connection> {
        let idx = if let Some(idx) = self.user_tracker.get(&tracker) {
            *idx
        } else {
            return None;
        };
        match &mut self.connections[idx] {
            None => None,
            Some(conn) => Some(conn),
        }
    }
    pub fn len(&self) -> usize {
        let mut total: usize = 0;
        for bitmap in self.ids_used.iter() {
            total = total + bitmap.len()
        }
        total
    }
    /// Only *Reserve* a connection,
    /// without actually tracking it in self.connections
    pub(crate) fn reserve_first(&mut self) -> IDRecv {
        // uhm... bad things are going on here:
        // * id must be initialized, but only because:
        // * rust does not understand that after the `!found` id is always
        //   initialized
        // * `ID::new_u64` is really safe only with >0, but here it always is
        // ...we should probably rewrite it in better, safer rust
        let mut id_in_thread: usize = 0;
        let mut found = false;
        for (i, b) in self.ids_used.iter_mut().enumerate() {
            match b.first_false_index() {
                Some(idx) => {
                    b.set(idx, true);
                    id_in_thread = (i * 1024) + idx;
                    found = true;
                    break;
                }
                None => {}
            }
        }
        if !found {
            let mut new_bitmap = ::bitmaps::Bitmap::<1024>::new();
            new_bitmap.set(0, true);
            id_in_thread = self.ids_used.len() * 1024;
            self.ids_used.push(new_bitmap);
        }
        // make sure we have enough space in self.connections
        let curr_capacity = self.connections.capacity();
        if self.connections.capacity() <= id_in_thread {
            // Fill with "None", assure 64 connections without reallocations
            let multiple = 64 + curr_capacity - 1;
            let new_capacity = multiple - (multiple % curr_capacity);
            self.connections.resize_with(new_capacity, || None);
        }
        // calculate the actual connection ID
        let actual_id = ((id_in_thread as u64) * (self.thread_id.total as u64))
            + (self.thread_id.id as u64);
        let new_id = IDRecv(ID::new_u64(actual_id));
        new_id
    }
    /// NOTE: does NOT check if the connection has been previously reserved!
    pub(crate) fn track(
        &mut self,
        conn: Connection,
    ) -> Result<UserConnTracker, ()> {
        let conn_id = match conn.id_recv {
            IDRecv(ID::Handshake) => {
                return Err(());
            }
            IDRecv(ID::ID(conn_id)) => conn_id,
        };
        let id_in_thread: usize =
            (conn_id.get() / (self.thread_id.total as u64)) as usize;
        self.connections[id_in_thread] = Some(conn);
        let mut tracked;
        loop {
            tracked = self.last_tracked.advance();
            if self.user_tracker.get(&tracked).is_none() {
                let _ = self.user_tracker.insert(tracked, id_in_thread);
                break;
            }
        }
        Ok(tracked)
    }
    pub(crate) fn remove(&mut self, id: IDRecv) {
        if let IDRecv(ID::ID(raw_id)) = id {
            let id_in_thread: usize =
                (raw_id.get() / (self.thread_id.total as u64)) as usize;
            let vec_index = id_in_thread / 1024;
            let bitmask_index = id_in_thread % 1024;
            if let Some(bitmask) = self.ids_used.get_mut(vec_index) {
                bitmask.set(bitmask_index, false);
                self.connections[id_in_thread] = None;
            }
        }
    }
}

/// return wether we already have a connection, we are waiting for one, or you
/// can start one
#[derive(Debug, Clone, Copy)]
pub(crate) enum Reservation {
    /// we already have a connection. use this ID.
    Present(IDSend),
    /// we don't have a connection, but we are waiting for one to be established.
    Waiting,
    /// we have reserved a spot for your connection.
    Reserved,
}

enum MapEntry {
    Present(IDSend),
    Reserved,
}

/// Link the public key of the authentication server to a connection id
/// so that we can reuse that connection to ask for more authentications
///
/// Note that a server can have multiple public keys,
/// and the handshake will only ever verify one.
/// To avoid malicious publication fo keys that are not yours,
/// on connection we:
/// * reserve all public keys of the server
/// * wait for the connection to finish
/// * remove all those reservations, exept the one key that actually succeded
/// While searching, we return a connection ID if just one key is a match
// TODO: can we shard this per-core by hashing the pubkey? or domain? or...???
//  This needs a mutex and it will be our goeal to avoid any synchronization
pub(crate) struct AuthServerConnections {
    conn_map: HashMap<PubKey, MapEntry>,
}

impl AuthServerConnections {
    pub(crate) fn new() -> Self {
        Self {
            conn_map: HashMap::with_capacity(32),
        }
    }
    /// add an ID to the reserved spot,
    /// and unlock the other pubkeys which have not been verified
    pub(crate) fn add(
        &mut self,
        pubkey: &PubKey,
        id: IDSend,
        record: &dnssec::Record,
    ) {
        let _ = self.conn_map.insert(*pubkey, MapEntry::Present(id));
        for (_, pk) in record.public_keys.iter() {
            if pk == pubkey {
                continue;
            }
            let _ = self.conn_map.remove(pk);
        }
    }
    /// remove a dropped connection
    pub(crate) fn remove_reserved(&mut self, record: &dnssec::Record) {
        for (_, pk) in record.public_keys.iter() {
            let _ = self.conn_map.remove(pk);
        }
    }
    /// remove a dropped connection
    pub(crate) fn remove_conn(&mut self, pubkey: &PubKey) {
        let _ = self.conn_map.remove(pubkey);
    }

    /// each dnssec::Record has multiple Pubkeys. reserve and ID for them all.
    /// later on, when `add` is called we will delete
    /// those that have not actually benn used
    pub(crate) fn get_or_reserve(
        &mut self,
        record: &dnssec::Record,
    ) -> Reservation {
        for (_, pk) in record.public_keys.iter() {
            match self.conn_map.get(pk) {
                None => {}
                Some(MapEntry::Reserved) => return Reservation::Waiting,
                Some(MapEntry::Present(id)) => {
                    return Reservation::Present(id.clone())
                }
            }
        }
        for (_, pk) in record.public_keys.iter() {
            let _ = self.conn_map.insert(*pk, MapEntry::Reserved);
        }
        Reservation::Reserved
    }
}