From 541f667a0e1561ad3fa940b118a78efc341762d3 Mon Sep 17 00:00:00 2001
From: DJMrTV <tvnebel@gmail.com>
Date: Wed, 22 Jan 2025 20:40:18 +0100
Subject: [PATCH] add signature to connection response packet

---
 src/prudp/sockaddr.rs | 17 +++++++++++++++++
 src/prudp/socket.rs   | 27 +++++++--------------------
 2 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/src/prudp/sockaddr.rs b/src/prudp/sockaddr.rs
index 5f87bc8..b48b8c1 100644
--- a/src/prudp/sockaddr.rs
+++ b/src/prudp/sockaddr.rs
@@ -1,8 +1,25 @@
+use std::io::Write;
 use std::net::SocketAddrV4;
+use hmac::{Hmac, Mac};
 use crate::prudp::packet::VirtualPort;
 
+type Md5Hmac = Hmac<md5::Md5>;
+
 #[derive(Eq, PartialEq, Hash, Debug, Copy, Clone)]
 pub struct PRUDPSockAddr{
     pub regular_socket_addr: SocketAddrV4,
     pub virtual_port: VirtualPort
+}
+
+impl PRUDPSockAddr{
+    pub fn calculate_connection_signature(&self) -> [u8; 16] {
+        let mut hmac = Md5Hmac::new_from_slice(&[0; 16]).expect("fuck");
+
+        let mut data = self.regular_socket_addr.ip().octets().to_vec();
+        data.extend_from_slice(&self.regular_socket_addr.port().to_be_bytes());
+
+        hmac.write_all(&data).expect("figuring this out was complete ass");
+        let result: [u8; 16] = hmac.finalize().into_bytes()[0..16].try_into().expect("fuck");
+        result
+    }
 }
\ No newline at end of file
diff --git a/src/prudp/socket.rs b/src/prudp/socket.rs
index 4c2bfe4..23fb119 100644
--- a/src/prudp/socket.rs
+++ b/src/prudp/socket.rs
@@ -19,7 +19,7 @@ use crate::prudp::router::{Error, Router};
 use crate::prudp::sockaddr::PRUDPSockAddr;
 
 
-type Md5Hmac = Hmac<md5::Md5>;
+
 
 /// PRUDP Socket for accepting connections to then send and recieve data from those clients
 pub struct Socket(Arc<SocketImpl>, Arc<Router>, Receiver<Connection>);
@@ -86,6 +86,8 @@ impl Deref for Socket{
     }
 }
 
+
+
 impl SocketImpl {
     fn new(router: &Router, connection_creation_sender: Sender<Connection>, port: VirtualPort, access_key: &'static str) -> Self {
         SocketImpl {
@@ -159,17 +161,11 @@ impl SocketImpl {
                 response_packet.header.types_and_flags.set_flag(flags::ACK);
                 response_packet.header.types_and_flags.set_flag(flags::HAS_SIZE);
 
-                let mut hmac = Md5Hmac::new_from_slice(&[0; 16]).expect("fuck");
 
-                let mut data = connection.regular_socket_addr.ip().octets().to_vec();
-                data.extend_from_slice(&connection.regular_socket_addr.port().to_be_bytes());
 
-                hmac.write_all(&data).expect("figuring this out was complete ass");
-                let result: [u8; 16] = hmac.finalize().into_bytes()[0..16].try_into().expect("fuck");
+                conn.signature = connection.calculate_connection_signature();
 
-                conn.signature = result;
-
-                response_packet.options.push(PacketOption::ConnectionSignature(result));
+                response_packet.options.push(ConnectionSignature(conn.signature));
 
 
 
@@ -207,6 +203,8 @@ impl SocketImpl {
                 response_packet.header.session_id = conn.session_id;
                 response_packet.header.sequence_id = 1;
 
+                response_packet.options.push(ConnectionSignature(conn.signature));
+
                 for option in &packet.options{
                     match option  {
                         MaximumSubstreamId(max_substream) => response_packet.options.push(MaximumSubstreamId(*max_substream)),
@@ -236,14 +234,3 @@ impl SocketImpl {
 
     }
 }
-
-#[cfg(test)]
-mod test {
-    use hmac::Mac;
-    use crate::prudp::socket::Md5Hmac;
-
-    #[test]
-    fn fuck() {
-        let hmac = Md5Hmac::new_from_slice(&[0; 16]).expect("fuck");
-    }
-}
\ No newline at end of file