fleck/src/fleck_core/crypto.rs

use std::{cell::RefCell, collections::HashMap, rc::Rc};

use aes_gcm::{
    aead::{Aead, KeyInit},
    Aes128Gcm, Nonce,
};
pub use ed25519_dalek::{Keypair, Signature, Signer};

use crate::prelude::*;

pub enum TrustLevel {
    Unknown,
    Distrusted,
    // Future extension
    // Transitively,
    Ultimately,
}

pub struct TrustDB {}

#[derive(serde::Serialize, serde::Deserialize, Default, Debug)]
pub enum PacketHeader {
    #[default]
    NoCrypto,
    Signed(Signature),
    // u128 here is the nonce
    Encrypted(u128),
}

#[derive(Default, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct PublicKey(ed25519_dalek::PublicKey);

impl PublicKey {
    pub fn wrap(pkey: ed25519_dalek::PublicKey) -> Self {
        Self(pkey)
    }
}

// *technically* there may be equivalence classes, but we really don't care about those in this
// context.
impl Eq for PublicKey {
    fn assert_receiver_is_total_eq(&self) {}
}

impl std::hash::Hash for PublicKey {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        state.write(self.0.as_bytes())
    }
}

impl std::fmt::Debug for PublicKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("PublicKey {")?;
        self.0.as_bytes().fmt(f)?;
        f.write_str("}")
    }
}

pub struct SymmetricKey {
    key: [u8; 16],
    nonce: u128,
}

pub(crate) struct VerifyPacketSignature {
    api: std::rc::Rc<crate::API>,
}

impl Service for VerifyPacketSignature {
    fn new(api: std::rc::Rc<crate::API>) -> Self {
        Self { api }
    }
    fn setup(self: &std::rc::Rc<Self>) {
        self.api
            .channel::<fleck_core::ReceivePacketChannel>()
            .sub_opt(fleck_core::ReceiveOrder::Verify, self, Self::check);
    }
}

impl VerifyPacketSignature {
    fn check(&self, mut message: fleck_core::msg::Message) -> Option<fleck_core::msg::Message> {
        // we don't care about the message if it isn't supposed to be signed
        if !message.saved_params.unwrap().signed {
            return Some(message);
        }

        if let PacketHeader::Signed(signature) = message.crypto_header {
            let pubkey = message.node.as_ref()?.pubkey();
            // reset the crypto header for verification
            message.crypto_header = PacketHeader::NoCrypto;
            let to_verify =
                bincode::serialize(&message).expect("assume that we can re-serialize the message");

            use ed25519_dalek::Verifier;
            pubkey.0.verify(&to_verify, &signature).ok()?;
            // verification passed!
            Some(message)
        } else {
            // it was supposed to have a signature, but it didn't...
            None
        }
    }
}

pub(crate) struct SignPacket {
    api: std::rc::Rc<crate::API>,
}

impl Service for SignPacket {
    fn new(api: std::rc::Rc<crate::API>) -> Self {
        Self { api }
    }
    fn setup(self: &std::rc::Rc<Self>) {
        self.api.channel::<fleck_core::SendPacketChannel>().sub_ref(
            fleck_core::SendOrder::Sign,
            self,
            Self::sign,
        );
    }
}

impl SignPacket {
    fn sign(&self, message: &mut fleck_core::msg::Message) {
        let params = message
            .saved_params
            .expect("packet got to signing stage without saved parameters?");

        // don't do anything if the packet doesn't want a signature
        if !params.signed {
            return;
        }

        self.api.with_service(|ns: &fleck_core::NodeService| {
            let ns = ns.self_node();
            let keypair = ns.keypair().unwrap();
            let sign_data = bincode::serialize(&message.content).unwrap();
            message.crypto_header = PacketHeader::Signed(keypair.sign(&sign_data));
        });
    }
}

#[derive(serde::Serialize, serde::Deserialize, Debug)]
pub struct KeyExchange {
    pubkey: ed25519_dalek::PublicKey,
    key: x25519_dalek::PublicKey,
}

impl fleck_core::msg::MessageParams for KeyExchange {
    const NAME: &'static str = "KeyExchange";
    const ENCRYPTED: bool = false;
    const SIGNED: bool = true;
}

pub enum KeyExchangeState {
    NoState,
    Given(x25519_dalek::PublicKey),
    WaitForResponse(x25519_dalek::EphemeralSecret),
    Completed(x25519_dalek::PublicKey, SymmetricKey),
}

impl std::fmt::Debug for KeyExchangeState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NoState => f.write_str("KeyExchangeState::NoState"),
            Self::Given(pkey) => f.write_fmt(format_args!(
                "KeyExchangeState::Given {{ pkey: {:?} }}",
                pkey
            )),
            Self::WaitForResponse(_) => {
                f.write_str("KeyExchangeState::WaitForResponse { redacted }")
            },
            Self::Completed(pkey, _) => f.write_fmt(format_args!(
                "KeyExchangeState::Completed {{ pkey: {:?}, redacted }}",
                pkey
            )),
        }
    }
}

#[derive(Clone)]
struct EncryptionStateKey(Rc<fleck_core::Node>);
impl PartialEq for EncryptionStateKey {
    fn eq(&self, other: &Self) -> bool {
        Rc::as_ptr(&self.0) == Rc::as_ptr(&other.0)
    }
}
impl Eq for EncryptionStateKey {
    fn assert_receiver_is_total_eq(&self) {}
}
impl std::hash::Hash for EncryptionStateKey {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        Rc::as_ptr(&self.0).hash(state);
    }
}

pub struct EncryptionService {
    api: Rc<crate::API>,
    state: RefCell<HashMap<EncryptionStateKey, KeyExchangeState>>,
}

impl Service for EncryptionService {
    fn new(api: std::rc::Rc<crate::API>) -> Self {
        Self {
            api,
            state: Default::default(),
        }
    }

    fn setup(self: &std::rc::Rc<Self>) {
        self.api.channel::<fleck_core::SendPacketChannel>().sub_opt(
            fleck_core::SendOrder::Encrypt,
            self,
            Self::encrypt,
        );
        self.api
            .channel::<fleck_core::ReceivePacketChannel>()
            .sub_opt(fleck_core::ReceiveOrder::Decrypt, self, Self::decrypt);

        self.api.with_service(|msg: &fleck_core::MessageService| {
            msg.add_message_type::<KeyExchange>();
        });
        self.api
            .channel::<fleck_core::node::NodeRegistrationChannel>()
            .sub_ref(self, Self::new_node);
        self.api
            .channel::<fleck_core::msg::MessageChannel<KeyExchange>>()
            .sub_eat(self, Self::incoming);
    }
}

impl EncryptionService {
    fn encrypt(&self, mut packet: fleck_core::Message) -> Option<fleck_core::Message> {
        if !packet.saved_params.unwrap().encrypted {
            return Some(packet);
        }

        let mut state = self.state.borrow_mut();
        let mut keyexchange = state.get_mut(&EncryptionStateKey(packet.node.as_ref().cloned()?))?;
        if let KeyExchangeState::Completed(_, symkey) = &mut keyexchange {
            let cipher = Aes128Gcm::new_from_slice(&symkey.key).expect("key length mismatch?");
            let nonce_bytes: [u8; 12] = u128::to_be_bytes(symkey.nonce)[0..12].try_into().unwrap();
            let nonce: Nonce<_> = nonce_bytes.into();

            // serialize+encrypt with the message type so we don't leak that
            let plaintext = bincode::serialize(&packet.content).ok()?;
            packet.content.ty = 0;
            packet.content.data = cipher.encrypt(&nonce, plaintext.as_slice()).ok()?;

            packet.crypto_header = PacketHeader::Encrypted(symkey.nonce);

            // advance the nonce
            symkey.nonce += 1;

            Some(packet)
        } else {
            log::info!("Tried to send packet before encryption handshake was finished");
            None
        }
    }

    fn decrypt(&self, mut packet: fleck_core::Message) -> Option<fleck_core::Message> {
        // no peer = pass along
        if packet.peer.is_none() {
            return Some(packet);
        }

        // check for an encryption header
        let nonce = match packet.crypto_header {
            PacketHeader::Encrypted(nonce) => nonce,
            _ => return Some(packet),
        };

        let mut state = self.state.borrow_mut();

        // we don't have saved_params or anything yet, or even know what the message type is, so
        // we'll just have to trust that if we have an encryption key and there's a crypto header,
        // we should decrypt it.

        // There shouldn't be anything wrong with decrypting a message that claims to be encrypted
        // and is authenticated, right?... TODO: double-check this assertion.
        if let KeyExchangeState::Completed(_, symkey) =
            state.get_mut(&EncryptionStateKey(packet.node.as_ref().cloned()?))?
        {
            let cipher = Aes128Gcm::new_from_slice(&symkey.key).expect("key length mismatch?");
            let nonce_bytes: [u8; 12] = u128::to_be_bytes(nonce)[0..12].try_into().unwrap();
            let nonce: Nonce<_> = nonce_bytes.into();

            let plaintext = cipher
                .decrypt(&nonce, packet.content.data.as_slice())
                .ok()?;
            packet.content = bincode::deserialize(&plaintext).ok()?;

            packet.crypto_header = PacketHeader::NoCrypto;

            Some(packet)
        } else {
            None
        }
    }

    fn new_node(&self, node: &mut std::rc::Rc<fleck_core::node::Node>) {
        // we may want to initiate a handshake, if we're not already doing that...
        if !self
            .state
            .borrow()
            .contains_key(&EncryptionStateKey(node.clone()))
        {
            self.begin_handshake(node)
        }
    }

    fn begin_handshake(&self, with: &Rc<fleck_core::Node>) {
        use rand::SeedableRng;
        let csprng = rand::rngs::StdRng::from_entropy();
        let secret = x25519_dalek::EphemeralSecret::new(csprng);

        // send notification with new secret
        let send = || {
            log::trace!("sending KeyExchange message");
            self.api.queue::<fleck_core::SendPacketChannel>(
                fleck_core::msg::Message::build(KeyExchange {
                    pubkey: self
                        .api
                        .with_service(|ns: &fleck_core::NodeService| ns.self_node().pubkey().0),
                    key: x25519_dalek::PublicKey::from(&secret),
                })
                .with_node(with.clone()),
            );
        };

        let mut states = self.state.borrow_mut();
        let key = EncryptionStateKey(with.clone());
        let state = states.remove(&key).unwrap_or(KeyExchangeState::NoState);
        match state {
            KeyExchangeState::NoState => {
                send();
                states.insert(key, KeyExchangeState::WaitForResponse(secret));
            },
            KeyExchangeState::Given(pubkey) => {
                send();
                let result = secret.diffie_hellman(&pubkey);

                let sk = SymmetricKey {
                    key: result.as_bytes()[0..16].try_into().unwrap(),
                    nonce: u128::from_be_bytes(result.as_bytes()[16..32].try_into().unwrap()),
                };

                log::trace!("Completed key negotiation with peer!");
                states.insert(key, KeyExchangeState::Completed(pubkey, sk));
            },
            KeyExchangeState::WaitForResponse(secret) => {
                send();
                states.insert(key, KeyExchangeState::WaitForResponse(secret));
            },
            KeyExchangeState::Completed(_pubkey, _symmetric_key) => {
                log::warn!("asked to begin handshake when already have a good encryption key --- did the peer die?");
            },
        }
    }

    fn incoming(&self, msg: (fleck_core::msg::Metadata, KeyExchange)) {
        let mut states = self.state.borrow_mut();

        // cases:
        // 1. we already know about this node, and have a Node instance for it
        // 2. we don't know about this node yet, and it has a Peer
        // 3. we don't know about this node yet, and it has no Peer

        let node = if let Some(node) = msg.0.node {
            node
        } else {
            // exciting! let's tell the node service about it, with all the info we have right now:
            // the pubkey, and potentially a peer

            self.api.with_service(|ns: &fleck_core::NodeService| {
                let pk = PublicKey::wrap(msg.1.pubkey);
                ns.inform_of(pk.clone(), msg.0.peer);
                ns.node_by_pubkey(&pk)
                    .expect("we just told the service about this node")
            })
        };

        let skey = EncryptionStateKey(node.clone());
        let state = states.remove(&skey).unwrap_or(KeyExchangeState::NoState);
        log::trace!(
            "incoming key exchange message: {:?}, state: {:?}",
            msg.1,
            &state
        );

        match state {
            KeyExchangeState::NoState => {
                states.insert(skey.clone(), KeyExchangeState::Given(msg.1.key));

                drop(states);
                self.begin_handshake(&node);
            },
            KeyExchangeState::Given(peer_pubkey) => {
                if peer_pubkey != msg.1.key {
                    log::info!("peer DH key changed!");
                    states.insert(skey.clone(), KeyExchangeState::Given(msg.1.key));
                }

                drop(states);
                self.begin_handshake(&node);
            },
            KeyExchangeState::WaitForResponse(secret) => {
                // key handshake is finished!
                let result = secret.diffie_hellman(&msg.1.key);

                let sk = SymmetricKey {
                    key: result.as_bytes()[0..16].try_into().unwrap(),
                    nonce: u128::from_be_bytes(result.as_bytes()[16..32].try_into().unwrap()),
                };

                states.insert(skey.clone(), KeyExchangeState::Completed(msg.1.key, sk));
                drop(states);
            },
            KeyExchangeState::Completed(peer_pubkey, _symmetric_key) => {
                if peer_pubkey == msg.1.key {
                    // it's a repeat, so nothing to do
                    return;
                }
                // it's a new pubkey, so we want to restart the handshake
                states.insert(skey.clone(), KeyExchangeState::Given(msg.1.key));
                drop(states);
                self.begin_handshake(&node);
            },
        }
        log::trace!(
            "after processing incoming key exchange message, state: {:?}",
            self.state.borrow().get(&skey)
        );
    }
}