feat: Add KeyExchange and Encrypted tcpstream

From https://github.com/rustdesk/rustdesk-server/pull/400/commits/7e6ddf06a988d31beaf80dc6ea12ad70068b38cb

------

libs/hbb_common/src/tcp.rs

@@ -3,7 +3,10 @@ use anyhow::Context as AnyhowCtx;
 use bytes::{BufMut, Bytes, BytesMut};
 use futures::{SinkExt, StreamExt};
 use protobuf::Message;
-use sodiumoxide::crypto::secretbox::{self, Key, Nonce};
+use sodiumoxide::crypto::{
+    box_,
+    secretbox::{self, Key, Nonce},
+};
 use std::{
     io::{self, Error, ErrorKind},
     net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
@@ -21,10 +24,13 @@ use tokio_util::codec::Framed;
 pub trait TcpStreamTrait: AsyncRead + AsyncWrite + Unpin {}
 pub struct DynTcpStream(Box<dyn TcpStreamTrait + Send + Sync>);
 
+#[derive(Clone)]
+pub struct Encrypt(Key, u64, u64);
+
 pub struct FramedStream(
     Framed<DynTcpStream, BytesCodec>,
     SocketAddr,
-    Option<(Key, u64, u64)>,
+    Option<Encrypt>,
     u64,
 );
 
@@ -130,7 +136,7 @@ impl FramedStream {
                     ms_timeout,
                     Socks5Stream::connect_with_socket(stream, target),
                 )
-                .await??
+                    .await??
             } else {
                 super::timeout(
                     ms_timeout,
@@ -138,7 +144,7 @@ impl FramedStream {
                         stream, target, username, password,
                     ),
                 )
-                .await??
+                    .await??
             };
             let addr = stream.local_addr()?;
             return Ok(Self(
@@ -185,10 +191,10 @@ impl FramedStream {
     #[inline]
     pub async fn send_raw(&mut self, msg: Vec<u8>) -> ResultType<()> {
         let mut msg = msg;
+        log::debug!("Send message: {:?}", msg);
         if let Some(key) = self.2.as_mut() {
-            key.1 += 1;
-            let nonce = Self::get_nonce(key.1);
-            msg = secretbox::seal(&msg, &nonce, &key.0);
+            log::debug!("Encrypting message");
+            msg = key.enc(&msg);
         }
         self.send_bytes(bytes::Bytes::from(msg)).await?;
         Ok(())
@@ -207,18 +213,10 @@ impl FramedStream {
     #[inline]
     pub async fn next(&mut self) -> Option<Result<BytesMut, Error>> {
         let mut res = self.0.next().await;
-        if let Some(key) = self.2.as_mut() {
-            if let Some(Ok(bytes)) = res.as_mut() {
-                key.2 += 1;
-                let nonce = Self::get_nonce(key.2);
-                match secretbox::open(bytes, &nonce, &key.0) {
-                    Ok(res) => {
-                        bytes.clear();
-                        bytes.put_slice(&res);
-                    }
-                    Err(()) => {
-                        return Some(Err(Error::new(ErrorKind::Other, "decryption error")));
-                    }
+        if let Some(Ok(bytes)) = res.as_mut() {
+            if let Some(key) = self.2.as_mut() {
+                if let Err(err) = key.dec(bytes) {
+                    return Some(Err(err));
                 }
             }
         }
@@ -235,7 +233,7 @@ impl FramedStream {
     }
 
     pub fn set_key(&mut self, key: Key) {
-        self.2 = Some((key, 0, 0));
+        self.2 = Some(Encrypt::new(key));
     }
 
     fn get_nonce(seqnum: u64) -> Nonce {
@@ -274,6 +272,8 @@ pub async fn listen_any(port: u16, reuse: bool) -> ResultType<TcpListener> {
         }
         #[cfg(unix)]
         {
+            socket.set_reuseport(true).ok();
+            socket.set_reuseaddr(true).ok();
             use std::os::unix::io::{FromRawFd, IntoRawFd};
             let raw_fd = socket.into_raw_fd();
             let sock2 = unsafe { socket2::Socket::from_raw_fd(raw_fd) };
@@ -297,9 +297,11 @@ pub async fn listen_any(port: u16, reuse: bool) -> ResultType<TcpListener> {
             }
         }
     }
-    let s = TcpSocket::new_v4()?;
-    s.bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), port))?;
-    Ok(s.listen(DEFAULT_BACKLOG)?)
+    Ok(new_socket(
+        SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), port),
+        true,
+    )?
+        .listen(DEFAULT_BACKLOG)?)
 }
 
 impl Unpin for DynTcpStream {}
@@ -333,3 +335,53 @@ impl AsyncWrite for DynTcpStream {
 }
 
 impl<R: AsyncRead + AsyncWrite + Unpin> TcpStreamTrait for R {}
+
+impl Encrypt {
+    pub fn new(key: Key) -> Self {
+        Self(key, 0, 0)
+    }
+
+    pub fn dec(&mut self, bytes: &mut bytes::BytesMut) -> Result<(), Error> {
+        if bytes.len() <= 1 {
+            return Ok(());
+        }
+        self.2 += 1;
+        let nonce = FramedStream::get_nonce(self.2);
+        match secretbox::open(bytes, &nonce, &self.0) {
+            Ok(res) => {
+                bytes.clear();
+                bytes.put_slice(&res);
+                Ok(())
+            }
+            Err(()) => Err(Error::new(ErrorKind::Other, "decryption error")),
+        }
+    }
+
+    pub fn enc(&mut self, data: &[u8]) -> Vec<u8> {
+        self.1 += 1;
+        let nonce = FramedStream::get_nonce(self.1);
+        secretbox::seal(&data, &nonce, &self.0)
+    }
+
+    pub fn decode(
+        symmetric_data: &[u8],
+        their_pk_b: &[u8],
+        our_sk_b: &box_::SecretKey,
+    ) -> ResultType<Key> {
+        if their_pk_b.len() != box_::PUBLICKEYBYTES {
+            anyhow::bail!("Handshake failed: pk length {}", their_pk_b.len());
+        }
+        let nonce = box_::Nonce([0u8; box_::NONCEBYTES]);
+        let mut pk_ = [0u8; box_::PUBLICKEYBYTES];
+        pk_[..].copy_from_slice(their_pk_b);
+        let their_pk_b = box_::PublicKey(pk_);
+        let symmetric_key = box_::open(symmetric_data, &nonce, &their_pk_b, &our_sk_b)
+            .map_err(|_| anyhow::anyhow!("Handshake failed: box decryption failure"))?;
+        if symmetric_key.len() != secretbox::KEYBYTES {
+            anyhow::bail!("Handshake failed: invalid secret key length from peer");
+        }
+        let mut key = [0u8; secretbox::KEYBYTES];
+        key[..].copy_from_slice(&symmetric_key);
+        Ok(Key(key))
+    }
+}