Create encrypted UDP transport with streaming and firewall hole-punching

crates/core/Cargo.toml

@@ -14,10 +14,12 @@ path = "src/bin/freenet.rs"
 
 [dependencies]
 anyhow = "1"
+arc-swap = "1.6"
 asynchronous-codec = "0.7"
+aes-gcm = "0.10.3"
 async-trait = "0.1"
 axum = { default-features = false, features = ["http1", "matched-path", "query", "tower-log", "ws", "json"], workspace = true }
-bincode = "1"
+bincode = "1.3.3"
 blake3 = { workspace = true }
 bs58 = "0.5"
 byteorder = "1"
@@ -54,6 +56,7 @@ serde_with = { workspace = true }
 sqlx = { features = ["runtime-tokio-rustls", "sqlite"], optional = true, version = "0.7" }
 stretto = { features = ["async", "sync"], version = "0.8" }
 tar = { version = "0.4.38" }
+time = "0.3.30"
 thiserror = "1"
 tokio = { features = ["fs", "macros", "rt-multi-thread", "sync", "process"], version = "1" }
 tokio-tungstenite = "0.21"
@@ -63,7 +66,7 @@ unsigned-varint = { version = "0.8", features = ["codec", "asynchronous_codec"]
 wasmer = { features = ["sys"], workspace = true }
 xz2 = { version = "0.1" }
 reqwest = { version = "0.11.23", features = ["json"] }
-# enum-iterator = "1.4.1"
+rsa = { version = "0.9.6", features = ["serde"] }
 
 # Tracing deps
 opentelemetry = "0.21.0"
@@ -74,7 +77,6 @@ tracing-subscriber = { optional = true, version = "0.3.16" }
 
 # internal deps
 freenet-stdlib = { features = ["net"], workspace = true }
-time = "0.3.30"
 
 [dev-dependencies]
 arbitrary = { features = ["derive"], version = "1" }

crates/core/src/lib.rs

@@ -23,6 +23,8 @@ pub mod server;
 mod topology;
 /// Tracing and loging infrastructure. Includes our custom event log register. Tracing collectors, etc.
 mod tracing;
+/// Code for communicating with other peers over UDP, handles hole-punching, error handling, etc.
+mod transport;
 pub mod util;
 /// WASM code execution runtime, tailored for the contract and delegate APIs.
 mod wasm_runtime;

crates/core/src/topology.rs

@@ -546,7 +546,7 @@ mod tests {
         // Simulate a bunch of random requests clustered around 0.35
         for _ in 0..NUM_REQUESTS {
             let requested_location = random_location(&random_location(&this_peer_location));
-            // FIXME: Is PeerKeyLocation unimportant for this test?
+            // todo: Is PeerKeyLocation unimportant for this test?
             topology_manager.record_request(
                 PeerKeyLocation::random(),
                 requested_location,

crates/core/src/transport.rs

@@ -0,0 +1,82 @@
+#![allow(dead_code)] // TODO: Remove before integration
+//! Freenet Transport protocol implementation.
+//!
+//! Please see `docs/architecture/transport.md` for more information.
+
+mod bw;
+mod connection_handler;
+mod crypto;
+mod packet_data;
+mod peer_connection;
+mod received_packet_tracker;
+mod sent_packet_tracker;
+mod symmetric_message;
+
+type MessagePayload = Vec<u8>;
+type MessageId = u32;
+
+use self::packet_data::PacketData;
+use std::time::Duration;
+
+/// We can wait up to 100ms to confirm a message was received, this allows us to batch
+/// receipts together and send them in a single message.
+const MAX_CONFIRMATION_DELAY: Duration = Duration::from_millis(100);
+
+struct BytesPerSecond(f64);
+
+impl BytesPerSecond {
+    pub fn new(bytes_per_second: f64) -> Self {
+        assert!(bytes_per_second >= 0.0);
+        Self(bytes_per_second)
+    }
+
+    pub fn as_f64(&self) -> f64 {
+        self.0
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::transport::received_packet_tracker::ReportResult;
+    use crate::transport::sent_packet_tracker::{ResendAction, MESSAGE_CONFIRMATION_TIMEOUT};
+
+    #[test]
+    fn test_packet_send_receive_acknowledge_flow() {
+        let mut sent_tracker = sent_packet_tracker::tests::mock_sent_packet_tracker();
+        let mut received_tracker = received_packet_tracker::tests::mock_received_packet_tracker();
+
+        // Simulate sending packets
+        for id in 1..=5 {
+            sent_tracker.report_sent_packet(id, vec![id as u8]);
+        }
+
+        // Simulate receiving some packets
+        for id in [1, 3, 5] {
+            assert_eq!(
+                received_tracker.report_received_packet(id),
+                ReportResult::Ok
+            );
+        }
+
+        // Get receipts and simulate acknowledging them
+        let receipts = received_tracker.get_receipts();
+        assert_eq!(receipts, vec![1, 3, 5]);
+        sent_tracker.report_received_receipts(&receipts);
+
+        // Check resend action for lost packets
+        sent_tracker
+            .time_source
+            .advance_time(MESSAGE_CONFIRMATION_TIMEOUT);
+        for id in [2, 4] {
+            match sent_tracker.get_resend() {
+                ResendAction::Resend(message_id, packet) => {
+                    assert_eq!(message_id, id);
+                    // Simulate resending packet
+                    sent_tracker.report_sent_packet(id, packet);
+                }
+                _ => panic!("Expected resend action for packet {}", id),
+            }
+        }
+    }
+}

crates/core/src/transport/bw.rs

@@ -0,0 +1,156 @@
+use crate::util::{CachingSystemTimeSrc, TimeSource};
+use std::collections::VecDeque;
+use std::time::{Duration, Instant};
+
+/// Keeps track of the bandwidth used in the last window_size. Recommend a `window_size` of
+/// 10 seconds.
+pub(super) struct PacketBWTracker<T: TimeSource> {
+    packets: VecDeque<(usize, Instant)>,
+    window_size: Duration,
+    current_bandwidth: usize,
+    time_source: T,
+}
+
+impl PacketBWTracker<CachingSystemTimeSrc> {
+    pub(super) fn new(window_size: Duration) -> Self {
+        PacketBWTracker {
+            packets: VecDeque::new(),
+            window_size,
+            current_bandwidth: 0,
+            time_source: CachingSystemTimeSrc::new(),
+        }
+    }
+}
+
+impl<T: TimeSource> PacketBWTracker<T> {
+    /// Report that a packet was sent
+    pub(super) fn add_packet(&mut self, packet_size: usize) {
+        let now = self.time_source.now();
+        self.packets.push_back((packet_size, now));
+        self.current_bandwidth += packet_size;
+        self.cleanup();
+    }
+
+    /// Removes packets that are older than the window size.
+    fn cleanup(&mut self) {
+        let now = self.time_source.now();
+        while self
+            .packets
+            .front()
+            .map_or(false, |&(_, time)| now - time > self.window_size)
+        {
+            let expired = self.packets.pop_front();
+            if let Some((size, _)) = expired {
+                self.current_bandwidth -= size;
+            }
+        }
+    }
+
+    /// Returns none if the packet can be sent immediately without `bandwidth_limit` being
+    /// exceeded within the `window_size`. Otherwise returns Some(wait_time) where wait_time is the
+    /// amount of time that should be waited before sending the packet.
+    ///
+    /// `bandwidth_limit` should be set to 50% higher than the target upstream bandwidth the
+    /// [topology manager](crate::topology::TopologyManager) is aiming for, as it serves
+    /// as a hard limit which we'd prefer not to hit.
+    pub(super) fn can_send_packet(
+        &mut self,
+        bandwidth_limit: usize,
+        packet_size: usize,
+    ) -> Option<Duration> {
+        self.cleanup();
+
+        if self.current_bandwidth + packet_size <= bandwidth_limit {
+            return None;
+        }
+
+        let mut temp_bandwidth = self.current_bandwidth;
+        let mut wait_time = None;
+
+        for &(size, time) in self.packets.iter() {
+            temp_bandwidth -= size;
+            if temp_bandwidth + packet_size <= bandwidth_limit {
+                wait_time = Some(self.window_size - (self.time_source.now() - time));
+                break;
+            }
+        }
+
+        wait_time
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::util::MockTimeSource;
+
+    fn mock_tracker(window_size: Duration) -> PacketBWTracker<MockTimeSource> {
+        PacketBWTracker {
+            packets: VecDeque::new(),
+            window_size,
+            current_bandwidth: 0,
+            time_source: MockTimeSource::new(Instant::now()),
+        }
+    }
+
+    fn verify_bandwidth_match<T: TimeSource>(tracker: &PacketBWTracker<T>) {
+        let mut total_bandwidth = 0;
+        for &(size, _) in tracker.packets.iter() {
+            total_bandwidth += size;
+        }
+        assert_eq!(total_bandwidth, tracker.current_bandwidth);
+    }
+
+    #[test]
+    fn test_adding_packets() {
+        let mut tracker = PacketBWTracker::new(Duration::from_secs(1));
+        verify_bandwidth_match(&tracker);
+        tracker.add_packet(1500);
+        verify_bandwidth_match(&tracker);
+        assert_eq!(tracker.packets.len(), 1);
+    }
+
+    #[test]
+    fn test_bandwidth_calculation() {
+        let mut tracker = PacketBWTracker::new(Duration::from_secs(1));
+        tracker.add_packet(1500);
+        tracker.add_packet(2500);
+        verify_bandwidth_match(&tracker);
+        assert_eq!(
+            tracker.packets.iter().map(|&(size, _)| size).sum::<usize>(),
+            4000
+        );
+    }
+
+    #[test]
+    fn test_packet_expiry() {
+        let mut tracker = mock_tracker(Duration::from_millis(200));
+        tracker.add_packet(1500);
+        verify_bandwidth_match(&tracker);
+        tracker.time_source.advance_time(Duration::from_millis(300));
+        tracker.cleanup();
+        verify_bandwidth_match(&tracker);
+        assert!(tracker.packets.is_empty());
+    }
+
+    #[test]
+    fn test_wait_time_calculation() {
+        let mut tracker = mock_tracker(Duration::from_secs(1));
+        tracker.add_packet(5000);
+        verify_bandwidth_match(&tracker);
+        tracker.time_source.advance_time(Duration::from_millis(500));
+        tracker.add_packet(4000);
+        verify_bandwidth_match(&tracker);
+        match tracker.can_send_packet(10000, 2000) {
+            None => panic!("Should require waiting"),
+            Some(wait_time) => assert_eq!(wait_time, Duration::from_millis(500)),
+        }
+    }
+
+    #[test]
+    fn test_immediate_send() {
+        let mut tracker = PacketBWTracker::new(Duration::from_millis(10));
+        tracker.add_packet(3000);
+        assert_eq!(tracker.can_send_packet(10000, 2000), None);
+    }
+}