wip: improve bandwidth limiting

Author: cilki

Cargo.lock

@@ -10,8 +10,8 @@ version = "0.0.8"
 
 [dependencies]
 anyhow = "1.0.86"
-askama = { version = "0.12.1", features = ["with-axum"] }
-askama_axum = "0.4.0"
+askama = { version = "0.12.1", features = ["with-axum"], optional = true }
+askama_axum = { version = "0.4.0", optional = true }
 atty = "0.2"
 aws-config = { version = "1.5.3", optional = true, features = [
   "behavior-version-latest",
@@ -20,17 +20,19 @@ aws-sdk-ec2 = { version = "1.55.0", optional = true }
 aws-sdk-cloudformation = { version = "1.0", optional = true }
 aws-sdk-route53 = { version = "1.0", optional = true }
 aws-sdk-ssm = { version = "1.0", optional = true }
-axum = "0.7.5"
+axum = { version = "0.7.5", optional = true }
 base64 = "0.22"
-chrono = "0.4"
+chrono = { version = "0.4", optional = true }
 clap = { version = "4.5.4", features = ["string", "derive", "env"] }
-rust-embed = { version = "8.4.0", features = ["axum", "debug-embed"] }
+rust-embed = { version = "8.4.0", features = ["axum", "debug-embed"], optional = true }
 futures = "0.3.30"
-indicatif = "0.17"
-mime_guess = "2.0"
-nix = { version = "0.29", features = ["net"] }
-rand = "0.8"
-reqwest = { version = "0.12", default-features = false, features = ["rustls-tls"] }
+indicatif = "0.18.3"
+mime_guess = { version = "2.0", optional = true }
+nix = { version = "0.30.1", features = ["net"] }
+rand = "0.9.2"
+reqwest = { version = "0.12", default-features = false, features = [
+  "rustls-tls",
+] }
 serde = { version = "1.0.203", features = ["derive"] }
 serde_json = "1.0.117"
 serde_yaml = "0.9.27"
@@ -39,12 +41,17 @@ tokio = { version = "1.38.0", features = ["full"] }
 tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 url = "2.5"
+validator = { version = "0.18", features = ["derive"] }
 
 [build-dependencies]
 anyhow = "1.0.86"
-reqwest = { version = "0.12.4", default-features = false, features = ["blocking", "rustls-tls"] }
+reqwest = { version = "0.12.4", default-features = false, features = [
+  "blocking",
+  "rustls-tls",
+] }
 
 [features]
+default = ["aws", "dashboard"]
 cloudflare = []
 aws = [
   "dep:aws-config",
@@ -53,3 +60,27 @@ aws = [
   "dep:aws-sdk-route53",
   "dep:aws-sdk-ssm",
 ]
+dashboard = [
+  "dep:axum",
+  "dep:askama",
+  "dep:askama_axum",
+  "dep:rust-embed",
+  "dep:mime_guess",
+  "dep:chrono",
+]
+
+[[example]]
+name = "ui_aws"
+required-features = ["dashboard"]
+
+[[example]]
+name = "ui_cloudflare"
+required-features = ["dashboard"]
+
+[[example]]
+name = "ui_no_proxy"
+required-features = ["dashboard"]
+
+[[example]]
+name = "ui_high_traffic"
+required-features = ["dashboard"]

Cargo.toml

@@ -1,4 +1,4 @@
-FROM rust:slim as build
+FROM rust:slim AS build
 RUN rustup target add x86_64-unknown-linux-musl
 
 RUN apt-get update && apt-get install -y musl-tools

Dockerfile

@@ -30,6 +30,8 @@ async fn index() -> Html<String> {
         tunnel_stats: stats,
         proxy_info,
         cloudfront_info: None,
+        upload_limit: Some(100),
+        download_limit: Some(100),
     };
 
     Html(template.render().unwrap())

examples/ui_aws.rs

@@ -24,6 +24,9 @@ async fn index() -> Html<String> {
     let template = IndexTemplate {
         tunnel_stats: stats,
         proxy_info,
+        cloudfront_info: None,
+        upload_limit: None,
+        download_limit: None,
     };
 
     Html(template.render().unwrap())

examples/ui_cloudflare.rs

@@ -30,6 +30,8 @@ async fn index() -> Html<String> {
         tunnel_stats: stats,
         proxy_info,
         cloudfront_info: None,
+        upload_limit: Some(1000),
+        download_limit: Some(1000),
     };
 
     Html(template.render().unwrap())

examples/ui_high_traffic.rs

@@ -18,6 +18,9 @@ async fn index() -> Html<String> {
     let template = IndexTemplate {
         tunnel_stats: stats,
         proxy_info: None,
+        cloudfront_info: None,
+        upload_limit: None,
+        download_limit: None,
     };
 
     Html(template.render().unwrap())

examples/ui_no_proxy.rs

@@ -1,4 +1,4 @@
-{ pkgs ? import <nixpkgs> { } }:
+{ pkgs ? import (fetchTarball "https://github.com/NixOS/nixpkgs/archive/nixos-unstable.tar.gz") { } }:
 
 with pkgs;
 

shell.nix

@@ -14,11 +14,27 @@ use tokio::sync::RwLock;
 #[folder = "assets/"]
 struct Assets;
 
-#[derive(Clone, Default)]
+#[derive(Clone)]
 pub struct AppState {
     pub stats: Arc<RwLock<TunnelStats>>,
     pub proxy_info: Arc<RwLock<Option<ProxyInfo>>>,
     pub cloudfront_info: Arc<RwLock<Option<CloudFrontInfo>>>,
+    pub tunnel: Arc<RwLock<Option<Arc<crate::wireguard::OriginTunnel>>>>,
+    pub upload_limit: Option<u32>,
+    pub download_limit: Option<u32>,
+}
+
+impl Default for AppState {
+    fn default() -> Self {
+        Self {
+            stats: Arc::new(RwLock::new(TunnelStats::default())),
+            proxy_info: Arc::new(RwLock::new(None)),
+            cloudfront_info: Arc::new(RwLock::new(None)),
+            tunnel: Arc::new(RwLock::new(None)),
+            upload_limit: None,
+            download_limit: None,
+        }
+    }
 }
 
 #[derive(Clone, Default)]
@@ -166,6 +182,8 @@ pub struct IndexTemplate {
     pub tunnel_stats: TunnelStats,
     pub proxy_info: Option<ProxyInfo>,
     pub cloudfront_info: Option<CloudFrontInfo>,
+    pub upload_limit: Option<u32>,
+    pub download_limit: Option<u32>,
 }
 
 pub async fn assets(axum::extract::Path(file): axum::extract::Path<String>) -> Response {
@@ -188,12 +206,58 @@ pub async fn assets(axum::extract::Path(file): axum::extract::Path<String>) -> R
 pub fn router(state: AppState) -> Router {
     Router::new()
         .route("/", get(index))
+        .route("/api/stats", get(stats_api))
         .route("/assets/*file", get(assets))
         .with_state(state)
 }
 
+pub async fn stats_api(State(state): State<AppState>) -> impl IntoResponse {
+    use axum::Json;
+    use serde::Serialize;
+
+    #[derive(Serialize)]
+    struct StatsResponse {
+        bytes_sent: u64,
+        bytes_received: u64,
+        bytes_sent_formatted: String,
+        bytes_received_formatted: String,
+        packets_sent: u64,
+        packets_received: u64,
+    }
+
+    let mut stats = state.stats.read().await.clone();
+
+    // Update stats from iptables if tunnel is available
+    if let Some(tunnel) = state.tunnel.read().await.as_ref() {
+        if let Ok((bytes_sent, bytes_received)) = tunnel.get_traffic_stats().await {
+            stats.bytes_sent = bytes_sent;
+            stats.bytes_received = bytes_received;
+        }
+    }
+
+    stats.format_sizes();
+
+    Json(StatsResponse {
+        bytes_sent: stats.bytes_sent,
+        bytes_received: stats.bytes_received,
+        bytes_sent_formatted: stats.bytes_sent_formatted,
+        bytes_received_formatted: stats.bytes_received_formatted,
+        packets_sent: stats.packets_sent,
+        packets_received: stats.packets_received,
+    })
+}
+
 pub async fn index(State(state): State<AppState>) -> impl IntoResponse {
     let mut stats = state.stats.read().await.clone();
+
+    // Update stats from iptables if tunnel is available
+    if let Some(tunnel) = state.tunnel.read().await.as_ref() {
+        if let Ok((bytes_sent, bytes_received)) = tunnel.get_traffic_stats().await {
+            stats.bytes_sent = bytes_sent;
+            stats.bytes_received = bytes_received;
+        }
+    }
+
     stats.format_sizes();
     let mut proxy_info = state.proxy_info.read().await.clone();
     let cloudfront_info = state.cloudfront_info.read().await.clone();
@@ -212,6 +276,8 @@ pub async fn index(State(state): State<AppState>) -> impl IntoResponse {
         tunnel_stats: stats,
         proxy_info,
         cloudfront_info,
+        upload_limit: state.upload_limit,
+        download_limit: state.download_limit,
     };
 
     Html(template.render().unwrap())

src/api.rs

@@ -5,8 +5,9 @@ pub struct CloudFormationTemplate {
     pub stack_name: String,
     pub region: String,
     pub ingress_host: String,
-    pub ingress_port: u16,
+    pub ingress_port: u16, // Primary port (first ingress) for backwards compat
     pub ingress_protocol: String,
+    pub port_mappings: Vec<(u16, String)>, // All port mappings (port, protocol)
     pub origin_host: String,
     pub origin_port: u16,
     pub origin_ip: String,
@@ -334,15 +335,19 @@ impl CloudFormationTemplate {
                 "CidrIp": format!("{}/32", self.origin_ip),
                 "Description": "WireGuard from origin"
             }),
-            json!({
-                "IpProtocol": self.ingress_protocol.as_str(),
-                "FromPort": self.ingress_port,
-                "ToPort": self.ingress_port,
-                "CidrIp": "0.0.0.0/0",
-                "Description": "Ingress traffic"
-            }),
         ];
 
+        // Add rules for each port mapping
+        for (port, protocol) in &self.port_mappings {
+            rules.push(json!({
+                "IpProtocol": protocol.to_lowercase(),
+                "FromPort": port,
+                "ToPort": port,
+                "CidrIp": "0.0.0.0/0",
+                "Description": format!("Ingress {} traffic on port {}", protocol.to_uppercase(), port)
+            }));
+        }
+
         if self.debug {
             rules.push(json!({
                 "IpProtocol": "tcp",
@@ -386,16 +391,32 @@ impl CloudFormationTemplate {
             .map(|s| format!("{}.0", s))
             .unwrap_or_else(|| "172.17.0.0".to_string());
 
+        // Generate Nix list expression for port mappings
+        // Format: [ { port = 80; protocol = "tcp"; } { port = 443; protocol = "tcp"; } ]
+        let port_mappings_nix = if self.port_mappings.is_empty() {
+            "[ ]".to_string()
+        } else {
+            let mappings: Vec<String> = self.port_mappings
+                .iter()
+                .map(|(port, protocol)| {
+                    format!(
+                        "{{ port = {}; protocol = \"{}\"; }}",
+                        port,
+                        protocol.to_lowercase()
+                    )
+                })
+                .collect();
+            format!("[\n    {}\n  ]", mappings.join("\n    "))
+        };
+
         // Replace placeholders in the Nix template
         let nix_config = NIX_TEMPLATE
             .replace(
                 "debug = false",
                 &format!("debug = {}", if self.debug { "true" } else { "false" }),
             )
             .replace("{PROXY_WG_PRIVATE_KEY}", &self.proxy_wg_private_key)
-            .replace("{PROTOCOL}", &self.ingress_protocol)
-            .replace("{INGRESS_PORT}", &self.ingress_port.to_string())
-            .replace("{ORIGIN_PORT}", &self.origin_port.to_string())
+            .replace("{PORT_MAPPINGS}", &port_mappings_nix)
             .replace("{ORIGIN_WG_PUBLIC_KEY}", &self.origin_wg_public_key)
             .replace("{PRESHARED_KEY}", &self.preshared_key)
             .replace("{ORIGIN_IP}", &self.wg_origin_ip)
@@ -420,6 +441,7 @@ mod tests {
             ingress_host: "test.example.com".to_string(),
             ingress_port: 80,
             ingress_protocol: "tcp".to_string(),
+            port_mappings: vec![(80, "tcp".to_string())],
             origin_host: "localhost".to_string(),
             origin_port: 8080,
             origin_ip: "1.2.3.4".to_string(),
@@ -445,6 +467,7 @@ mod tests {
             ingress_host: "test.example.com".to_string(),
             ingress_port: 80,
             ingress_protocol: "tcp".to_string(),
+            port_mappings: vec![(80, "tcp".to_string())],
             origin_host: "localhost".to_string(),
             origin_port: 8080,
             origin_ip: "1.2.3.4".to_string(),
@@ -470,6 +493,7 @@ mod tests {
             ingress_host: "test.example.com".to_string(),
             ingress_port: 80,
             ingress_protocol: "tcp".to_string(),
+            port_mappings: vec![(80, "tcp".to_string())],
             origin_host: "localhost".to_string(),
             origin_port: 8080,
             origin_ip: "1.2.3.4".to_string(),
@@ -486,10 +510,11 @@ mod tests {
 
         let userdata = template.generate_userdata();
         let userdata_str = serde_json::to_string(&userdata).unwrap();
-        // Check for NixOS configuration syntax and TCP protocol
+        // Check for NixOS configuration syntax and port mappings
         assert!(userdata_str.contains("{ config, pkgs, lib, ... }:"));
         assert!(userdata_str.contains("debug = false"));
-        assert!(userdata_str.contains("-p tcp --dport 80"));
+        assert!(userdata_str.contains("port = 80"));
+        assert!(userdata_str.contains("protocol = \\\"tcp\\\""));
     }
 
     #[test]
@@ -500,6 +525,7 @@ mod tests {
             ingress_host: "test.example.com".to_string(),
             ingress_port: 53,
             ingress_protocol: "udp".to_string(),
+            port_mappings: vec![(53, "udp".to_string())],
             origin_host: "localhost".to_string(),
             origin_port: 53,
             origin_ip: "1.2.3.4".to_string(),
@@ -516,10 +542,11 @@ mod tests {
 
         let userdata = template.generate_userdata();
         let userdata_str = serde_json::to_string(&userdata).unwrap();
-        // Check for NixOS configuration syntax and UDP protocol
+        // Check for NixOS configuration syntax and port mappings
         assert!(userdata_str.contains("{ config, pkgs, lib, ... }:"));
         assert!(userdata_str.contains("debug = false"));
-        assert!(userdata_str.contains("-p udp --dport 53"));
+        assert!(userdata_str.contains("port = 53"));
+        assert!(userdata_str.contains("protocol = \\\"udp\\\""));
     }
 
     #[test]
@@ -530,6 +557,7 @@ mod tests {
             ingress_host: "test.example.com".to_string(),
             ingress_port: 80,
             ingress_protocol: "tcp".to_string(),
+            port_mappings: vec![(80, "tcp".to_string())],
             origin_host: "localhost".to_string(),
             origin_port: 8080,
             origin_ip: "1.2.3.4".to_string(),