Attribute wireguard listening ports

Cargo.toml

@@ -15,8 +15,10 @@ anstyle = "1.0.8"
 anyhow = "1.0.86"
 itertools = "0.13.0"
 netlink-packet-core = "=0.7.0"
+netlink-packet-generic = "0.3.3"
 netlink-packet-route = "=0.17.1"
 netlink-packet-sock-diag = "=0.4.2"
+netlink-packet-wireguard = "0.2.3"
 netlink-sys = "=0.8.5"
 procfs = "0.16.0"
 terminal_size = "0.3.0"

src/main.rs

@@ -6,10 +6,13 @@ mod termtree;
 
 use anyhow::Result;
 use itertools::Itertools;
-use netlink::sock::{Family, SockInfo};
+use netlink::{
+    sock::{Family, SockInfo},
+    wg::wireguards,
+};
 use procfs::process::all_processes;
 use std::{
-    collections::BTreeMap,
+    collections::{BTreeMap, HashMap},
     env::var_os,
     io::{stdout, BufWriter, Write},
     net::{IpAddr, Ipv4Addr, Ipv6Addr},
@@ -21,14 +24,14 @@ pub type Ino = u64;
 
 fn main() -> Result<()> {
     let users_cache = UsersCache::new();
-    let (interfaces, local_routes) = interfaces_routes();
+    let iface_info = interfaces_routes();
 
-    let filters = options::parse_args(&interfaces, &local_routes, &users_cache)?;
+    let filters = options::parse_args(&iface_info, &users_cache)?;
 
-    let socks = netlink::sock::all_sockets(&interfaces, &local_routes); // TODO no clone, pass filters
+    let socks = netlink::sock::all_sockets(&iface_info); // TODO no clone, pass filters
     let mut socks = match socks {
         Ok(socks) => socks,
-        Err(netlink_err) => match sockets_procfs::all_sockets(&interfaces, &local_routes) {
+        Err(netlink_err) => match sockets_procfs::all_sockets(&iface_info) {
             Ok(socks) => socks,
             Err(proc_err) => {
                 eprintln!(
@@ -42,6 +45,8 @@ fn main() -> Result<()> {
     };
     let mut output = termtree::Tree::new();
     let self_user_ns = procs::get_user_ns(&procs::ourself()?).ok();
+
+    // output known processes/sockets
     let mut lps = all_processes()?
         .filter_map(|p| procs::ProcDesc::inspect_ps(p, &mut socks, &users_cache, self_user_ns).ok())
         .filter(|p| !p.sockets.is_empty())
@@ -60,6 +65,31 @@ fn main() -> Result<()> {
             );
         }
     }
+
+    // output wireguards
+    let mut wireguard_sockets = HashMap::<_, Vec<_>>::new();
+    socks.retain(|_sockid, sockinfo| {
+        if let Some(if_id) = iface_info.wireguard_ports.get(&sockinfo.port) {
+            wireguard_sockets
+                .entry(if_id)
+                .or_default()
+                .push(sockinfo.to_owned());
+            false
+        } else {
+            true
+        }
+    });
+    for (if_id, socks) in &wireguard_sockets {
+        if filters.accept_wg() {
+            let name = match iface_info.id2name.get(if_id) {
+                Some(ifname) => format!("[wireguard {ifname}]"),
+                None => format!("wireguard, index {if_id}"),
+            };
+            output.node(name, sockets_tree(socks, &filters));
+        }
+    }
+
+    // output unknown sockets
     let mut socks = socks
         .values()
         .into_group_map_by(|s| s.uid)
@@ -76,8 +106,10 @@ fn main() -> Result<()> {
             }
         }
         false => {
-            eprintln!("WARNING: Some listening sockets hidden:");
-            eprintln!("Not all sockets could not be matched to a process, process-based filtering not fully possible.");
+            if !socks.is_empty() {
+                eprintln!("WARNING: Some listening sockets hidden:");
+                eprintln!("Not all sockets could not be matched to a process, process-based filtering not fully possible.");
+            }
         }
     }
 
@@ -91,13 +123,28 @@ fn main() -> Result<()> {
     Ok(())
 }
 
-fn interfaces_routes() -> (std::collections::HashMap<u32, String>, netlink::route::Rtbl) {
+#[derive(Default)]
+struct IfaceInfo {
+    id2name: HashMap<u32, String>,
+    wireguard_ports: HashMap<u16, u32>,
+    local_routes: netlink::route::Rtbl,
+}
+
+fn interfaces_routes() -> IfaceInfo {
     let Ok(ref route_socket) = netlink::route::socket() else {
         return Default::default();
     };
-    let interfaces = netlink::route::interface_names(route_socket).unwrap_or_default();
+    let netlink::route::Interfaces {
+        id2name,
+        wireguard_ids,
+    } = netlink::route::interface_names(route_socket).unwrap_or_default();
     let local_routes = netlink::route::local_routes(route_socket).unwrap_or_default();
-    (interfaces, local_routes)
+    let wireguard_ports = wireguards(&wireguard_ids).unwrap_or_default();
+    IfaceInfo {
+        id2name,
+        wireguard_ports,
+        local_routes,
+    }
 }
 
 fn sockets_tree<'a>(

src/netlink/mod.rs

@@ -1,5 +1,6 @@
 pub mod route;
 pub mod sock;
+pub mod wg;
 
 use anyhow::{Context, Result};
 use netlink_packet_core::{
@@ -32,7 +33,10 @@ where
             match rx_packet.payload {
                 NetlinkPayload::Done(_) => return Ok(()),
                 NetlinkPayload::InnerMessage(inner) => recv(inner),
-                NetlinkPayload::Error(err) => return Err(err.to_io()).context("Netlink error"),
+                NetlinkPayload::Error(err) => match err.code {
+                    Some(_) => return Err(err.to_io()).context("Netlink error"),
+                    None => return Ok(()),
+                },
                 p => todo!("Unexpected netlink payload {:?}", p.message_type()),
             }
 

src/netlink/route.rs

@@ -4,13 +4,22 @@ use netlink_packet_core::{
     NetlinkHeader, NetlinkMessage, NetlinkPayload, NLM_F_DUMP, NLM_F_REQUEST,
 };
 use netlink_packet_route::{
-    constants::*, link::nlas::Nla as LinkNla, route::nlas::Nla as RouteNla, LinkMessage,
-    RouteMessage, RtnlMessage,
+    constants::*,
+    link::nlas::Nla as LinkNla,
+    nlas::link::{Info, InfoKind},
+    route::nlas::Nla as RouteNla,
+    LinkMessage, RouteMessage, RtnlMessage,
 };
 use netlink_sys::{protocols::NETLINK_ROUTE, Socket, SocketAddr};
 use std::{cmp::Reverse, collections::HashMap, net::IpAddr};
 
-pub fn interface_names(socket: &Socket) -> Result<HashMap<u32, String>> {
+#[derive(Default)]
+pub struct Interfaces {
+    pub id2name: HashMap<u32, String>,
+    pub wireguard_ids: Vec<u32>,
+}
+
+pub fn interface_names(socket: &Socket) -> Result<Interfaces> {
     let mut packet = NetlinkMessage::new(
         NetlinkHeader::default(),
         NetlinkPayload::from(RtnlMessage::GetLink(LinkMessage::default())),
@@ -19,19 +28,32 @@ pub fn interface_names(socket: &Socket) -> Result<HashMap<u32, String>> {
     packet.header.sequence_number = 1;
 
     let mut map = HashMap::new();
+    let mut wg_ids = Vec::new();
     drive_req(packet, socket, |inner| {
         if let RtnlMessage::NewLink(nl) = inner {
-            if let Some(name) = nl.nlas.into_iter().find_map(|s| match s {
-                LinkNla::IfName(n) => Some(n),
-                _ => None,
-            }) {
-                map.insert(nl.header.index, name);
+            for nla in nl.nlas {
+                match nla {
+                    LinkNla::IfName(name) => {
+                        map.insert(nl.header.index, name);
+                    }
+                    LinkNla::Info(infos) => {
+                        for info in infos {
+                            if info == Info::Kind(InfoKind::Wireguard) {
+                                wg_ids.push(nl.header.index);
+                            }
+                        }
+                    }
+                    _ => (),
+                }
             }
         }
     })
     .context("Get interface names")?;
 
-    Ok(map)
+    Ok(Interfaces {
+        id2name: map,
+        wireguard_ids: wg_ids,
+    })
 }
 
 pub fn socket() -> Result<Socket> {

src/netlink/sock.rs

@@ -1,5 +1,5 @@
 use super::{drive_req, nl_hdr_flags, route::Rtbl};
-use crate::Ino;
+use crate::{IfaceInfo, Ino};
 use anyhow::{Context, Result};
 use netlink_packet_core::{NetlinkMessage, NLM_F_DUMP, NLM_F_REQUEST};
 use netlink_packet_sock_diag::{
@@ -11,8 +11,11 @@ use netlink_sys::{protocols::NETLINK_SOCK_DIAG, Socket, SocketAddr};
 use std::{collections::HashMap, fmt::Display, net::IpAddr};
 
 pub fn all_sockets<'i>(
-    interfaces: &'i HashMap<u32, String>,
-    local_routes: &Rtbl,
+    IfaceInfo {
+        id2name: interfaces,
+        local_routes,
+        ..
+    }: &'i IfaceInfo,
 ) -> Result<HashMap<Ino, SockInfo<'i>>> {
     let mut socket =
         Socket::new(NETLINK_SOCK_DIAG).context("Construct netlink socket information socket")?;
@@ -149,7 +152,7 @@ impl std::str::FromStr for Protocol {
     }
 }
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq, Clone)]
 pub struct SockInfo<'a> {
     pub family: Family,
     pub protocol: Protocol,

src/options.rs

@@ -1,7 +1,7 @@
 use crate::netlink::route::Prefix;
-use crate::netlink::route::Rtbl;
 use crate::netlink::sock::Protocol;
 use crate::procs;
+use crate::IfaceInfo;
 use anyhow::bail;
 use anyhow::Context;
 use anyhow::Result;
@@ -83,6 +83,10 @@ impl Filters {
             || self.pfxs.iter().any(|pfx| pfx.matches(addr))
             || addr.is_unspecified()
     }
+
+    pub(crate) fn accept_wg(&self) -> bool {
+        self.cmd.is_empty() && self.pid.is_empty()
+    }
 }
 
 pub fn match_arg(arg: &str, args: &mut std::env::Args) -> Result<Option<(char, String)>> {
@@ -115,8 +119,11 @@ pub fn match_arg(arg: &str, args: &mut std::env::Args) -> Result<Option<(char, S
 }
 
 pub fn parse_args(
-    ifaces: &HashMap<u32, String>,
-    local_routes: &Rtbl,
+    IfaceInfo {
+        id2name: ifaces,
+        local_routes,
+        ..
+    }: &IfaceInfo,
     users: &UsersCache,
 ) -> Result<Filters> {
     let ifaces = ifaces

src/procs.rs

@@ -351,7 +351,8 @@ impl Ord for ProcDesc<'_> {
 pub fn get_user_ns(p: &Process) -> Result<u64> {
     Ok(p.namespaces()
         .context("Namespaces inaccessible")?
-        .0.get(&OsString::from_vec(b"user".to_vec()))
+        .0
+        .get(&OsString::from_vec(b"user".to_vec()))
         .context("No user ns")?
         .identifier)
 }

src/sockets_procfs.rs

@@ -1,14 +1,17 @@
 use super::Ino;
-use crate::netlink::{
-    route::Rtbl,
-    sock::{Family, Protocol, SockInfo},
+use crate::{
+    netlink::sock::{Family, Protocol, SockInfo},
+    IfaceInfo,
 };
 use anyhow::{Context, Result};
 use std::collections::HashMap;
 
 pub fn all_sockets<'i>(
-    interfaces: &'i HashMap<u32, String>,
-    local_routes: &Rtbl,
+    IfaceInfo {
+        id2name: interfaces,
+        local_routes,
+        ..
+    }: &'i IfaceInfo,
 ) -> Result<HashMap<Ino, SockInfo<'i>>> {
     eprintln!("WARNING: Falling back to parsing info from procfs, limited to TCP and UDP");
     let mut ret = HashMap::new();

src/termtree.rs

@@ -91,10 +91,7 @@ fn render_entry(
     ret(b"\n");
     if collapsed.is_none() {
         for (pos, child) in tree.children.0.iter().with_position() {
-            let last = matches!(
-                pos,
-                itertools::Position::Last | itertools::Position::Only
-            );
+            let last = matches!(pos, itertools::Position::Last | itertools::Position::Only);
             let prefix = Prefix { last, prefix };
             render_entry(child, mw, color, ret, Some(&prefix));
         }