Add ability to block wireguard traffic

3 files changed, 144 insertions, 68 deletions

diff --git a/ci/ios/test-router/raas/src/block_list/mod.rs b/ci/ios/test-router/raas/src/block_list/mod.rs
index 8aab004126..82eece7bb1 100644
--- a/ci/ios/test-router/raas/src/block_list/mod.rs
+++ b/ci/ios/test-router/raas/src/block_list/mod.rs
@@ -6,7 +6,7 @@ static TABLE_NAME: Lazy<CString> = Lazy::new(|| CString::new("raas").unwrap());
 static FORWARD_CHAIN_NAME: Lazy<CString> = Lazy::new(|| CString::new("forward").unwrap());
 
 mod rule;
-pub use rule::BlockRule;
+pub use rule::{BlockRule, Endpoints};
 
 #[derive(Default)]
 pub struct BlockList {
diff --git a/ci/ios/test-router/raas/src/block_list/rule.rs b/ci/ios/test-router/raas/src/block_list/rule.rs
index e6ce90f20c..5bd876fe96 100644
--- a/ci/ios/test-router/raas/src/block_list/rule.rs
+++ b/ci/ios/test-router/raas/src/block_list/rule.rs
@@ -1,14 +1,24 @@
 use crate::web::routes::TransportProtocol;
 use mnl::mnl_sys::libc;
-use nftnl::{expr, nft_expr, Chain, Rule};
+use nftnl::{expr, nft_expr, nft_expr_payload, Chain, Rule};
 
 use std::{collections::BTreeSet, iter, net::IpAddr};
 
 #[derive(Clone, serde::Serialize)]
-pub struct BlockRule {
+pub enum BlockRule {
+    Host {
+        endpoints: Endpoints,
+        protocols: BTreeSet<TransportProtocol>,
+    },
+    WireGuard {
+        endpoints: Endpoints,
+    },
+}
+
+#[derive(Clone, serde::Serialize)]
+pub struct Endpoints {
     pub src: IpAddr,
     pub dst: IpAddr,
-    pub protocols: BTreeSet<TransportProtocol>,
 }
 
 impl BlockRule {
@@ -16,15 +26,15 @@ impl BlockRule {
         &'a self,
         chain: &'a Chain<'a>,
     ) -> Box<dyn Iterator<Item = Rule<'a>> + 'a> {
-        if self.protocols.is_empty() {
-            return Box::new(iter::once(self.create_nft_rule_inner(chain, None)));
+        match self {
+            BlockRule::Host { protocols, .. } if !protocols.is_empty() => {
+                let iter = protocols
+                    .iter()
+                    .map(|protocol| self.create_nft_rule_inner(chain, Some(*protocol)));
+                Box::new(iter)
+            }
+            _ => Box::new(iter::once(self.create_nft_rule_inner(chain, None))),
         }
-
-        let iter = self
-            .protocols
-            .iter()
-            .map(|protocol| self.create_nft_rule_inner(chain, Some(*protocol)));
-        Box::new(iter)
     }
 
     fn create_nft_rule_inner<'a>(
@@ -33,33 +43,29 @@ impl BlockRule {
         transport_protocol: Option<TransportProtocol>,
     ) -> Rule<'a> {
         let mut rule = Rule::new(chain);
-        check_l3proto(&mut rule, self.src);
-        if let Some(protocol) = transport_protocol {
-            check_l4proto(&mut rule, protocol);
-        };
 
-        // Add source checking
-        rule.add_expr(match self.src {
-            IpAddr::V4(_) => &nft_expr!(payload ipv4 saddr),
-            IpAddr::V6(_) => &nft_expr!(payload ipv6 saddr),
-        });
-        match self.src {
-            IpAddr::V4(addr) => rule.add_expr(&nft_expr!(cmp == addr)),
-            IpAddr::V6(addr) => rule.add_expr(&nft_expr!(cmp == addr)),
-        };
+        match *self {
+            BlockRule::Host {
+                endpoints: Endpoints { src, dst },
+                ..
+            } => {
+                check_l3proto(&mut rule, src);
+                if let Some(protocol) = transport_protocol {
+                    check_l4proto(&mut rule, protocol);
+                };
+                check_ip_addrs(&mut rule, src, dst);
+            }
+            BlockRule::WireGuard {
+                endpoints: Endpoints { src, dst },
+            } => {
+                check_l3proto(&mut rule, src);
+                check_ip_addrs(&mut rule, src, dst);
+                check_wireguard_traffic(&mut rule);
+            }
+        }
 
-        // Add destination check
-        rule.add_expr(match self.dst {
-            IpAddr::V4(_) => &nft_expr!(payload ipv4 daddr),
-            IpAddr::V6(_) => &nft_expr!(payload ipv6 daddr),
-        });
-        match self.dst {
-            IpAddr::V4(addr) => rule.add_expr(&nft_expr!(cmp == addr)),
-            IpAddr::V6(addr) => rule.add_expr(&nft_expr!(cmp == addr)),
-        };
         rule.add_expr(&nft_expr!(counter));
         rule.add_expr(&expr::Verdict::Drop);
-
         rule
     }
 }
@@ -80,3 +86,42 @@ fn check_l4proto(rule: &mut Rule<'_>, protocol: TransportProtocol) {
     rule.add_expr(&nft_expr!(meta l4proto));
     rule.add_expr(&nft_expr!(cmp == protocol.as_ipproto()));
 }
+
+fn check_ip_addrs(rule: &mut Rule, src: IpAddr, dst: IpAddr) {
+    // Add source checking
+    rule.add_expr(match src {
+        IpAddr::V4(_) => &nft_expr!(payload ipv4 saddr),
+        IpAddr::V6(_) => &nft_expr!(payload ipv6 saddr),
+    });
+    match src {
+        IpAddr::V4(addr) => rule.add_expr(&nft_expr!(cmp == addr)),
+        IpAddr::V6(addr) => rule.add_expr(&nft_expr!(cmp == addr)),
+    };
+
+    // Add destination check
+    rule.add_expr(match dst {
+        IpAddr::V4(_) => &nft_expr!(payload ipv4 daddr),
+        IpAddr::V6(_) => &nft_expr!(payload ipv6 daddr),
+    });
+    match dst {
+        IpAddr::V4(addr) => rule.add_expr(&nft_expr!(cmp == addr)),
+        IpAddr::V6(addr) => rule.add_expr(&nft_expr!(cmp == addr)),
+    };
+}
+
+fn check_wireguard_traffic(rule: &mut Rule) {
+    rule.add_expr(&nft_expr!(meta l4proto));
+    rule.add_expr(&nft_expr!(cmp == libc::IPPROTO_UDP));
+
+    // UDP header is 8 bytes, after which we have the WireGuard header which is 4 bytes,
+    // where the first byte can be 1 to 4 (inclusive) and the last 3 bytes are 0.
+    // See: https://wiki.wireshark.org/WireGuard
+    rule.add_expr(&nft_expr_payload!(th 8, 1));
+    rule.add_expr(&nft_expr!(cmp >= 1));
+
+    rule.add_expr(&nft_expr_payload!(th 8, 1));
+    rule.add_expr(&nft_expr!(cmp <= 4));
+
+    rule.add_expr(&nft_expr_payload!(th 9, 3));
+    rule.add_expr(&nft_expr!(cmp == 0));
+}
diff --git a/ci/ios/test-router/raas/src/web/routes.rs b/ci/ios/test-router/raas/src/web/routes.rs
index a2fae72739..2b6432c5a2 100644
--- a/ci/ios/test-router/raas/src/web/routes.rs
+++ b/ci/ios/test-router/raas/src/web/routes.rs
@@ -1,28 +1,32 @@
-use std::{
-    collections::{BTreeMap, BTreeSet},
-    net::IpAddr,
-};
-
 use axum::{
     extract::{Json, Path, State},
     http::StatusCode,
     response::IntoResponse,
 };
 use mnl::mnl_sys::libc;
+use std::sync::MutexGuard;
+use std::{
+    collections::{BTreeMap, BTreeSet},
+    net::IpAddr,
+};
 use uuid::Uuid;
 
-use crate::block_list::BlockRule;
+use crate::block_list::{BlockList, BlockRule, Endpoints};
+use crate::web;
 
 #[derive(serde::Deserialize, Clone)]
 pub struct NewRule {
     pub src: IpAddr,
     pub dst: IpAddr,
     pub protocols: Option<BTreeSet<TransportProtocol>>,
+    #[serde(default)]
+    pub block_wireguard: bool,
     pub label: Uuid,
 }
 
 #[derive(serde::Deserialize, serde::Serialize, PartialOrd, Ord, PartialEq, Eq, Clone, Copy)]
 #[serde(rename_all = "snake_case")]
+#[derive(Debug)]
 pub enum TransportProtocol {
     Tcp,
     Udp,
@@ -43,30 +47,29 @@ impl TransportProtocol {
 
 pub async fn add_rule(
     State(state): State<super::State>,
-    Json(rule): Json<NewRule>,
+    Json(json): Json<NewRule>,
 ) -> impl IntoResponse {
-    let result = tokio::task::spawn_blocking(move || -> anyhow::Result<()> {
-        let label = rule.label;
-        let rule = BlockRule {
-            src: rule.src,
-            dst: rule.dst,
-            protocols: rule.protocols.unwrap_or_default(),
-        };
-        let Ok(mut fw) = state.block_list.lock() else {
-            return Err(anyhow::anyhow!("Firewall thread panicked"));
+    let result = access_firewall(state, move |mut fw| {
+        let label = json.label;
+        let src = json.src;
+        let dst = json.dst;
+
+        let rule = if json.block_wireguard {
+            BlockRule::WireGuard {
+                endpoints: Endpoints { src, dst },
+            }
+        } else {
+            BlockRule::Host {
+                endpoints: Endpoints { src, dst },
+                protocols: json.protocols.unwrap_or_default(),
+            }
         };
 
         fw.add_rule(rule.clone(), label)?;
-        log::info!(
-            "Successfully added a rule to block {} from {} for test {}",
-            rule.src,
-            rule.dst,
-            label,
-        );
+        log_rule(&rule, &label);
         Ok(())
     })
-    .await
-    .expect("failed to join blocking task");
+    .await;
 
     respond_with_result(result, StatusCode::CREATED)
 }
@@ -75,20 +78,30 @@ pub async fn delete_rules(
     Path(label): Path<Uuid>,
     State(state): State<super::State>,
 ) -> impl IntoResponse {
-    let result = tokio::task::spawn_blocking(move || -> anyhow::Result<()> {
-        let Ok(mut fw) = state.block_list.lock() else {
-            return Err(anyhow::anyhow!("Firewall thread panicked"));
-        };
-
+    let result = access_firewall(state, move |mut fw| {
         fw.clear_rules_with_label(&label)?;
-        log::info!("Successfully removed all rules for test {label}",);
+        log::info!("Successfully removed all rules for test {label}");
         Ok(())
     })
-    .await
-    .expect("failed to join blocking task");
+    .await;
+
     respond_with_result(result, StatusCode::OK)
 }
 
+pub async fn access_firewall<F>(state: web::State, run: F) -> anyhow::Result<()>
+where
+    F: FnOnce(MutexGuard<BlockList>) -> anyhow::Result<()> + Send + 'static,
+{
+    tokio::task::spawn_blocking(move || -> anyhow::Result<()> {
+        let Ok(fw) = state.block_list.lock() else {
+            return Err(anyhow::anyhow!("Firewall thread panicked"));
+        };
+        run(fw)
+    })
+    .await
+    .expect("failed to join blocking task")
+}
+
 fn respond_with_result(result: anyhow::Result<()>, success_code: StatusCode) -> impl IntoResponse {
     match result {
         Ok(_) => (success_code, String::new()),
@@ -108,7 +121,25 @@ pub async fn list_all_rules(State(state): State<super::State>) -> impl IntoRespo
     .expect("failed to join blocking task");
 
     match all_rules {
-        Ok(all_rules) => axum::Json(all_rules).into_response(),
+        Ok(all_rules) => Json(all_rules).into_response(),
         Err(err) => (StatusCode::SERVICE_UNAVAILABLE, format!("{err}\n")).into_response(),
     }
 }
+
+fn log_rule(rule: &BlockRule, label: &Uuid) {
+    match rule {
+        BlockRule::Host {
+            protocols,
+            endpoints: Endpoints { src, dst },
+        } => {
+            log::info!(
+                "Successfully added a rule to block {src} from {dst} for test {label} for protocols {protocols:?}",
+            );
+        }
+        BlockRule::WireGuard {
+            endpoints: Endpoints { src, dst },
+        } => {
+            log::info!("Successfully added a rule to block {src} from {dst} WireGuard traffic for test {label}",);
+        }
+    }
+}