Add public key validation to encryption layer options

Signed-off-by: trivernis <trivernis@protonmail.com>

Cargo.lock

@@ -120,7 +120,7 @@ dependencies = [
 
 [[package]]
 name = "bromine"
-version = "0.20.1"
+version = "0.21.0"
 dependencies = [
  "async-trait",
  "bincode",

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "bromine"
-version = "0.20.1"
+version = "0.21.0"
 authors = ["trivernis <trivernis@protonmail.com>"]
 edition = "2018"
 readme = "README.md"

src/error.rs

@@ -1,6 +1,8 @@
 use crate::error_event::ErrorEventData;
 use thiserror::Error;
 use tokio::sync::oneshot;
+#[cfg(feature = "encryption_layer")]
+use x25519_dalek::PublicKey;
 
 pub type Result<T> = std::result::Result<T, Error>;
 
@@ -39,6 +41,10 @@ pub enum Error {
 
     #[error("Invalid state")]
     InvalidState,
+
+    #[cfg(feature = "encryption_layer")]
+    #[error("Connection of unknown peer with key {0:?} refused")]
+    UnknownPeer(PublicKey),
 }
 
 impl Error {

src/lib.rs

@@ -119,6 +119,7 @@ pub mod ipc;
 mod macros;
 mod namespaces;
 pub mod protocol;
+pub mod utils;
 
 /// Reexported for usage in payload implementations
 pub use bytes;

src/protocol/encrypted/crypt_handling.rs

@@ -1,5 +1,5 @@
-use crate::prelude::encrypted::EncryptedStream;
-use crate::prelude::IPCResult;
+use crate::prelude::encrypted::{EncryptedStream, Keys};
+use crate::prelude::{IPCError, IPCResult};
 use crate::protocol::AsyncProtocolStream;
 use bytes::Bytes;
 use chacha20poly1305::aead::{Aead, NewAead};
@@ -105,15 +105,21 @@ impl<T: AsyncProtocolStream> EncryptedStream<T> {
     /// 4. The server generates a new secret
     /// 5. The server sends the secret to the client
     /// 6. The connection is upgraded with the new shared key
-    pub async fn from_server_key_exchange(mut inner: T, secret: StaticSecret) -> IPCResult<Self> {
+    pub async fn from_server_key_exchange(mut inner: T, keys: &Keys) -> IPCResult<Self> {
         let other_pub = receive_public_key(&mut inner).await?;
-        send_public_key(&mut inner, &secret).await?;
-        let shared_secret = secret.diffie_hellman(&other_pub);
+        tracing::debug!("received peer public key {:?}", other_pub);
+
+        if !keys.allow_unknown && !keys.known_peers.contains(&other_pub) {
+            return Err(IPCError::UnknownPeer(other_pub));
+        }
+        send_public_key(&mut inner, &keys.secret).await?;
+        let shared_secret = keys.secret.diffie_hellman(&other_pub);
         let mut stream = Self::new(inner, shared_secret);
         let permanent_secret = generate_secret();
         stream.write_all(&permanent_secret).await?;
         stream.flush().await?;
         stream.update_key(permanent_secret.into());
+        tracing::debug!("Connection established");
 
         Ok(stream)
     }
@@ -124,14 +130,20 @@ impl<T: AsyncProtocolStream> EncryptedStream<T> {
     /// 3. The client creates an intermediary encrypted connection
     /// 4. The client receives the new key from the server
     /// 5. The connection is upgraded with the new shared key
-    pub async fn from_client_key_exchange(mut inner: T, secret: StaticSecret) -> IPCResult<Self> {
-        send_public_key(&mut inner, &secret).await?;
+    pub async fn from_client_key_exchange(mut inner: T, keys: &Keys) -> IPCResult<Self> {
+        send_public_key(&mut inner, &keys.secret).await?;
         let other_pub = receive_public_key(&mut inner).await?;
-        let shared_secret = secret.diffie_hellman(&other_pub);
+        tracing::debug!("received peer public key {:?}", other_pub);
+
+        if !keys.allow_unknown && !keys.known_peers.contains(&other_pub) {
+            return Err(IPCError::UnknownPeer(other_pub));
+        }
+        let shared_secret = keys.secret.diffie_hellman(&other_pub);
         let mut stream = Self::new(inner, shared_secret);
         let mut key_buf = vec![0u8; 32];
         stream.read_exact(&mut key_buf).await?;
         stream.update_key(key_buf.into());
+        tracing::debug!("Connection established");
 
         Ok(stream)
     }

src/protocol/encrypted/mod.rs

@@ -10,7 +10,7 @@ use std::future::Future;
 use std::io;
 use std::pin::Pin;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite};
-use x25519_dalek::{SharedSecret, StaticSecret};
+use x25519_dalek::{PublicKey, SharedSecret, StaticSecret};
 
 use crate::prelude::encrypted::crypt_handling::CipherBox;
 use crate::prelude::{AsyncProtocolStream, AsyncStreamProtocolListener};
@@ -20,7 +20,14 @@ pub type OptionalFuture<T> = Option<Pin<Box<dyn Future<Output = T> + Send + Sync
 #[derive(Clone)]
 pub struct EncryptionOptions<T: Clone + Default> {
     pub inner_options: T,
+    pub keys: Keys,
+}
+
+#[derive(Clone)]
+pub struct Keys {
     pub secret: StaticSecret,
+    pub known_peers: Vec<PublicKey>,
+    pub allow_unknown: bool,
 }
 
 impl<T: Clone + Default> Default for EncryptionOptions<T> {
@@ -30,7 +37,11 @@ impl<T: Clone + Default> Default for EncryptionOptions<T> {
         rng.fill_bytes(&mut secret);
 
         Self {
-            secret: StaticSecret::from(secret),
+            keys: Keys {
+                known_peers: Vec::new(),
+                allow_unknown: false,
+                secret: StaticSecret::from(secret),
+            },
             inner_options: T::default(),
         }
     }
@@ -38,12 +49,12 @@ impl<T: Clone + Default> Default for EncryptionOptions<T> {
 
 pub struct EncryptedListener<T: AsyncStreamProtocolListener> {
     inner: T,
-    secret: StaticSecret,
+    keys: Keys,
 }
 
 impl<T: AsyncStreamProtocolListener> EncryptedListener<T> {
-    pub fn new(inner: T, secret: StaticSecret) -> Self {
-        Self { inner, secret }
+    pub fn new(inner: T, keys: Keys) -> Self {
+        Self { inner, keys }
     }
 }
 

src/protocol/encrypted/protocol_impl.rs

@@ -18,13 +18,12 @@ impl<T: AsyncStreamProtocolListener> AsyncStreamProtocolListener for EncryptedLi
     ) -> IPCResult<Self> {
         let inner = T::protocol_bind(address, options.inner_options).await?;
 
-        Ok(EncryptedListener::new(inner, options.secret))
+        Ok(EncryptedListener::new(inner, options.keys))
     }
 
     async fn protocol_accept(&self) -> IPCResult<(Self::Stream, Self::RemoteAddressType)> {
         let (inner_stream, remote_addr) = self.inner.protocol_accept().await?;
-        let stream =
-            Self::Stream::from_server_key_exchange(inner_stream, self.secret.clone()).await?;
+        let stream = Self::Stream::from_server_key_exchange(inner_stream, &self.keys).await?;
 
         Ok((stream, remote_addr))
     }
@@ -40,7 +39,7 @@ impl<T: AsyncProtocolStream> AsyncProtocolStream for EncryptedStream<T> {
         options: Self::StreamOptions,
     ) -> Result<Self> {
         let inner = T::protocol_connect(address, options.inner_options).await?;
-        EncryptedStream::from_client_key_exchange(inner, options.secret).await
+        EncryptedStream::from_client_key_exchange(inner, &options.keys).await
     }
 }
 

src/utils.rs

@@ -0,0 +1,11 @@
+#[cfg(feature = "encryption_layer")]
+/// Generates a secret that can be passed to the options of the encryption layer and for creating
+/// a public key
+pub fn generate_secret() -> x25519_dalek::StaticSecret {
+    let mut rng = rand::thread_rng();
+    use rand_core::RngCore;
+    let mut secret = [0u8; 32];
+    rng.fill_bytes(&mut secret);
+
+    x25519_dalek::StaticSecret::from(secret)
+}

tests/test_encryption.rs

@@ -1,19 +1,38 @@
 #![cfg(feature = "encryption_layer")]
+
 use crate::utils::call_counter::increment_counter_for_event;
 use crate::utils::protocol::TestProtocolListener;
 use crate::utils::{get_free_port, start_server_and_client};
-use bromine::prelude::encrypted::EncryptedListener;
+use bromine::prelude::encrypted::{EncryptedListener, EncryptionOptions, Keys};
 use bromine::prelude::*;
+use bromine::utils::generate_secret;
 use bromine::IPCBuilder;
 use byteorder::{BigEndian, ReadBytesExt};
 use bytes::{BufMut, Bytes, BytesMut};
+use dashmap::DashMap;
 use futures::StreamExt;
+use lazy_static::lazy_static;
 use rand_core::RngCore;
 use std::io::Read;
 use std::time::Duration;
+use x25519_dalek::{PublicKey, StaticSecret};
 
 mod utils;
 
+pub fn get_secret<S: AsRef<str>>(name: S) -> StaticSecret {
+    lazy_static! {
+        static ref KEYS: DashMap<String, StaticSecret> = DashMap::new();
+    }
+    if KEYS.contains_key(name.as_ref()) {
+        KEYS.get(name.as_ref()).as_ref().unwrap().value().clone()
+    } else {
+        let secret = generate_secret();
+        KEYS.insert(name.as_ref().to_string(), secret.clone());
+
+        secret
+    }
+}
+
 #[tokio::test]
 async fn it_sends_and_receives_smaller_packages() {
     send_and_receive_bytes(140).await.unwrap();
@@ -66,7 +85,27 @@ async fn get_client_with_server() -> Context {
 }
 
 fn get_builder(port: u8) -> IPCBuilder<EncryptedListener<TestProtocolListener>> {
+    let server_secret = get_secret(format!("server-{}", port));
+    let client_secret = get_secret(format!("client-{}", port));
+    let client_keys = Keys {
+        secret: client_secret.clone(),
+        known_peers: vec![PublicKey::from(&server_secret)],
+        allow_unknown: false,
+    };
+    let server_keys = Keys {
+        secret: server_secret.clone(),
+        known_peers: vec![PublicKey::from(&client_secret)],
+        allow_unknown: false,
+    };
     IPCBuilder::new()
+        .client_options(EncryptionOptions {
+            keys: client_keys,
+            inner_options: (),
+        })
+        .server_options(EncryptionOptions {
+            keys: server_keys,
+            inner_options: (),
+        })
         .address(port)
         .on("bytes", callback!(handle_bytes))
         .on("string", callback!(handle_string))