Add 1GB/day/user bandwidth cap

common/authenticator-requests/src/v1/request.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use nym_sphinx::addressing::Recipient;
-use nym_wireguard_types::{GatewayClient, InitMessage};
+use nym_wireguard_types::{GatewayClient, InitMessage, PeerPublicKey};
 use serde::{Deserialize, Serialize};
 
 use crate::make_bincode_serializer;
@@ -57,6 +57,19 @@ impl AuthenticatorRequest {
         )
     }
 
+    pub fn new_query_request(peer_public_key: PeerPublicKey, reply_to: Recipient) -> (Self, u64) {
+        let request_id = generate_random();
+        (
+            Self {
+                version: VERSION,
+                data: AuthenticatorRequestData::QueryBandwidth(peer_public_key),
+                reply_to,
+                request_id,
+            },
+            request_id,
+        )
+    }
+
     pub fn to_bytes(&self) -> Result<Vec<u8>, bincode::Error> {
         use bincode::Options;
         make_bincode_serializer().serialize(self)
@@ -67,4 +80,5 @@ impl AuthenticatorRequest {
 pub enum AuthenticatorRequestData {
     Initial(InitMessage),
     Final(GatewayClient),
+    QueryBandwidth(PeerPublicKey),
 }

common/authenticator-requests/src/v1/response.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use nym_sphinx::addressing::Recipient;
-use nym_wireguard_types::registration::RegistrationData;
+use nym_wireguard_types::registration::{RegistrationData, RegistredData, RemainingBandwidthData};
 use serde::{Deserialize, Serialize};
 
 use crate::make_bincode_serializer;
@@ -33,10 +33,31 @@ impl AuthenticatorResponse {
         }
     }
 
-    pub fn new_registered(reply_to: Recipient, request_id: u64) -> Self {
+    pub fn new_registered(
+        registred_data: RegistredData,
+        reply_to: Recipient,
+        request_id: u64,
+    ) -> Self {
         Self {
             version: VERSION,
             data: AuthenticatorResponseData::Registered(RegisteredResponse {
+                reply: registred_data,
+                reply_to,
+                request_id,
+            }),
+            reply_to,
+        }
+    }
+
+    pub fn new_remaining_bandwidth(
+        remaining_bandwidth_data: Option<RemainingBandwidthData>,
+        reply_to: Recipient,
+        request_id: u64,
+    ) -> Self {
+        Self {
+            version: VERSION,
+            data: AuthenticatorResponseData::RemainingBandwidth(RemainingBandwidthResponse {
+                reply: remaining_bandwidth_data,
                 reply_to,
                 request_id,
             }),
@@ -64,6 +85,7 @@ impl AuthenticatorResponse {
         match &self.data {
             AuthenticatorResponseData::PendingRegistration(response) => Some(response.request_id),
             AuthenticatorResponseData::Registered(response) => Some(response.request_id),
+            AuthenticatorResponseData::RemainingBandwidth(response) => Some(response.request_id),
         }
     }
 }
@@ -72,6 +94,7 @@ impl AuthenticatorResponse {
 pub enum AuthenticatorResponseData {
     PendingRegistration(PendingRegistrationResponse),
     Registered(RegisteredResponse),
+    RemainingBandwidth(RemainingBandwidthResponse),
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -85,4 +108,12 @@ pub struct PendingRegistrationResponse {
 pub struct RegisteredResponse {
     pub request_id: u64,
     pub reply_to: Recipient,
+    pub reply: RegistredData,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct RemainingBandwidthResponse {
+    pub request_id: u64,
+    pub reply_to: Recipient,
+    pub reply: Option<RemainingBandwidthData>,
 }

common/wireguard-types/src/registration.rs

@@ -27,12 +27,15 @@ pub type HmacSha256 = Hmac<Sha256>;
 pub type Nonce = u64;
 pub type Taken = Option<SystemTime>;
 
+pub const BANDWIDTH_CAP_PER_DAY: u64 = 1024 * 1024 * 1024; // 1 GB
+
 #[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(tag = "type", rename_all = "camelCase")]
 #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
 pub enum ClientMessage {
     Initial(InitMessage),
     Final(GatewayClient),
+    Query(PeerPublicKey),
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -60,6 +63,19 @@ pub struct RegistrationData {
     pub wg_port: u16,
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct RegistredData {
+    pub pub_key: PeerPublicKey,
+    pub private_ip: IpAddr,
+    pub wg_port: u16,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct RemainingBandwidthData {
+    pub available_bandwidth: u64,
+    pub suspended: bool,
+}
+
 /// Client that wants to register sends its PublicKey bytes mac digest encrypted with a DH shared secret.
 /// Gateway/Nym node can then verify pub_key payload using the same process
 #[derive(Serialize, Deserialize, Debug, Clone)]

common/wireguard/Cargo.toml

@@ -12,6 +12,7 @@ license.workspace = true
 
 [dependencies]
 base64 = { workspace = true }
+chrono = { workspace = true }
 dashmap = { workspace = true }
 defguard_wireguard_rs = { workspace = true }
 # The latest version on crates.io at the time of writing this (6.0.0) has a
@@ -24,5 +25,6 @@ nym-crypto = { path = "../crypto", features = ["asymmetric"] }
 nym-network-defaults = { path = "../network-defaults" }
 nym-task = { path = "../task" }
 nym-wireguard-types = { path = "../wireguard-types" }
+thiserror = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread", "net", "io-util"] }
 tokio-stream = { workspace = true }

common/wireguard/src/error.rs

@@ -0,0 +1,11 @@
+// Copyright 2024 - Nym Technologies SA <[email protected]>
+// SPDX-License-Identifier: Apache-2.0
+
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("peers in wireguard don't match with in-memory ")]
+    PeerMismatch,
+
+    #[error("{0}")]
+    Defguard(#[from] defguard_wireguard_rs::error::WireguardInterfaceError),
+}

common/wireguard/src/lib.rs

@@ -1,3 +1,6 @@
+// Copyright 2024 - Nym Technologies SA <[email protected]>
+// SPDX-License-Identifier: Apache-2.0
+
 #![cfg_attr(not(target_os = "linux"), allow(dead_code))]
 // #![warn(clippy::pedantic)]
 // #![warn(clippy::expect_used)]
@@ -6,13 +9,14 @@
 use dashmap::DashMap;
 use defguard_wireguard_rs::{host::Peer, key::Key, net::IpAddrMask, WGApi};
 use nym_crypto::asymmetric::encryption::KeyPair;
-use nym_wireguard_types::{Config, Error, GatewayClient, GatewayClientRegistry};
-use peer_controller::PeerControlMessage;
+use nym_wireguard_types::{Config, Error, GatewayClient, GatewayClientRegistry, PeerPublicKey};
+use peer_controller::PeerControlRequest;
 use std::sync::Arc;
-use tokio::sync::mpsc::{self, UnboundedReceiver};
+use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
 
 const WG_TUN_NAME: &str = "nymwg";
 
+pub(crate) mod error;
 pub mod peer_controller;
 
 pub struct WgApiWrapper {
@@ -39,14 +43,14 @@ pub struct WireguardGatewayData {
     config: Config,
     keypair: Arc<KeyPair>,
     client_registry: Arc<GatewayClientRegistry>,
-    peer_tx: mpsc::UnboundedSender<PeerControlMessage>,
+    peer_tx: UnboundedSender<PeerControlRequest>,
 }
 
 impl WireguardGatewayData {
     pub fn new(
         config: Config,
         keypair: Arc<KeyPair>,
-    ) -> (Self, mpsc::UnboundedReceiver<PeerControlMessage>) {
+    ) -> (Self, UnboundedReceiver<PeerControlRequest>) {
         let (peer_tx, peer_rx) = mpsc::unbounded_channel();
         (
             WireguardGatewayData {
@@ -75,27 +79,34 @@ impl WireguardGatewayData {
         let mut peer = Peer::new(Key::new(client.pub_key.to_bytes()));
         peer.allowed_ips
             .push(IpAddrMask::new(client.private_ip, 32));
-        let msg = PeerControlMessage::AddPeer(peer);
+        let msg = PeerControlRequest::AddPeer(peer);
         self.peer_tx.send(msg).map_err(|_| Error::PeerModifyStopped)
     }
 
     pub fn remove_peer(&self, client: &GatewayClient) -> Result<(), Error> {
         let key = Key::new(client.pub_key().to_bytes());
-        let msg = PeerControlMessage::RemovePeer(key);
+        let msg = PeerControlRequest::RemovePeer(key);
+        self.peer_tx.send(msg).map_err(|_| Error::PeerModifyStopped)
+    }
+
+    pub fn query_bandwidth(&self, peer_public_key: PeerPublicKey) -> Result<(), Error> {
+        let key = Key::new(peer_public_key.to_bytes());
+        let msg = PeerControlRequest::QueryBandwidth(key);
         self.peer_tx.send(msg).map_err(|_| Error::PeerModifyStopped)
     }
 }
 
 pub struct WireguardData {
     pub inner: WireguardGatewayData,
-    pub peer_rx: UnboundedReceiver<PeerControlMessage>,
+    pub peer_rx: UnboundedReceiver<PeerControlRequest>,
 }
 
 /// Start wireguard device
 #[cfg(target_os = "linux")]
 pub async fn start_wireguard(
     task_client: nym_task::TaskClient,
     wireguard_data: WireguardData,
+    control_tx: UnboundedSender<peer_controller::PeerControlResponse>,
 ) -> Result<std::sync::Arc<WgApiWrapper>, Box<dyn std::error::Error + Send + Sync + 'static>> {
     use base64::{prelude::BASE64_STANDARD, Engine};
     use defguard_wireguard_rs::{InterfaceConfiguration, WireguardInterfaceApi};
@@ -135,13 +146,13 @@ pub async fn start_wireguard(
     wg_api.configure_peer_routing(&[catch_all_peer])?;
 
     let wg_api = std::sync::Arc::new(WgApiWrapper::new(wg_api));
-    let mut controller = PeerController::new(wg_api.clone(), wireguard_data.peer_rx);
+    let mut controller = PeerController::new(
+        wg_api.clone(),
+        interface_config.peers,
+        wireguard_data.peer_rx,
+        control_tx,
+    );
     tokio::spawn(async move { controller.run(task_client).await });
 
     Ok(wg_api)
 }
-
-#[cfg(not(target_os = "linux"))]
-pub async fn start_wireguard() {
-    todo!("WireGuard is currently only supported on Linux");
-}