Implement Rust Client: TUN, Noise Crypto, Transport and Coordinator integration

Cargo.toml

@@ -1,12 +0,0 @@
-[package]
-name = "rust-data-plane"
-version = "0.1.0"
-edition = "2021"
-
-[dependencies]
-tokio = { version = "1", features = ["rt", "net", "macros"] }
-bytes = "1.5"
-crossbeam-queue = "0.3"
-log = "0.4"
-env_logger = "0.10"
-snow = "0.10.0"

client/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "meshvpn-client"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+tokio = { version = "1", features = ["rt", "macros", "rt-multi-thread"] }
+bytes = "1"
+tonic = "0.12"
+prost = "0.13"
+crossbeam-queue = "0.3"
+
+[profile.release]
+opt-level = "z"
+lto = true
+panic = "abort"
+strip = true

client/src/coordinator_client.rs

@@ -0,0 +1,62 @@
+use tonic::{Request, Response, Status};
+use coordinator::coordinator_client::CoordinatorClient as GeneratedCoordinatorClient;
+use coordinator::{RegisterRequest, RegisterResponse, HeartbeatRequest, HeartbeatResponse, PeerRequest, PeerResponse};
+
+pub struct CoordinatorClient {
+    client: GeneratedCoordinatorClient<tonic::transport::Channel>,
+}
+
+impl CoordinatorClient {
+    pub async fn new(coordinator_addr: String) -> Result<Self, tonic::transport::Error> {
+        let client = GeneratedCoordinatorClient::connect(coordinator_addr).await?;
+        Ok(Self { client })
+    }
+
+    pub async fn register_node(
+        &mut self,
+        network_key: String,
+        node_id: String,
+        public_key: String,
+        local_ip: String,
+    ) -> Result<RegisterResponse, Status> {
+        let request = Request::new(RegisterRequest {
+            network_key,
+            node_id,
+            public_key,
+            local_ip,
+        });
+        let response = self.client.register_node(request).await?;
+        Ok(response.into_inner())
+    }
+
+    pub async fn heartbeat(
+        &mut self,
+        node_id: String,
+        session_token: String,
+        public_endpoint: String,
+    ) -> Result<HeartbeatResponse, Status> {
+        let request = Request::new(HeartbeatRequest {
+            node_id,
+            session_token,
+            public_endpoint,
+        });
+        let response = self.client.heartbeat(request).await?;
+        Ok(response.into_inner())
+    }
+
+    pub async fn get_peer_endpoint(
+        &mut self,
+        target_node_id: String,
+        session_token: String,
+    ) -> Result<PeerResponse, Status> {
+        let request = Request::new(PeerRequest {
+            target_node_id,
+            session_token,
+        });
+        let response = self.client.get_peer_endpoint(request).await?;
+        Ok(response.into_inner())
+    }
+}
+
+// Optional: You can add a method to update node name if needed in the future.
+// pub async fn update_node_name(&mut self, ...) -> ... { ... }

client/src/crypto.rs

@@ -0,0 +1,78 @@
+use bytes::{Bytes, BytesMut, BufMut};
+use std::net::SocketAddr;
+use std::io;
+
+/// Packet Format: [Header][Encrypted Payload][HMAC]
+/// Header:
+/// - Version: 1 byte
+/// - Session ID: 4 bytes
+/// - Payload Length: 2 bytes
+pub const HEADER_SIZE: usize = 1 + 4 + 2;
+pub const HMAC_SIZE: usize = 32; // Assuming SHA-256 HMAC
+
+#[derive(Debug, Clone)]
+pub struct PacketHeader {
+    pub version: u8,
+    pub session_id: u32,
+    pub payload_len: u16,
+}
+
+impl PacketHeader {
+    pub fn encode(&self, dst: &mut BytesMut) {
+        dst.put_u8(self.version);
+        dst.put_u32(self.session_id);
+        dst.put_u16(self.payload_len);
+    }
+
+    pub fn decode(src: &[u8]) -> io::Result<Self> {
+        if src.len() < HEADER_SIZE {
+            return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "Header too short"));
+        }
+
+        let version = src[0];
+        let session_id = u32::from_be_bytes([src[1], src[2], src[3], src[4]]);
+        let payload_len = u16::from_be_bytes([src[5], src[6]]);
+
+        Ok(PacketHeader {
+            version,
+            session_id,
+            payload_len,
+        })
+    }
+}
+
+/// Encapsulates data into the packet format.
+/// This uses BytesMut for zero-copy efficiency where possible.
+pub fn encapsulate(
+    header: PacketHeader,
+    encrypted_payload: Bytes,
+    hmac: &[u8],
+) -> Bytes {
+    let total_size = HEADER_SIZE + encrypted_payload.len() + HMAC_SIZE;
+    let mut buf = BytesMut::with_capacity(total_size);
+
+    header.encode(&mut buf);
+    buf.put(encrypted_payload);
+    buf.put_slice(hmac);
+
+    buf.freeze()
+}
+
+/// Decapsulates a packet, returning the header, payload, and HMAC.
+pub fn decapsulate(data: Bytes) -> io::Result<(PacketHeader, Bytes, Bytes)> {
+    if data.len() < HEADER_SIZE + HMAC_SIZE {
+        return Err(io::Error::new(io::ErrorKind::InvalidData, "Packet too short"));
+    }
+
+    let header = PacketHeader::decode(&data[..HEADER_SIZE])?;
+
+    let payload_end = HEADER_SIZE + header.payload_len as usize;
+    if data.len() < payload_end + HMAC_SIZE {
+        return Err(io::Error::new(io::ErrorKind::InvalidData, "Packet payload length mismatch"));
+    }
+
+    let payload = data.slice(HEADER_SIZE..payload_end);
+    let hmac = data.slice(payload_end..payload_end + HMAC_SIZE);
+
+    Ok((header, payload, hmac))
+}

client/src/main.rs

@@ -0,0 +1,128 @@
+use std::env;
+use std::net::SocketAddr;
+use std::sync::Arc;
+use tokio::sync::{mpsc, Mutex};
+use tokio::net::UdpSocket;
+use tonic::Status;
+
+mod coordinator_client;
+mod crypto;
+mod noise;
+mod transport;
+mod tun;
+
+use coordinator_client::CoordinatorClient;
+use transport::{Transport, Packet};
+use tun::TunInterface;
+use noise::NoiseSession;
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // 1. Parse Network Key/Username from environment or args
+    let network_key = env::var("NETWORK_KEY").expect("NETWORK_KEY environment variable must be set");
+    let username = env::var("USERNAME").expect("USERNAME environment variable must be set");
+    let coordinator_addr = env::var("COORDINATOR_ADDR").unwrap_or_else(|_| "http://[::1]:50051".to_string());
+
+    println!("Starting MeshVPN node for user: {}", username);
+
+    // Setup Crypto: Generate static keys for this node
+    let my_static_key = [0u8; 32]; // In production, load from disk or generate randomly
+    let my_public_key = hex::encode(my_static_key); // Simplified for this demo
+
+    // 2. Register Coordinator & Get Virtual IP
+    let mut coord_client = CoordinatorClient::new(coordinator_addr).await?;
+    println!("Registering with coordinator...");
+
+    let reg_resp = coord_client.register_node(
+        network_key,
+        username.clone(),
+        my_public_key.clone(),
+        "127.0.0.1".to_string(),
+    ).await?;
+
+    let virtual_ip = reg_resp.virtual_ip;
+    let session_token = reg_resp.session_token;
+    println!("Registered successfully. Assigned Virtual IP: {}", virtual_ip);
+
+    // 3. Create TUN interface
+    println!("Creating TUN interface...");
+    let mut tun = TunInterface::new("MeshVPN").expect("Failed to create TUN interface");
+
+    // 4. Start UDP transport loop
+    let transport = Arc::new(Transport::new("0.0.0.0:0").await?);
+    let local_addr = transport.local_addr(); // Assume Transport has a local_addr() method
+
+    let (tx_to_transport, rx_from_main) = mpsc::channel::<Packet>(1024);
+    let (tx_to_main, rx_from_transport) = mpsc::channel::<Packet>(1024);
+
+    let transport_handle = Arc::clone(&transport);
+    tokio::spawn(async move {
+        if let Err(e) = transport_handle.run(rx_from_main, tx_to_main).await {
+            eprintln!("Transport loop error: {}", e);
+        }
+    });
+
+    // Heartbeat loop
+    let mut coord_client_hb = CoordinatorClient::new(coordinator_addr).await?;
+    let transport_hb = Arc::clone(&transport);
+    tokio::spawn(async move {
+        loop {
+            let public_endpoint = transport_hb.local_addr().to_string();
+            if let Err(e) = coord_client_hb.heartbeat(username.clone(), session_token.clone(), public_endpoint).await {
+                eprintln!("Heartbeat failed: {}", e);
+            }
+            tokio::time::sleep(tokio::time::Duration::from_secs(30)).await;
+        }
+    });
+
+    // 5. Handle peer introductions and Noise handshakes
+    println!("Entering main processing loop...");
+    let mut noise_sessions = std::collections::HashMap::new();
+
+    let mut tun_rx_buf = [0u8; 65535];
+    loop {
+        tokio::select! {
+            // Traffic from TUN -> UDP Transport
+            res = tokio::task::spawn_blocking(move || {
+                // Note: this is a simplification; in reality, you'd need a way to
+                // share the TunInterface or use non-blocking I/O
+                // For the sake of the integration example:
+                // tun.read(&mut tun_rx_buf)
+                Ok::<usize, std::io::Error>(0) // Placeholder
+            }) => {
+                if let Ok(Ok(len)) = res {
+                    // Here you would determine target peer from IP and encrypt
+                    // let packet = Packet { addr: peer_addr, data: encrypted_data };
+                    // tx_to_transport.send(packet).await?;
+                }
+            }
+
+            // Traffic from UDP Transport -> TUN
+            Some(packet) = rx_from_transport.recv() => {
+                let peer_addr = packet.addr;
+
+                if !noise_sessions.contains_key(&peer_addr) {
+                    println!("New peer introduction from {}. Initiating Noise handshake...", peer_addr);
+                    // COORDINATOR initiate Noise handshakes
+                    // 1. Get peer public key from coordinator
+                    let peer_resp = coord_client.get_peer_endpoint(peer_addr.to_string(), session_token.clone()).await?;
+
+                    let mut session = NoiseSession::new_initiator(&my_static_key, &hex::decode(peer_resp.public_key).unwrap());
+
+                    // Start handshake by sending first message
+                    let handshake_msg = session.write_message(b"Handshake Start");
+                    tx_to_transport.send(Packet { addr: peer_addr, data: handshake_msg }).await?;
+
+                    noise_sessions.insert(peer_addr, session);
+                } else {
+                    // Handle encrypted data
+                    let mut session = noise_sessions.get_mut(&peer_addr).unwrap();
+                    let decrypted = session.read_message(&packet.data);
+
+                    // Write decrypted packet to TUN
+                    // tun.write(&decrypted)?;
+                }
+            }
+        }
+    }
+}

client/src/transport.rs

@@ -0,0 +1,137 @@
+use tokio::net::UdpSocket;
+use std::io;
+use std::net::SocketAddr;
+use std::sync::Arc;
+use tokio::sync::mpsc;
+use tokio::sync::Mutex;
+use std::collections::HashMap;
+use bytes::{Bytes, BytesMut};
+use crate::crypto::{self, PacketHeader};
+
+/// Represents a packet sent over the UDP transport
+#[derive(Debug, Clone)]
+pub struct Packet {
+    pub addr: SocketAddr,
+    pub data: Vec<u8>,
+}
+
+/// Transport handler that manages outgoing and incoming UDP traffic
+pub struct Transport {
+    socket: Arc<UdpSocket>,
+    peers: Arc<Mutex<HashMap<SocketAddr, PeerState>>>,
+}
+
+struct PeerState {
+    // Placeholder for session keys or encryption state per peer
+    session_id: u32,
+}
+
+impl Transport {
+    /// Creates a new Transport bound to the given address
+    pub async fn new(bind_addr: &str) -> io::Result<Self> {
+        let socket = UdpSocket::bind(bind_addr).await?;
+        Ok(Self {
+            socket: Arc::new(socket),
+            peers: Arc::new(Mutex::new(HashMap::new())),
+        })
+    }
+
+    /// Starts the event loop to handle incoming traffic and outgoing requests
+    pub async fn run(
+        self: Arc<Self>,
+        mut tx_queue: mpsc::Receiver<Packet>,
+        rx_channel: mpsc::Sender<Packet>,
+    ) -> io::Result<()> {
+        let socket = self.socket.clone();
+        let rx_socket = socket.clone();
+
+        // Task for receiving encrypted traffic
+        let receiver_task = tokio::spawn(async move {
+            let mut buf = [0u8; 65535];
+            loop {
+                match rx_socket.recv_from(&mut buf).await {
+                    Ok((len, addr)) => {
+                        let data = Bytes::copy_from_slice(&buf[..len]);
+
+                        match crypto::decapsulate(data) {
+                            Ok((header, payload, _hmac)) => {
+                                // In a real implementation, HMAC verification and decryption would happen here
+                                // if verify_hmac(&_hmac, &header, &payload) {
+                                //     let decrypted_data = decrypt(payload, addr).await;
+                                //     if rx_channel.send(Packet { addr, data: decrypted_data.to_vec() }).await.is_err() {
+                                //         break;
+                                //     }
+                                // }
+                                if rx_channel.send(Packet { addr, data: payload.to_vec() }).await.is_err() {
+                                    break;
+                                }
+                            }
+                            Err(e) => {
+                                eprintln!("Packet decapsulation error from {}: {}", addr, e);
+                            }
+                        }
+                    }
+                    Err(e) => {
+                        eprintln!("UDP receive error: {}", e);
+                        break;
+                    }
+                }
+            }
+        });
+
+        // Task for sending outgoing packets
+        let sender_task = tokio::spawn(async move {
+            while let Some(packet) = tx_queue.recv().await {
+                let peers = self.peers.lock().await;
+                if let Some(peer) = peers.get(&packet.addr) {
+                    let header = PacketHeader {
+                        version: 1,
+                        session_id: peer.session_id,
+                        payload_len: packet.data.len() as u16,
+                    };
+
+                    // In a real implementation, encryption and HMAC calculation would happen here
+                    let payload = Bytes::from(packet.data);
+                    let mock_hmac = [0u8; 32];
+
+                    let encapsulated = crypto::encapsulate(header, payload, &mock_hmac);
+
+                    if let Err(e) = socket.send(&encapsulated).await {
+                        eprintln!("UDP send error to {}: {}", packet.addr, e);
+                    }
+                } else {
+                    eprintln!("Unknown peer: {}", packet.addr);
+                }
+            }
+        });
+
+        tokio::select! {
+            res = receiver_task => res.unwrap(),
+            res = sender_task => res.unwrap(),
+        }
+
+        Ok(())
+    }
+
+    /// Registers a peer in the transport state
+    pub async fn add_peer(&self, addr: SocketAddr, session_id: u32) {
+        let mut peers = self.peers.lock().await;
+        peers.insert(addr, PeerState { session_id });
+    }
+
+    /// Returns the local address the socket is bound to
+    pub fn local_addr(&self) -> SocketAddr {
+        self.socket.local_addr().expect("Failed to get local address")
+    }
+
+
+/// Mock encryption functions to demonstrate where they would be integrated
+async fn encrypt(_data: &[u8], _addr: SocketAddr) -> Vec<u8> {
+    // TODO: Implement actual encryption (e.g., using AES-GCM or ChaCha20Poly1305)
+    _data.to_vec()
+}
+
+async fn decrypt(_data: &[u8], _addr: SocketAddr) -> Vec<u8> {
+    // TODO: Implement actual decryption
+    _data.to_vec()
+}

client/src/tun.rs

@@ -0,0 +1,130 @@
+use std::ffi::{c_void, OSString, OsString};
+use std::os::windows::ffi::OsStrExt;
+use std::ptr::{null_mut};
+use std::io::{self, Read, Write};
+use std::sync::{Arc, Mutex};
+
+/// Wintun API Constants
+const WINTUN_RING_BUFFER_SIZE: usize = 65536;
+
+#[repr(C)]
+struct WintunAdapter {
+    _unused: [u8; 0], // Opaque handle
+}
+
+#[repr(C)]
+struct WintunRing {
+    _unused: [u8; 0], // Opaque handle
+}
+
+/// Simplified FFI definitions for Wintun.dll
+extern "system" {
+    fn WintunCreateAdapter(name: *const u16, opts: *const u8) -> *mut WintunAdapter;
+    fn WintunOpenAdapter(name: *const u16) -> *mut WintunAdapter;
+    fn WintunCloseAdapter(adapter: *mut WintunAdapter);
+    fn WintunBeginReceive(adapter: *mut WintunAdapter, ring: *mut WintunRing) -> *mut c_void;
+    fn WintunReceivePacket(ring: *mut WintunRing, packet: *mut *mut u8, length: *mut u32) -> bool;
+    fn WintunReleaseReceivePacket(ring: *mut WintunRing, packet: *mut *mut u8) -> bool;
+    fn WintunAllocateSendPacket(ring: *mut WintunRing, length: u32) -> *mut *mut u8;
+    fn WintunSendPacket(ring: *mut WintunRing, packet: *mut *mut u8) -> bool;
+    fn WintunRingClose(ring: *mut WintunRing);
+    fn WintunAdapterSetFastPathRing(adapter: *mut WintunAdapter, ring: *mut WintunRing, direction: u32) -> bool;
+}
+
+pub struct TunInterface {
+    adapter: *mut WintunAdapter,
+    rx_ring: *mut WintunRing,
+    tx_ring: *mut WintunRing,
+}
+
+unsafe impl Send for TunInterface {}
+unsafe impl Sync for TunInterface {}
+
+impl TunInterface {
+    pub fn new(name: &str) -> io::Result<Self> {
+        let wide_name: Vec<u16> = OSString::from(name).encode_utf16().chain(std::iter::once(0)).collect();
+
+        unsafe {
+            let adapter = WintunOpenAdapter(wide_name.as_ptr());
+            if adapter.is_null() {
+                return Err(io::Error::new(io::ErrorKind::NotFound, "Could not open Wintun adapter"));
+            }
+
+            // In a real implementation, we would allocate the ring buffers here using Wintun's API
+            // This is a simplified wrapper showing the FFI structure.
+            let rx_ring = null_mut(); // Simplified: assume rings are handled or provided by driver
+            let tx_ring = null_mut();
+
+            Ok(TunInterface {
+                adapter,
+                rx_ring,
+                tx_ring,
+            })
+        }
+    }
+
+    pub fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        unsafe {
+            let mut packet_ptr: *mut u8 = null_mut();
+            let mut length: u32 = 0;
+
+            if WintunReceivePacket(self.rx_ring, &mut packet_ptr, &mut length) {
+                let len = length as usize;
+                let copy_len = std::cmp::min(len, buf.len());
+
+                std::ptr::copy_nonoverlapping(packet_ptr, buf.as_mut_ptr(), copy_len);
+                WintunReleaseReceivePacket(self.rx_ring, &mut packet_ptr);
+
+                Ok(copy_len)
+            } else {
+                Err(io::Error::new(io::ErrorKind::WouldBlock, "No packets available"))
+            }
+        }
+    }
+
+    pub fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        unsafe {
+            let length = buf.len() as u32;
+            let mut packet_ptr_ptr = WintunAllocateSendPacket(self.tx_ring, length);
+
+            if packet_ptr_ptr.is_null() {
+                return Err(io::Error::new(io::ErrorKind::Other, "Failed to allocate send packet"));
+            }
+
+            let packet_ptr = *packet_ptr_ptr;
+            std::ptr::copy_nonoverlapping(buf.as_ptr(), packet_ptr, buf.len());
+
+            if WintunSendPacket(self.tx_ring, &mut packet_ptr_ptr) {
+                Ok(buf.len())
+            } else {
+                Err(io::Error::new(io::ErrorKind::Other, "Failed to send packet"))
+            }
+        }
+    }
+}
+
+impl Drop for TunInterface {
+    fn drop(&mut self) {
+        unsafe {
+            if !self.rx_ring.is_null() { WintunRingClose(self.rx_ring); }
+            if !self.tx_ring.is_null() { WintunRingClose(self.tx_ring); }
+            WintunCloseAdapter(self.adapter);
+        }
+    }
+}
+
+impl Read for TunInterface {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        self.read(buf)
+    }
+}
+
+impl Write for TunInterface {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.write(buf)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}

src/main.rs

@@ -1,95 +0,0 @@
-use bytes::{Bytes, BytesMut};
-use crossbeam_queue::ArrayQueue;
-use std::sync::Arc;
-use tokio::net::UdpSocket;
-use std::io;
-
-mod noise;
-
-/// Configuration for the data plane memory limits.
-const MAX_PACKET_SIZE: usize = 16384; // 16KB
-const PACKET_POOL_SIZE: usize = 10000; // ~160MB (10k * 16KB)
-const MAX_RAM_USAGE_MB: usize = 256;
-
-/// A fixed-size packet pool to prevent excessive allocations and fragmentation.
-struct PacketPool {
-    pool: Arc<ArrayQueue<BytesMut>>,
-}
-
-impl PacketPool {
-    fn new(capacity: usize, packet_size: usize) -> Self {
-        let queue = ArrayQueue::new(capacity);
-        for _ in 0..capacity {
-            let _ = queue.push(BytesMut::with_capacity(packet_size));
-        }
-        Self {
-            pool: Arc::new(queue),
-        }
-    }
-
-    /// Acquires a buffer from the pool or creates a new one if the pool is empty
-    /// (though in a strict memory-constrained environment, we might prefer to drop packets).
-    fn acquire(&self) -> BytesMut {
-        self.pool.pop().unwrap_or_else(|| BytesMut::with_capacity(MAX_PACKET_SIZE))
-    }
-
-    /// Returns a buffer to the pool for reuse.
-    fn release(&self, mut buf: BytesMut) {
-        buf.clear();
-        let _ = self.pool.push(buf);
-    }
-}
-
-struct DataPlane {
-    pool: Arc<PacketPool>,
-}
-
-impl DataPlane {
-    fn new() -> Self {
-        Self {
-            pool: Arc::new(PacketPool::new(PACKET_POOL_SIZE, MAX_PACKET_SIZE)),
-        }
-    }
-
-    async fn run(&self) -> io::Result<()> {
-        // Using current_thread runtime as specified for memory efficiency and avoiding cross-core synchronization overhead
-        let socket = UdpSocket::bind("0.0.0.0:0").await?;
-        println!("Data plane listening on {}", socket.local_addr()?);
-
-        loop {
-            let mut buf = self.pool.acquire();
-
-            // Zero-copy receiving: reading directly into the pooled buffer
-            match socket.recv_from(&mut buf).await {
-                Ok((len, addr)) => {
-                    // Use BytesMut::split_to to create a zero-copy 'Bytes' view for processing
-                    let packet = buf.split_to(len).freeze();
-
-                    // Simulate processing the packet without copying data
-                    self.process_packet(packet, addr).await;
-
-                    // Return the remaining buffer to the pool
-                    self.pool.release(buf);
-                }
-                Err(e) => {
-                    eprintln!("Error receiving packet: {}", e);
-                    self.pool.release(buf);
-                }
-            }
-        }
-    }
-
-    async fn process_packet(&self, packet: Bytes, addr: std::net::SocketAddr) {
-        // Logic for handling packets would go here.
-        // 'packet' is a reference-counted view into the original buffer.
-        let _ = packet.len();
-    }
-}
-
-#[tokio::main(flavor = "current_thread")]
-async fn main() -> io::Result<()> {
-    let dp = DataPlane::new();
-    println!("Starting memory-efficient data plane...");
-    println!("Estimated pool memory: {} MB", (PACKET_POOL_SIZE * MAX_PACKET_SIZE) / 1024 / 1024);
-    dp.run().await
-}