Add possibility to use multiple private keys for server

examples/echo_server.rs

@@ -46,7 +46,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
         let private_key = private_key.clone();
         tokio::spawn(async move {
             // ADNL: handle handshake
-            let mut adnl_server = AdnlPeer::handle_handshake(socket, &private_key).await.expect("handshake failed");
+            let mut adnl_server = AdnlPeer::handle_handshake(socket, |_| Some(private_key.clone())).await.expect("handshake failed");
 
             // In a loop, read data from the socket and write the data back.
             while let Some(Ok(packet)) = adnl_server.next().await {

src/helper_types.rs

@@ -203,3 +203,23 @@ pub enum AdnlError {
     #[error("End of stream")]
     EndOfStream,
 }
+
+/// Information about connected peers.
+pub struct AdnlConnectionInfo {
+    local_address: AdnlAddress,
+    remote_address: AdnlAddress,
+}
+
+impl AdnlConnectionInfo {
+    pub fn new(local_address: AdnlAddress, remote_address: AdnlAddress) -> Self {
+        Self { local_address, remote_address }
+    }
+
+    pub fn local_address(&self) -> &AdnlAddress {
+        &self.local_address
+    }
+
+    pub fn remote_address(&self) -> &AdnlAddress {
+        &self.remote_address
+    }
+}

src/lib.rs

@@ -1,5 +1,5 @@
 pub use helper_types::{
-    AdnlAddress, AdnlAesParams, AdnlError, AdnlPrivateKey, AdnlPublicKey, AdnlSecret, AdnlRawPublicKey,
+    AdnlAddress, AdnlAesParams, AdnlError, AdnlPrivateKey, AdnlPublicKey, AdnlSecret, AdnlRawPublicKey, AdnlConnectionInfo
 };
 pub use primitives::handshake::AdnlHandshake;
 pub use primitives::codec::AdnlCodec;

src/primitives/codec.rs

@@ -6,6 +6,7 @@ use crate::{AdnlAesParams, AdnlError};
 
 use super::AdnlAes;
 
+/// Implementation of ADNL protocol. Connection must be first initialized with [`AdnlHandshake`] to exchange keys.
 pub struct AdnlCodec {
     aes_rx: AdnlAes,
     aes_tx: AdnlAes,

src/primitives/handshake.rs

@@ -99,14 +99,17 @@ impl<P: AdnlPublicKey> AdnlHandshake<P> {
 }
 
 impl AdnlHandshake<AdnlRawPublicKey> {
-    /// Deserialize and decrypt handshake
-    pub fn decrypt_from_raw<S: AdnlPrivateKey>(packet: &[u8; 256], key: &S) -> Result<Self, AdnlError> {
+    /// Deserialize and decrypt handshake using private key from `private_key_selector` function
+    pub fn decrypt_from_raw<S: AdnlPrivateKey, F: Fn(&AdnlAddress) -> Option<S>>(packet: &[u8; 256], private_key_selector: F) -> Result<Self, AdnlError> {
         let receiver = packet[..32].try_into().unwrap();
         let sender = packet[32..64].try_into().unwrap();
         let hash: [u8; 32] = packet[64..96].try_into().unwrap();
         let mut raw_params: [u8; 160] = packet[96..256].try_into().unwrap();
 
+        let key = private_key_selector(&receiver).ok_or_else(|| AdnlError::UnknownAddr(receiver.clone()))?;
+
         if key.public().address() != receiver {
+            log::error!("private key selector returned wrong key, expected address: {:?}, got: {:?}", &receiver, &key.public().address());
             return Err(AdnlError::UnknownAddr(receiver))
         }
 

src/tests.rs

@@ -67,6 +67,7 @@ fn test_handshake(
     );
 
     // test deserializing
+    #[derive(Clone)]
     struct DummyKey {
         ecdh: AdnlSecret,
         public: AdnlRawPublicKey
@@ -85,7 +86,7 @@ fn test_handshake(
     }
 
     let key = DummyKey { ecdh: ecdh, public: remote_public.clone() };
-    let handshake2 = AdnlHandshake::decrypt_from_raw(expected_handshake.as_slice().try_into().unwrap(), &key).expect("invalid handshake");
+    let handshake2 = AdnlHandshake::decrypt_from_raw(expected_handshake.as_slice().try_into().unwrap(), |_| Some(key.clone())).expect("invalid handshake");
     assert_eq!(handshake2.aes_params().to_bytes(), aes_params_raw, "aes_params mismatch");
     assert_eq!(handshake2.receiver(), &remote_public.address(), "receiver mismatch");
     assert_eq!(handshake2.sender().edwards_repr(), local_public.edwards_repr(), "sender mismatch");
@@ -183,7 +184,7 @@ async fn integrity_test() {
             let (socket, _) = listener.accept().await.unwrap();
             let private_key = server_private.clone();
             tokio::spawn(async move {
-                let mut adnl_server = AdnlPeer::handle_handshake(socket, &private_key).await.expect("handshake failed");
+                let mut adnl_server = AdnlPeer::handle_handshake(socket, |_| Some(private_key.clone())).await.expect("handshake failed");
                 while let Some(Ok(packet)) = adnl_server.next().await {
                     let _ = adnl_server.send(packet).await;
                 }

src/wrappers/peer.rs

@@ -1,7 +1,9 @@
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
-use crate::{AdnlBuilder, AdnlError, AdnlHandshake, AdnlPrivateKey, AdnlPublicKey};
+use crate::helper_types::AdnlConnectionInfo;
+use crate::primitives::handshake;
+use crate::{AdnlAddress, AdnlBuilder, AdnlError, AdnlHandshake, AdnlPrivateKey, AdnlPublicKey};
 use pin_project::pin_project;
 use tokio::io::{AsyncRead, AsyncWrite, AsyncReadExt, AsyncWriteExt};
 use tokio::net::{TcpStream, ToSocketAddrs};
@@ -17,6 +19,7 @@ use crate::primitives::codec::AdnlCodec;
 pub struct AdnlPeer<T> where T: AsyncRead + AsyncWrite {
     #[pin]
     stream: Framed<T, AdnlCodec>,
+    connection_info: AdnlConnectionInfo,
 }
 
 impl AdnlPeer<TcpStream> {
@@ -57,24 +60,24 @@ impl<T: AsyncReadExt + AsyncWriteExt + Unpin> AdnlPeer<T> {
         // receive empty message to ensure that server knows our AES keys
         if let Some(x) = stream.next().await {
             x?;
-            Ok(Self { stream })
+            let connection_info = AdnlConnectionInfo::new(handshake.sender().address(), handshake.receiver().clone());
+            Ok(Self { stream, connection_info })
         } else {
             Err(AdnlError::EndOfStream)
         }
     }
 
-    /// Act as a server: receive handshake over transport. 
-    /// Verifies following things:
-    /// 1) target ADNL address matches associated with provided private key
-    /// 2) integrity of handshake is not compromised
-    pub async fn handle_handshake<S: AdnlPrivateKey>(mut transport: T, private_key: &S) -> Result<Self, AdnlError> {
+    /// Act as a server: receive handshake over transport using private key provided by `private_key_selector`.
+    pub async fn handle_handshake<S: AdnlPrivateKey, F: Fn(&AdnlAddress) -> Option<S>>(mut transport: T, private_key_selector: F) -> Result<Self, AdnlError> {
         // receive handshake
         let mut packet = [0u8; 256];
         transport.read_exact(&mut packet).await.map_err(AdnlError::IoError)?;
-        let handshake = AdnlHandshake::decrypt_from_raw(&packet, private_key)?;
+        let handshake = AdnlHandshake::decrypt_from_raw(&packet, private_key_selector)?;
+        let connection_info = AdnlConnectionInfo::new(handshake.receiver().clone(), handshake.sender().address());
 
         let mut server = Self {
             stream: handshake.make_server_codec().framed(transport),
+            connection_info,
         };
 
         // send empty packet to proof knowledge of AES keys