Add automatic encryption layer based on the key exchange.

src/fleck_core.rs

@@ -38,7 +38,7 @@ pub type SendPacketChannel = (channel_tags::SendPacketTag, msg::Message, SendOrd
 pub enum ReceiveOrder {
     /// Annotate with various properties based on message source
     Annotate,
-    Decryption,
+    Decrypt,
     Verify,
     Parse,
     Dispatch,

src/fleck_core/crypto.rs

@@ -1,6 +1,7 @@
 use std::{cell::RefCell, collections::HashMap};
 
-use aes_gcm::aes::cipher::BlockEncrypt;
+// use aes_gcm::aes::cipher::BlockEncrypt;
+use aes_gcm::{Aes128Gcm, aead::{KeyInit, Aead}, Nonce};
 pub use ed25519_dalek::{Keypair, PublicKey, Signature, Signer};
 
 use crate::prelude::*;
@@ -158,6 +159,11 @@ impl Service for EncryptionService {
     }
 
     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>();
         });
@@ -171,6 +177,73 @@ impl Service for EncryptionService {
 }
 
 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(packet.peer.as_ref()?)?;
+        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(packet.peer.as_ref()?)? {
+            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 let Some(peer) = node.peer().as_ref() {

src/test/encryption.rs

@@ -0,0 +1,91 @@
+use std::{cell::Cell, rc::Rc};
+
+use crate::fleck_core;
+
+const SECRET : u32 = 0x12345678;
+
+#[derive(Debug, serde::Serialize, serde::Deserialize)]
+struct EncryptedMessage {
+    secret_payload: u32
+}
+
+impl fleck_core::MessageParams for EncryptedMessage {
+    const NAME: &'static str = "EncryptedMessage";
+    const SIGNED: bool = false;
+    const ENCRYPTED: bool = true;
+}
+
+#[derive(Default)]
+struct ReceiveService {
+    counter: Cell<u32>,
+}
+
+impl crate::DefaultService for ReceiveService {
+    fn setup(self: &Rc<Self>, api: Rc<crate::API>) {
+        api.channel::<fleck_core::MessageChannel<EncryptedMessage>>().sub_eat(self, Self::message);
+    }
+}
+
+impl ReceiveService {
+    fn message(&self, msg: (fleck_core::msg::Metadata, EncryptedMessage)) {
+        if msg.1 .secret_payload == SECRET {
+            self.counter.set(self.counter.get() + 1);
+        }
+    }
+}
+
+#[test_log::test]
+fn one_way_test() {
+    let skeleton = super::TestSkeleton::new();
+
+    let span_f1 = tracing::span!(tracing::Level::TRACE, "f1");
+    let span_f2 = tracing::span!(tracing::Level::TRACE, "f2");
+
+    let fvec = skeleton.initialize_nodes(2);
+    let f1 = fvec[0].clone();
+    let f2 = fvec[1].clone();
+
+    f1.with_service(|ms: &fleck_core::MessageService| {
+        ms.add_message_type::<EncryptedMessage>();
+    });
+    f2.with_service(|ms: &fleck_core::MessageService| {
+        ms.add_message_type::<EncryptedMessage>();
+    });
+
+    f2.add_service::<ReceiveService>();
+
+    // let the encryption handshake happen
+    for _ in 0..5 {
+        {
+            let _guard = span_f1.enter();
+            f1.run_once();
+        }
+        {
+            let _guard = span_f2.enter();
+            f2.run_once();
+        }
+    }
+
+    // grab f1's node for f2
+    let node_for_f2 = f1.with_service(|ns: &fleck_core::NodeService|
+        ns.node_by_pubkey(&f2.with_service(|ns: &fleck_core::NodeService| *ns.self_node().pubkey()))).expect("handshake failure?");
+
+    log::info!("--- test triggering a message to be sent ---");
+
+    // trigger f1 to send a packet
+    f1.queue::<fleck_core::SendPacketChannel>(fleck_core::msg::Message::build(EncryptedMessage { secret_payload: SECRET }).with_peer(node_for_f2.peer().unwrap()).with_node(node_for_f2));
+
+    // let the communication commence...
+    for _ in 0..5 {
+        {
+            let _guard = span_f1.enter();
+            f1.run_once();
+        }
+        {
+            let _guard = span_f2.enter();
+            f2.run_once();
+        }
+    }
+
+    assert_eq!(1, f2.with_service(|rs: &ReceiveService| rs.counter.get()));
+}

src/test/handshake.rs

@@ -2,42 +2,14 @@ use crate::fleck_core;
 
 #[test_log::test]
 fn end_to_end() {
-    tracing_log::LogTracer::init().unwrap();
     let skeleton = super::TestSkeleton::new();
 
     let span_f1 = tracing::span!(tracing::Level::TRACE, "f1");
     let span_f2 = tracing::span!(tracing::Level::TRACE, "f2");
 
-    let f1 = crate::API::new();
-    let f2 = crate::API::new();
-
-    let f1pk = f1.with_service(|ns: &fleck_core::NodeService| {
-        ns.build_ephemeral_self_node();
-        ns.self_node().pubkey().clone()
-    });
-    let f2pk = f2.with_service(|ns: &fleck_core::NodeService| {
-        ns.build_ephemeral_self_node();
-        ns.self_node().pubkey().clone()
-    });
-
-    f1.with_service(|io: &fleck_core::io::IOService| {
-        fleck_core::io::UnixSocketBuilder::default()
-            .set_path(skeleton.socket_path().to_str().unwrap())
-            .serve_mode()
-            .build(io);
-    });
-
-    f2.with_service(|pd: &fleck_core::discovery::PeerDiscovery| {
-        pd.new_peer(
-            f2.with_service(|io: &fleck_core::io::IOService| {
-                fleck_core::io::UnixSocketBuilder::default()
-                    .set_path(skeleton.socket_path().to_str().unwrap())
-                    .connect_mode()
-                    .build(io)
-            })
-            .connected_peer(),
-        )
-    });
+    let fvec = skeleton.initialize_nodes(2);
+    let f1 = fvec[0].clone();
+    let f2 = fvec[1].clone();
 
     // run for a few rounds of the event loops, without firing any ticks.
     for _ in 0..5 {
@@ -51,6 +23,9 @@ fn end_to_end() {
         }
     }
 
+    let f1pk = f1.with_service(|ns: &fleck_core::NodeService| *ns.self_node().pubkey());
+    let f2pk = f1.with_service(|ns: &fleck_core::NodeService| *ns.self_node().pubkey());
+
     // check that they know about each other.
     f1.with_service(|ns: &fleck_core::NodeService| {
         assert!(ns.node_by_pubkey(&f2pk).is_some());

src/test/mod.rs

@@ -1,4 +1,15 @@
 mod handshake;
+mod encryption;
+
+static TRACING_SETUP_FLAG: std::sync::Mutex<bool> = std::sync::Mutex::new(false);
+
+fn setup_tracing() {
+    let mut guard = TRACING_SETUP_FLAG.lock().unwrap();
+    if !*guard {
+        tracing_log::LogTracer::init().unwrap();
+        *guard = true;
+    }
+}
 
 struct TestSkeleton {
     tmpfile: tempfile::NamedTempFile,
@@ -6,6 +17,7 @@ struct TestSkeleton {
 
 impl TestSkeleton {
     fn new() -> Self {
+        setup_tracing();
         Self {
             tmpfile: tempfile::NamedTempFile::new().expect("couldn't create temporary file?"),
         }
@@ -14,4 +26,57 @@ impl TestSkeleton {
     pub fn socket_path(&self) -> &std::path::Path {
         self.tmpfile.path()
     }
+
+    /// Builds a single UNIX pipe 'server' node and a given number of 'client' nodes
+    /// Note that the nodes are initialized and the client nodes are informed of the server's
+    /// existence, but they are not stepped forwards at all. If you want the handshake process to
+    /// finish, you will have to step them forwards yourself.
+    pub fn initialize_nodes(&self, count: usize) -> Vec<std::rc::Rc<crate::API>> {
+        assert!(count >= 1);
+
+        let mut result = vec![];
+
+        // initialize server
+        result.push({
+            use crate::fleck_core;
+            let fleck = crate::API::new();
+
+            fleck.with_service(|ns: &fleck_core::NodeService| { ns.build_ephemeral_self_node(); });
+
+            fleck.with_service(|io: &fleck_core::io::IOService| {
+                fleck_core::io::UnixSocketBuilder::default()
+                    .set_path(self.socket_path().to_str().unwrap())
+                    .serve_mode()
+                    .build(io);
+            });
+
+            fleck
+        });
+
+        // initialize clients
+        for _ in 1..count {
+            result.push({
+                use crate::fleck_core;
+                let fleck = crate::API::new();
+
+                fleck.with_service(|ns: &fleck_core::NodeService| { ns.build_ephemeral_self_node(); });
+
+                fleck.with_service(|pd: &fleck_core::discovery::PeerDiscovery| {
+                    pd.new_peer(
+                        fleck.with_service(|io: &fleck_core::io::IOService| {
+                            fleck_core::io::UnixSocketBuilder::default()
+                                .set_path(self.socket_path().to_str().unwrap())
+                                .connect_mode()
+                                .build(io)
+                        })
+                        .connected_peer(),
+                    )
+                });
+
+                fleck
+            });
+        }
+
+        result
+    }
 }