udp: start work on HMAC connection ID generation and validation

aquatic_udp/src/common.rs

@@ -1,19 +1,94 @@
 use std::collections::BTreeMap;
 use std::hash::Hash;
-use std::net::{Ipv4Addr, Ipv6Addr};
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
+use std::time::{Duration, Instant};
 
-use aquatic_common::CanonicalSocketAddr;
 use crossbeam_channel::{Sender, TrySendError};
+use getrandom::getrandom;
 
 use aquatic_common::access_list::AccessListArcSwap;
+use aquatic_common::CanonicalSocketAddr;
 use aquatic_udp_protocol::*;
 
 use crate::config::Config;
 
 pub const MAX_PACKET_SIZE: usize = 8192;
 
+pub struct ConnectionIdHandler {
+    start_time: Instant,
+    max_connection_age: Duration,
+    hmac: blake3::Hasher,
+}
+
+impl ConnectionIdHandler {
+    pub fn new(config: &Config) -> anyhow::Result<Self> {
+        let mut key = [0; 32];
+
+        getrandom(&mut key)?;
+
+        let hmac = blake3::Hasher::new_keyed(&key);
+
+        let start_time = Instant::now();
+        let max_connection_age = Duration::from_secs(config.cleaning.max_connection_age);
+
+        Ok(Self {
+            hmac,
+            start_time,
+            max_connection_age,
+        })
+    }
+
+    pub fn create_connection_id(&mut self, source_ip: IpAddr) -> ConnectionId {
+        // Seconds elapsed since server start, as bytes
+        let elapsed_time_bytes = (self.start_time.elapsed().as_secs() as u32).to_ne_bytes();
+
+        self.create_connection_id_inner(elapsed_time_bytes, source_ip)
+    }
+
+    pub fn connection_id_valid(&mut self, source_ip: IpAddr, connection_id: ConnectionId) -> bool {
+        let elapsed_time_bytes = connection_id.0.to_ne_bytes()[..4].try_into().unwrap();
+
+        // i64 comparison should be constant-time
+        let hmac_valid =
+            connection_id == self.create_connection_id_inner(elapsed_time_bytes, source_ip);
+
+        if !hmac_valid {
+            return false;
+        }
+
+        let connection_elapsed_since_start =
+            Duration::from_secs(u32::from_ne_bytes(elapsed_time_bytes) as u64);
+
+        connection_elapsed_since_start + self.max_connection_age > self.start_time.elapsed()
+    }
+
+    fn create_connection_id_inner(
+        &mut self,
+        elapsed_time_bytes: [u8; 4],
+        source_ip: IpAddr,
+    ) -> ConnectionId {
+        // The first 4 bytes is the elapsed time since server start in seconds. The last 4 is a
+        // truncated message authentication code.
+        let mut connection_id_bytes = [0u8; 8];
+
+        (&mut connection_id_bytes[..4]).copy_from_slice(&elapsed_time_bytes);
+
+        self.hmac.update(&elapsed_time_bytes);
+
+        match source_ip {
+            IpAddr::V4(ip) => self.hmac.update(&ip.octets()),
+            IpAddr::V6(ip) => self.hmac.update(&ip.octets()),
+        };
+
+        self.hmac.finalize_xof().fill(&mut connection_id_bytes[4..]);
+        self.hmac.reset();
+
+        ConnectionId(i64::from_ne_bytes(connection_id_bytes))
+    }
+}
+
 #[derive(Debug)]
 pub struct PendingScrapeRequest {
     pub slab_key: usize,