From efd0ce51cbb77b8cb8073e33812915188accf4bb Mon Sep 17 00:00:00 2001
From: Prabhpreet Dua <615318+prabhpreet@users.noreply.github.com>
Date: Sun, 21 Jan 2024 13:53:05 +0530
Subject: [PATCH] On-stack allocated host identification

---
 rosenpass/src/app_server.rs | 30 +++++++++++++++++++++---------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/rosenpass/src/app_server.rs b/rosenpass/src/app_server.rs
index dfb1378..c3d5edd 100644
--- a/rosenpass/src/app_server.rs
+++ b/rosenpass/src/app_server.rs
@@ -199,8 +199,15 @@ impl std::fmt::Display for SocketBoundEndpoint {
 }
 
 impl SocketBoundEndpoint {
-    pub fn to_bytes(&self) -> Vec<u8> {
-        let mut buf = Vec::with_capacity(26);
+
+    const SOCKET_SIZE: usize = usize::BITS as usize/8;
+    const IPV6_SIZE: usize = 16;
+    const PORT_SIZE: usize = 2;
+    const SCOPE_ID_SIZE: usize = 4;
+
+    const BUFFER_SIZE: usize = SocketBoundEndpoint::SOCKET_SIZE + SocketBoundEndpoint::IPV6_SIZE + SocketBoundEndpoint::PORT_SIZE + SocketBoundEndpoint::SCOPE_ID_SIZE;
+    pub fn to_bytes(&self) -> (usize,[u8; SocketBoundEndpoint::BUFFER_SIZE]) {
+        let mut buf = [0u8;SocketBoundEndpoint::BUFFER_SIZE];
         let addr = match self.addr {
             SocketAddr::V4(addr) => {
                 //Map IPv4-mapped to IPv6 addresses
@@ -209,11 +216,16 @@ impl SocketBoundEndpoint {
             }
             SocketAddr::V6(addr) => addr,
         };
-        buf.extend_from_slice(&self.socket.0.to_be_bytes());
-        buf.extend_from_slice(&addr.ip().octets());
-        buf.extend_from_slice(&addr.port().to_be_bytes());
-        buf.extend_from_slice(&addr.scope_id().to_be_bytes());
-        buf
+        let mut len: usize = 0;
+        buf[len..len+SocketBoundEndpoint::SOCKET_SIZE].copy_from_slice(&self.socket.0.to_be_bytes());
+        len += SocketBoundEndpoint::SOCKET_SIZE;
+        buf[len..len+SocketBoundEndpoint::IPV6_SIZE].copy_from_slice(&addr.ip().octets());
+        len += SocketBoundEndpoint::IPV6_SIZE;
+        buf[len..len+SocketBoundEndpoint::PORT_SIZE].copy_from_slice(&addr.port().to_be_bytes());
+        len += SocketBoundEndpoint::PORT_SIZE;
+        buf[len..len+SocketBoundEndpoint::SCOPE_ID_SIZE].copy_from_slice(&addr.scope_id().to_be_bytes());
+        len += SocketBoundEndpoint::SCOPE_ID_SIZE;
+        (len,buf)
     }
 }
 
@@ -640,9 +652,9 @@ impl AppServer {
     ) -> Result<crate::protocol::HandleMsgResult> {
         match endpoint {
             Endpoint::SocketBoundAddress(socket) => {
-                let host_identification = socket.to_bytes();
+                let (hi_len,host_identification )= socket.to_bytes();
                 self.crypt
-                    .handle_msg_under_load(&rx, &mut *tx, &host_identification)
+                    .handle_msg_under_load(&rx, &mut *tx, &host_identification[0..hi_len])
             }
             Endpoint::Discovery(_) => {
                 anyhow::bail!("Host-path discovery is not supported under load")