Track auth and service connections client side

Signed-off-by: Luca Fulchir <luca.fulchir@runesauth.com>

src/auth/mod.rs

@@ -103,4 +103,8 @@ impl ServiceID {
     pub const fn len() -> usize {
         16
     }
+    /// read the service id as bytes
+    pub fn as_bytes(&self) -> &[u8; 16] {
+        &self.0
+    }
 }

src/connection/handshake/dirsync.rs

@@ -425,7 +425,7 @@ pub struct RespData {
     /// Server Connection ID
     pub id: ID,
     /// Service Connection ID
-    pub service_id: ID,
+    pub service_connection_id: ID,
     /// Service encryption key
     pub service_key: Secret,
 }
@@ -448,7 +448,7 @@ impl RespData {
         self.id.serialize(&mut out[start..end]);
         start = end;
         end = end + Self::NONCE_LEN;
-        self.service_id.serialize(&mut out[start..end]);
+        self.service_connection_id.serialize(&mut out[start..end]);
         start = end;
         end = end + Self::NONCE_LEN;
         out[start..end].copy_from_slice(self.service_key.as_ref());

src/connection/handshake/mod.rs

@@ -36,6 +36,8 @@ pub(crate) struct HandshakeServer {
 pub(crate) struct HandshakeClient {
     pub id: crate::enc::asym::KeyID,
     pub key: crate::enc::asym::PrivKey,
+    pub service_id: crate::auth::ServiceID,
+    pub service_conn_id: connection::IDRecv,
     pub connection: Rc<crate::connection::Connection>,
 }
 

src/connection/mod.rs

@@ -20,10 +20,10 @@ use crate::{
 };
 
 /// strong typedef for receiving connection id
-#[derive(Debug, Copy, Clone)]
+#[derive(Debug, Copy, Clone, PartialEq)]
 pub struct IDRecv(pub ID);
 /// strong typedef for sending connection id
-#[derive(Debug, Copy, Clone)]
+#[derive(Debug, Copy, Clone, PartialEq)]
 pub struct IDSend(pub ID);
 
 /// Version of the fenrir protocol in use
@@ -86,8 +86,8 @@ impl Connection {
                 (hkdf.get_secret(b"to_client"), hkdf.get_secret(b"to_server"))
             }
         };
-        let mut cipher_recv = CipherRecv::new(cipher, secret_recv);
-        let mut cipher_send = CipherSend::new(cipher, secret_send, rand);
+        let cipher_recv = CipherRecv::new(cipher, secret_recv);
+        let cipher_send = CipherSend::new(cipher, secret_send, rand);
 
         Self {
             id_recv: IDRecv(ID::Handshake),
@@ -111,13 +111,17 @@ pub(crate) struct ConnList {
 impl ConnList {
     pub(crate) fn new(thread_id: ThreadTracker) -> Self {
         let mut bitmap_id = ::bitmaps::Bitmap::<1024>::new();
-        bitmap_id.set(0, true); // ID(0) == handshake
-        Self {
+        const INITIAL_CAP: usize = 128;
+        let mut ret = Self {
             thread_id,
-            connections: Vec::with_capacity(128),
+            connections: Vec::with_capacity(INITIAL_CAP),
             ids_used: vec![bitmap_id],
-        }
+        };
+        ret.connections.resize_with(INITIAL_CAP, || None);
+        ret
     }
+    /// Only *Reserve* a connection,
+    /// without actually tracking it in self.connections
     pub(crate) fn reserve_first(
         &mut self,
         mut conn: Connection,
@@ -128,13 +132,13 @@ impl ConnList {
         //   initialized
         // * `ID::new_u64` is really safe only with >0, but here it always is
         // ...we should probably rewrite it in better, safer rust
-        let mut id_in_thread: u64 = 0;
+        let mut id_in_thread: usize = 0;
         let mut found = false;
         for (i, b) in self.ids_used.iter_mut().enumerate() {
             match b.first_false_index() {
                 Some(idx) => {
                     b.set(idx, true);
-                    id_in_thread = ((i as u64) * 1024) + (idx as u64);
+                    id_in_thread = (i * 1024) + idx;
                     found = true;
                     break;
                 }
@@ -144,20 +148,48 @@ impl ConnList {
         if !found {
             let mut new_bitmap = ::bitmaps::Bitmap::<1024>::new();
             new_bitmap.set(0, true);
-            id_in_thread = (self.ids_used.len() as u64) * 1024;
+            id_in_thread = self.ids_used.len() * 1024;
             self.ids_used.push(new_bitmap);
         }
-        let actual_id = (id_in_thread * (self.thread_id.total as u64))
+        // make sure we have enough space in self.connections
+        let curr_capacity = self.connections.capacity();
+        if self.connections.capacity() <= id_in_thread {
+            // Fill with "None", assure 64 connections without reallocations
+            let multiple = 64 + curr_capacity - 1;
+            let new_capacity = multiple - (multiple % curr_capacity);
+            self.connections.resize_with(new_capacity, || None);
+        }
+        // calculate the actual connection ID
+        let actual_id = ((id_in_thread as u64) * (self.thread_id.total as u64))
             + (self.thread_id.id as u64);
         let new_id = IDRecv(ID::new_u64(actual_id));
         conn.id_recv = new_id;
-        let conn = Rc::new(conn);
-        if (self.connections.len() as u64) < id_in_thread {
-            self.connections.push(Some(conn.clone()));
-        } else {
-            // very probably redundant
-            self.connections[id_in_thread as usize] = Some(conn.clone());
+        // Return the new connection without tracking it
+        Rc::new(conn)
+    }
+    /// NOTE: does NOT check if the connection has been previously reserved!
+    pub(crate) fn track(&mut self, conn: Rc<Connection>) -> Result<(), ()> {
+        let conn_id = match conn.id_recv {
+            IDRecv(ID::Handshake) => {
+                return Err(());
+            }
+            IDRecv(ID::ID(conn_id)) => conn_id,
+        };
+        let id_in_thread: usize =
+            (conn_id.get() / (self.thread_id.total as u64)) as usize;
+        self.connections[id_in_thread] = Some(conn);
+        Ok(())
+    }
+    pub(crate) fn delete(&mut self, id: IDRecv) {
+        if let IDRecv(ID::ID(raw_id)) = id {
+            let id_in_thread: usize =
+                (raw_id.get() / (self.thread_id.total as u64)) as usize;
+            let vec_index = id_in_thread / 1024;
+            let bitmask_index = id_in_thread % 1024;
+            if let Some(bitmask) = self.ids_used.get_mut(vec_index) {
+                bitmask.set(bitmask_index, false);
+                self.connections[id_in_thread] = None;
+            }
         }
-        conn
     }
 }

src/enc/sym.rs

@@ -115,6 +115,11 @@ impl Cipher {
             }
         }
     }
+    pub fn kind(&self) -> CipherKind {
+        match self {
+            Cipher::XChaCha20Poly1305(_) => CipherKind::XChaCha20Poly1305,
+        }
+    }
     fn nonce_len(&self) -> HeadLen {
         match self {
             Cipher::XChaCha20Poly1305(_) => {
@@ -181,7 +186,7 @@ impl Cipher {
         aad: AAD,
         data: &mut [u8],
     ) -> Result<(), Error> {
-        // FIXME: check minimum buffer size 
+        // FIXME: check minimum buffer size
         match self {
             Cipher::XChaCha20Poly1305(cipher) => {
                 use ::chacha20poly1305::{
@@ -242,6 +247,10 @@ impl CipherRecv {
     ) -> Result<&'a [u8], Error> {
         self.0.decrypt(aad, data)
     }
+    /// return the underlying cipher id
+    pub fn kind(&self) -> CipherKind {
+        self.0.kind()
+    }
 }
 
 /// Allocate some data, with additional indexes to track
@@ -313,6 +322,10 @@ impl CipherSend {
         self.cipher.encrypt(&old_nonce, aad, data)?;
         Ok(())
     }
+    /// return the underlying cipher id
+    pub fn kind(&self) -> CipherKind {
+        self.cipher.kind()
+    }
 }
 
 /// XChaCha20Poly1305 cipher

src/inner/mod.rs

@@ -35,6 +35,10 @@ pub(crate) struct AuthNeededInfo {
 /// Client information needed to fully establish the conenction
 #[derive(Debug)]
 pub(crate) struct ClientConnectInfo {
+    /// The service ID that we are connecting to
+    pub service_id: auth::ServiceID,
+    /// The service ID that we are connecting to
+    pub service_connection_id: connection::IDRecv,
     /// Parsed handshake packet
     pub handshake: Handshake,
     /// Connection
@@ -90,7 +94,7 @@ impl HandshakeTracker {
         }
     }
     pub(crate) fn recv_handshake(
-        &self,
+        &mut self,
         mut handshake: Handshake,
         handshake_raw: &mut [u8],
     ) -> Result<HandshakeAction, Error> {
@@ -175,24 +179,28 @@ impl HandshakeTracker {
                     }));
                 }
                 DirSync::Resp(resp) => {
-                    let hshake = {
+                    let hshake_idx = {
                         match self
                             .hshake_cli
                             .iter()
-                            .find(|h| h.id == resp.client_key_id)
+                            .position(|h| h.id == resp.client_key_id)
                         {
                             Some(h) => Some(h.clone()),
                             None => None,
                         }
                     };
-                    if hshake.is_none() {
-                        ::tracing::debug!(
-                            "No such client key id: {:?}",
-                            resp.client_key_id
-                        );
-                        return Err(handshake::Error::UnknownKeyID.into());
-                    }
-                    let hshake = hshake.unwrap();
+                    let hshake_idx = {
+                        if let Some(real_idx) = hshake_idx {
+                            real_idx
+                        } else {
+                            ::tracing::debug!(
+                                "No such client key id: {:?}",
+                                resp.client_key_id
+                            );
+                            return Err(handshake::Error::UnknownKeyID.into());
+                        }
+                    };
+                    let hshake = &self.hshake_cli[hshake_idx];
                     let cipher_recv = &hshake.connection.cipher_recv;
                     use crate::enc::sym::AAD;
                     // no aad for now
@@ -208,8 +216,18 @@ impl HandshakeTracker {
                             return Err(handshake::Error::Key(e).into());
                         }
                     }
+                    // we can remove the handshake from the list
+                    let hshake: HandshakeClient = {
+                        let len = self.hshake_cli.len();
+                        if (hshake_idx + 1) != len {
+                            self.hshake_cli.swap(hshake_idx, len - 1);
+                        }
+                        self.hshake_cli.pop().unwrap()
+                    };
                     return Ok(HandshakeAction::ClientConnect(
                         ClientConnectInfo {
+                            service_id: hshake.service_id,
+                            service_connection_id: hshake.service_conn_id,
                             handshake,
                             connection: hshake.connection,
                         },

src/inner/worker.rs

@@ -11,7 +11,7 @@ use crate::{
         socket::{UdpClient, UdpServer},
         ConnList, Connection, IDSend, Packet, ID,
     },
-    enc::sym::Secret,
+    enc::{hkdf::HkdfSha3, sym::Secret},
     inner::{HandshakeAction, HandshakeTracker, ThreadTracker},
 };
 use ::std::{rc::Rc, sync::Arc, vec::Vec};
@@ -238,7 +238,7 @@ impl Worker {
                     let resp_data = dirsync::RespData {
                         client_nonce: req_data.nonce,
                         id: auth_conn.id_recv.0,
-                        service_id: srv_conn_id,
+                        service_connection_id: srv_conn_id,
                         service_key: srv_secret,
                     };
                     use crate::enc::sym::AAD;
@@ -296,10 +296,40 @@ impl Worker {
                         );
                         return;
                     }
-                    // FIXME: conn tracking and arc counting
-                    let conn = Rc::get_mut(&mut cci.connection).unwrap();
-                    conn.id_send = IDSend(resp_data.id);
-                    todo!();
+                    {
+                        let conn = Rc::get_mut(&mut cci.connection).unwrap();
+                        conn.id_send = IDSend(resp_data.id);
+                    }
+                    // track the connection to the authentication server
+                    if self.connections.track(cci.connection.clone()).is_err() {
+                        self.connections.delete(cci.connection.id_recv);
+                    }
+                    if cci.connection.id_recv.0
+                        == resp_data.service_connection_id
+                    {
+                        // the user asked a single connection
+                        // to the authentication server, without any additional
+                        // service. No more connections to setup
+                        return;
+                    }
+                    // create and track the connection to the service
+                    //FIXME: the Secret should be XORed with the client stored
+                    // secret (if any)
+                    let hkdf = HkdfSha3::new(
+                        cci.service_id.as_bytes(),
+                        resp_data.service_key,
+                    );
+                    let mut service_connection = Connection::new(
+                        hkdf,
+                        cci.connection.cipher_recv.kind(),
+                        connection::Role::Client,
+                        &self.rand,
+                    );
+                    service_connection.id_recv = cci.service_connection_id;
+                    service_connection.id_send =
+                        IDSend(resp_data.service_connection_id);
+                    self.connections.track(service_connection.into());
+                    return;
                 }
                 _ => {}
             };
@@ -326,6 +356,6 @@ impl Worker {
                 return;
             }
         };
-        src_sock.send_to(&data, client.0);
+        let _ = src_sock.send_to(&data, client.0).await;
     }
 }