From 19394279b93f3bb25735b35dd095f0540361b24b Mon Sep 17 00:00:00 2001 From: Vladimir Lebedev Date: Tue, 26 Mar 2024 14:59:11 +0700 Subject: [PATCH] Add possibility to use multiple private keys for server --- examples/echo_server.rs | 2 +- src/helper_types.rs | 20 ++++++++++++++++++++ src/lib.rs | 2 +- src/primitives/codec.rs | 1 + src/primitives/handshake.rs | 7 +++++-- src/tests.rs | 5 +++-- src/wrappers/peer.rs | 19 +++++++++++-------- 7 files changed, 42 insertions(+), 14 deletions(-) diff --git a/examples/echo_server.rs b/examples/echo_server.rs index 55e0b93..ae43df9 100644 --- a/examples/echo_server.rs +++ b/examples/echo_server.rs @@ -46,7 +46,7 @@ async fn main() -> Result<(), Box> { 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 { diff --git a/src/helper_types.rs b/src/helper_types.rs index a1864c3..883723d 100644 --- a/src/helper_types.rs +++ b/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 + } +} \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index 16794e7..e9d19cb 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/primitives/codec.rs b/src/primitives/codec.rs index 0e859a8..617fac8 100644 --- a/src/primitives/codec.rs +++ b/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, diff --git a/src/primitives/handshake.rs b/src/primitives/handshake.rs index f705605..672b367 100644 --- a/src/primitives/handshake.rs +++ b/src/primitives/handshake.rs @@ -99,14 +99,17 @@ impl AdnlHandshake

{ } impl AdnlHandshake { - /// Deserialize and decrypt handshake - pub fn decrypt_from_raw(packet: &[u8; 256], key: &S) -> Result { + /// Deserialize and decrypt handshake using private key from `private_key_selector` function + pub fn decrypt_from_raw Option>(packet: &[u8; 256], private_key_selector: F) -> Result { 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)) } diff --git a/src/tests.rs b/src/tests.rs index 89c10d3..088d63d 100644 --- a/src/tests.rs +++ b/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; } diff --git a/src/wrappers/peer.rs b/src/wrappers/peer.rs index db328cd..6a85f41 100644 --- a/src/wrappers/peer.rs +++ b/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 where T: AsyncRead + AsyncWrite { #[pin] stream: Framed, + connection_info: AdnlConnectionInfo, } impl AdnlPeer { @@ -57,24 +60,24 @@ impl AdnlPeer { // 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(mut transport: T, private_key: &S) -> Result { + /// Act as a server: receive handshake over transport using private key provided by `private_key_selector`. + pub async fn handle_handshake Option>(mut transport: T, private_key_selector: F) -> Result { // 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