chore: remove unused skill file for PR fixes

- switch back to let mut args: Args = Args::parse()

.cargo/config.toml

@@ -0,0 +1,2 @@
+[target.aarch64-apple-darwin]
+runner = 'sudo -E'

.cirrus.yml

@@ -1,39 +1,60 @@
+use_compute_credits: true
+
+macos_instance:
+  image: ghcr.io/cirruslabs/macos-runner:tahoe
+
 env:
   PATH: "$PATH:$HOME/.cargo/bin"
 
 task:
-  name: Build
-  macos_instance:
-    image: ghcr.io/cirruslabs/macos-sonoma-xcode:latest
-  install_rust_script:
-    - curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
-    - rustup toolchain install nightly-aarch64-apple-darwin
-    - rustup target add x86_64-apple-darwin
-  build_script:
-    - cargo build --target aarch64-apple-darwin --target x86_64-apple-darwin
+  name: Lint
+  install_rust_script: curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+  rustfmt_script: cargo fmt --check
+  clippy_script: cargo clippy --all-targets --all-features -- -D warnings
+
+task:
+  alias: Test
+  matrix:
+    - name: Test on macOS Sequoia
+      macos_instance:
+        image: ghcr.io/cirruslabs/macos-runner:sequoia
+    - name: Test on macOS Tahoe
+      macos_instance:
+        image: ghcr.io/cirruslabs/macos-runner:tahoe
+  install_rust_script: curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+  test_script: cargo test
+
+task:
+  name: Release (Dry Run)
+  only_if: $CIRRUS_TAG == ''
+  depends_on:
+    - Lint
+    - Test
+  install_rust_script: curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+  install_script: brew install go
+  install_goreleaser_script: brew install --cask goreleaser/tap/goreleaser-pro
+  build_script: goreleaser build --snapshot
+  goreleaser_artifacts:
+    path: "dist/**"
 
 task:
   name: Release
   only_if: $CIRRUS_TAG != ''
-  macos_instance:
-    image: ghcr.io/cirruslabs/macos-sonoma-xcode:latest
+  depends_on:
+    - Lint
+    - Test
   env:
     GITHUB_TOKEN: ENCRYPTED[!98ace8259c6024da912c14d5a3c5c6aac186890a8d4819fad78f3e0c41a4e0cd3a2537dd6e91493952fb056fa434be7c!]
     GORELEASER_KEY: ENCRYPTED[!9b80b6ef684ceaf40edd4c7af93014ee156c8aba7e6e5795f41c482729887b5c31f36b651491d790f1f668670888d9fd!]
     SENTRY_ORG: cirrus-labs
     SENTRY_PROJECT: persistent-workers
     SENTRY_AUTH_TOKEN: ENCRYPTED[!c16a5cf7da5f856b4bc2f21fe8cb7aa2a6c981f851c094ed4d3025fd02ea59a58a86cee8b193a69a1fc20fa217e56ac3!]
-  install_rust_script:
-    - curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
-    - rustup toolchain install nightly-aarch64-apple-darwin
-    - rustup target add x86_64-apple-darwin
-  install_script:
-    - brew install go goreleaser/tap/goreleaser-pro getsentry/tools/sentry-cli
-  build_script:
-    - cargo build --target aarch64-apple-darwin --target x86_64-apple-darwin --profile release-with-debug
+  install_rust_script: curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+  install_script: brew install go getsentry/tools/sentry-cli
+  install_goreleaser_script: brew install --cask goreleaser/tap/goreleaser-pro
   release_script: goreleaser
   upload_sentry_debug_files_script:
-    - cd target/aarch64-apple-darwin/release-with-debug/
+    - cd target/aarch64-apple-darwin/release/
     # Generate and upload symbols
     - dsymutil softnet
     - sentry-cli debug-files upload -o $SENTRY_ORG -p $SENTRY_PROJECT softnet.dSYM/

.goreleaser.yml

@@ -1,26 +1,27 @@
+---
+version: 2
 project_name: softnet
 
 builds:
-  - id: softnet
-    builder: prebuilt
-    goamd64: [v1]
-    goos:
-      - darwin
-    goarch:
-      - arm64
-      - amd64
-    prebuilt:
-      path: 'target/{{- if eq .Arch "arm64" }}aarch64{{- else }}x86_64{{ end }}-apple-darwin/release-with-debug/softnet'
+  - builder: rust
+    command: build
+    targets:
+      - aarch64-apple-darwin
+      - x86_64-apple-darwin
+
+universal_binaries:
+  - replace: true
 
 archives:
-  - id: regular
-    name_template: "{{ .ProjectName }}-{{ .Arch }}"
+  - name_template: "{{ .ProjectName }}"
+    formats:
+      - tar.gz
 
 release:
   prerelease: auto
 
 brews:
-  - name: softnet
+  - name: "{{ .ProjectName }}"
     repository:
       owner: cirruslabs
       name: homebrew-cli
@@ -28,3 +29,5 @@ brews:
     homepage: https://github.com/cirruslabs/softnet
     description: Software networking with isolation for Tart
     skip_upload: auto
+    custom_block: |
+      depends_on :macos => :sequoia

Cargo.toml

@@ -2,7 +2,7 @@
 name = "softnet"
 version = "0.1.0"
 publish = false
-edition = "2021"
+edition = "2024"
 
 [lib]
 path = "lib/mod.rs"
@@ -16,7 +16,7 @@ smoltcp = "0"
 libc = "0"
 polling = "3"
 dhcproto = { git = "https://github.com/bluecatengineering/dhcproto.git", branch = "master" }
-vmnet = "0"
+vmnet = "0.5.0"
 clap = { version = "4", features = ["derive"] }
 mac_address = "1"
 privdrop = "0"
@@ -27,8 +27,12 @@ system-configuration = "0"
 num_enum = "0"
 sentry = { version = "0", features = ["debug-images"] }
 sentry-anyhow = { version = "0", features = ["backtrace"] }
-nix = { version = "0", features = ["signal"] }
+nix = { version = "0", features = ["signal", "socket"] }
 prefix-trie = "0"
 ipnet = "2"
 oslog = "0.2.0"
-log = "0.4.27"
+log = "0.4.29"
+serial_test = "3"
+
+[profile.release]
+debug = true

lib/dhcp_snooper.rs

@@ -1,5 +1,5 @@
-use dhcproto::v4::{DhcpOption, MessageType, OptionCode};
 use dhcproto::Decodable;
+use dhcproto::v4::{DhcpOption, MessageType, OptionCode};
 use smoltcp::wire::Ipv4Address;
 use std::collections::HashSet;
 use std::time::{Duration, Instant};
@@ -45,6 +45,11 @@ impl DhcpSnooper {
         };
     }
 
+    #[cfg(test)]
+    pub(crate) fn set_lease(&mut self, vm_lease: Option<Lease>) {
+        self.vm_lease = vm_lease
+    }
+
     pub fn lease(&self) -> &Option<Lease> {
         &self.vm_lease
     }
@@ -58,6 +63,7 @@ impl DhcpSnooper {
     }
 }
 
+#[derive(Debug)]
 pub struct Lease {
     address: Ipv4Address,
     valid_until: Instant,
@@ -65,7 +71,7 @@ pub struct Lease {
 }
 
 impl Lease {
-    fn new(address: Ipv4Address, lease_time: Duration, dns_ips: HashSet<Ipv4Address>) -> Lease {
+    pub fn new(address: Ipv4Address, lease_time: Duration, dns_ips: HashSet<Ipv4Address>) -> Lease {
         Lease {
             address,
             valid_until: Instant::now() + lease_time,

lib/host.rs

@@ -1,15 +1,15 @@
-use anyhow::{anyhow, Context, Result};
+use anyhow::{Context, Result, anyhow};
 use clap::ValueEnum;
 use log::info;
 use std::net::Ipv4Addr;
 use std::os::unix::io::{AsRawFd, RawFd};
 use std::os::unix::net::UnixDatagram;
 use std::str::FromStr;
-use std::sync::mpsc::{sync_channel, SyncSender};
+use std::sync::mpsc::{SyncSender, sync_channel};
 use vmnet::mode::Mode;
 use vmnet::parameters::{Parameter, ParameterKind};
 use vmnet::port_forwarding::{AddressFamily, Protocol};
-use vmnet::{Events, Options};
+use vmnet::{Batch, Events, Options};
 
 #[derive(ValueEnum, Clone, Debug)]
 pub enum NetType {
@@ -29,6 +29,7 @@ pub struct Host {
     callback_can_continue_tx: SyncSender<()>,
     pub gateway_ip: smoltcp::wire::Ipv4Address,
     pub max_packet_size: u64,
+    pub read_max_packets: u64,
     finalized: bool,
 }
 
@@ -67,6 +68,15 @@ impl Host {
             ));
         };
 
+        // Retrieve read max packets for this interface
+        let Some(Parameter::ReadMaxPackets(read_max_packets)) =
+            interface.parameters().get(ParameterKind::ReadMaxPackets)
+        else {
+            return Err(anyhow!(
+                "failed to retrieve vmnet's interface read max packets"
+            ));
+        };
+
         // Set up a socketpair() to emulate polling of the vmnet interface
         let (new_packets_tx, new_packets_rx) = UnixDatagram::pair()?;
         new_packets_rx.set_nonblocking(true)?;
@@ -96,6 +106,7 @@ impl Host {
             callback_can_continue_tx,
             gateway_ip,
             max_packet_size,
+            read_max_packets,
             finalized: false,
         })
     }
@@ -108,7 +119,9 @@ impl Host {
         internal_addr: Ipv4Addr,
         internal_port: u16,
     ) -> Result<()> {
-        let details = format!("external_port={external_port}, internal_addr={internal_addr}, internal_port={internal_port}");
+        let details = format!(
+            "external_port={external_port}, internal_addr={internal_addr}, internal_port={internal_port}"
+        );
 
         self.interface
             .port_forwarding_rule_add(
@@ -131,14 +144,14 @@ impl Host {
             .map_err(|err| anyhow!("failed to remove port forwarding rule {details}: {err}"))
     }
 
-    pub fn read(&mut self, buf: &mut [u8]) -> vmnet::Result<usize> {
+    pub fn read(&mut self, batch: &mut Batch, bufs: &mut [Vec<u8>]) -> vmnet::Result<usize> {
         // Dequeue dummy datagram from the socket (if any)
         // to free up buffer space and reduce false-positives
         // when polling
         let mut buf_to_be_discarded: [u8; 1] = [0; 1];
         let _ = self.new_packets_rx.recv(&mut buf_to_be_discarded);
 
-        let result = self.interface.read(buf);
+        let result = self.interface.read_batch(batch, bufs);
 
         if let Err(vmnet::Error::VmnetReadNothing) = result {
             // We've emptied everything, unlock the callback

lib/poller.rs

@@ -1,7 +1,7 @@
 use anyhow::Result;
 use num_enum::IntoPrimitive;
-use polling::os::kqueue::PollerKqueueExt;
 use polling::PollMode;
+use polling::os::kqueue::PollerKqueueExt;
 use std::os::fd::{AsRawFd, BorrowedFd};
 use std::os::unix::io::RawFd;
 use std::time::Duration;

lib/proxy/exposed_port.rs

@@ -1,4 +1,4 @@
-use anyhow::{anyhow, Context, Error};
+use anyhow::{Context, Error, anyhow};
 use std::str::FromStr;
 
 #[derive(Debug, Clone, Copy, Default, PartialEq)]

lib/proxy/host.rs

@@ -1,5 +1,5 @@
-use crate::proxy::udp_packet_helper::UdpPacketHelper;
 use crate::proxy::Proxy;
+use crate::proxy::udp_packet_helper::UdpPacketHelper;
 use anyhow::{Context, Result};
 use smoltcp::wire::{EthernetFrame, EthernetProtocol, Ipv4Packet, UdpPacket};
 

lib/proxy/mod.rs

@@ -14,10 +14,11 @@ pub use exposed_port::ExposedPort;
 use ipnet::Ipv4Net;
 use mac_address::MacAddress;
 use port_forwarder::PortForwarder;
-use prefix_trie::{Prefix, PrefixSet};
+use prefix_trie::{Prefix, PrefixMap, PrefixSet};
 use smoltcp::wire::EthernetFrame;
 use std::io::ErrorKind;
 use std::os::unix::io::{AsRawFd, RawFd};
+use vmnet::Batch;
 
 pub struct Proxy<'proxy> {
     vm: VM,
@@ -25,38 +26,67 @@ pub struct Proxy<'proxy> {
     poller: Poller<'proxy>,
     vm_mac_address: smoltcp::wire::EthernetAddress,
     dhcp_snooper: DhcpSnooper,
-    allow: PrefixSet<Ipv4Net>,
+    rules: PrefixMap<Ipv4Net, Action>,
     enobufs_encountered: bool,
     port_forwarder: PortForwarder,
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) enum Action {
+    Block,
+    Allow,
+}
+
 impl Proxy<'_> {
     pub fn new<'proxy>(
         vm_fd: RawFd,
         vm_mac_address: MacAddress,
         vm_net_type: NetType,
         allow: PrefixSet<Ipv4Net>,
+        block: PrefixSet<Ipv4Net>,
         exposed_ports: Vec<ExposedPort>,
     ) -> Result<Proxy<'proxy>> {
         let vm = VM::new(vm_fd)?;
         let host = Host::new(vm_net_type, !allow.contains(&Ipv4Net::zero()))?;
         let poller = Poller::new(vm.as_raw_fd(), host.as_raw_fd())?;
 
+        // Craft packet filter rules
+        //
+        // SECURITY: blocking rules must always take precedence
+        // over allowing rules when prefixes are identical.
+        let mut rules = PrefixMap::new();
+
+        for allow_net in allow {
+            rules.insert(allow_net, Action::Allow);
+        }
+
+        for block_net in block {
+            rules.insert(block_net, Action::Block);
+        }
+
         Ok(Proxy {
             vm,
             host,
             poller,
             vm_mac_address: smoltcp::wire::EthernetAddress(vm_mac_address.bytes()),
             dhcp_snooper: Default::default(),
-            allow,
+            rules,
             enobufs_encountered: false,
             port_forwarder: PortForwarder::new(exposed_ports),
         })
     }
 
     pub fn run(&mut self) -> Result<()> {
+        // Create a single buffer from reading from the VM
         let mut buf: Vec<u8> = vec![0; self.host.max_packet_size as usize];
 
+        // Create multiple buffers and a batch for reading from the host
+        let mut bufs = vec![
+            vec![0u8; self.host.max_packet_size as usize];
+            self.host.read_max_packets as usize
+        ];
+        let mut batch = Batch::preallocate(bufs.len());
+
         self.poller.arm()?;
 
         loop {
@@ -67,7 +97,7 @@ impl Proxy<'_> {
             }
 
             if host_readable {
-                self.read_from_host(buf.as_mut_slice())?;
+                self.read_from_host(&mut batch, &mut bufs)?;
             }
 
             // Graceful termination
@@ -104,12 +134,14 @@ impl Proxy<'_> {
         }
     }
 
-    fn read_from_host(&mut self, buf: &mut [u8]) -> Result<()> {
+    fn read_from_host(&mut self, batch: &mut Batch, bufs: &mut [Vec<u8>]) -> Result<()> {
         loop {
-            match self.host.read(buf) {
-                Ok(n) => {
-                    if let Ok(pkt) = EthernetFrame::new_checked(&buf[..n]) {
-                        self.process_frame_from_host(&pkt)?;
+            match self.host.read(batch, bufs) {
+                Ok(pktcnt) => {
+                    for buf in batch.packet_sized_bufs(bufs).take(pktcnt) {
+                        if let Ok(pkt) = EthernetFrame::new_checked(buf) {
+                            self.process_frame_from_host(&pkt)?;
+                        }
                     }
                 }
                 Err(err) => {
@@ -123,3 +155,105 @@ impl Proxy<'_> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::NetType;
+    use crate::dhcp_snooper::Lease;
+    use crate::proxy::{Action, Proxy};
+    use ipnet::Ipv4Net;
+    use mac_address::MacAddress;
+    use nix::sys::socket::{AddressFamily, SockFlag, SockType, socketpair};
+    use prefix_trie::{PrefixMap, PrefixSet};
+    use serial_test::serial;
+    use smoltcp::wire::{Ipv4Address, Ipv4Packet};
+    use std::collections::HashSet;
+    use std::os::fd::AsRawFd;
+    use std::str::FromStr;
+    use std::time::Duration;
+
+    #[test]
+    #[serial]
+    fn test_blocking_takes_precedence() {
+        let vm_ip = Ipv4Address::from_str("192.168.0.2").unwrap();
+        let proxy = create_proxy(vm_ip, vec!["66.66.0.0/16"], vec!["66.66.0.0/16"]);
+
+        assert_eq!(
+            proxy.rules,
+            PrefixMap::<Ipv4Net, Action>::from_iter(vec![(
+                Ipv4Net::from_str("66.66.0.0/16").unwrap(),
+                Action::Block
+            ),])
+        );
+
+        assert!(allowed_from_vm_ipv4(&proxy, vm_ip, "66.66.66.66").is_none());
+    }
+
+    #[test]
+    #[serial]
+    fn test_longest_prefix_match_wins() {
+        let vm_ip = Ipv4Address::from_str("192.168.0.2").unwrap();
+        let proxy = create_proxy(vm_ip, vec!["33.33.33.33/32"], vec!["33.33.33.0/24"]);
+
+        assert_eq!(
+            proxy.rules,
+            PrefixMap::<Ipv4Net, Action>::from_iter(vec![
+                (Ipv4Net::from_str("33.33.33.33/32").unwrap(), Action::Allow),
+                (Ipv4Net::from_str("33.33.33.0/24").unwrap(), Action::Block),
+            ])
+        );
+
+        assert!(allowed_from_vm_ipv4(&proxy, vm_ip, "33.33.33.32").is_none());
+        assert!(allowed_from_vm_ipv4(&proxy, vm_ip, "33.33.33.33").is_some());
+        assert!(allowed_from_vm_ipv4(&proxy, vm_ip, "33.33.33.34").is_none());
+    }
+
+    fn create_proxy<'test>(vm_ip: Ipv4Address, allow: Vec<&str>, block: Vec<&str>) -> Proxy<'test> {
+        let (vm_fd, _) = socketpair(
+            AddressFamily::Unix,
+            SockType::Datagram,
+            None,
+            SockFlag::empty(),
+        )
+        .unwrap();
+        let vm_fd = Box::leak(Box::new(vm_fd));
+
+        let mut proxy = Proxy::new(
+            vm_fd.as_raw_fd(),
+            MacAddress::from_str("02:00:00:00:00:01").unwrap(),
+            NetType::Nat,
+            PrefixSet::from_iter(
+                allow
+                    .into_iter()
+                    .map(|cidr| Ipv4Net::from_str(cidr).unwrap()),
+            ),
+            PrefixSet::from_iter(
+                block
+                    .into_iter()
+                    .map(|cidr| Ipv4Net::from_str(cidr).unwrap()),
+            ),
+            Vec::default(),
+        )
+        .unwrap();
+
+        proxy.dhcp_snooper.set_lease(Some(Lease::new(
+            vm_ip,
+            Duration::from_secs(600),
+            HashSet::new(),
+        )));
+
+        proxy
+    }
+
+    fn allowed_from_vm_ipv4(proxy: &Proxy, src: Ipv4Address, dst: &str) -> Option<()> {
+        let mut buf = vec![0; 1500];
+
+        let mut ipv4_pkt_mut = Ipv4Packet::new_unchecked(&mut buf[..]);
+        ipv4_pkt_mut.set_src_addr(src);
+        ipv4_pkt_mut.set_dst_addr(Ipv4Address::from_str(dst).unwrap());
+
+        let ipv4_pkt = Ipv4Packet::new_unchecked(buf.as_slice());
+
+        proxy.allowed_from_vm_ipv4(ipv4_pkt)
+    }
+}

lib/proxy/vm.rs

@@ -1,5 +1,5 @@
 use crate::proxy::udp_packet_helper::UdpPacketHelper;
-use crate::proxy::Proxy;
+use crate::proxy::{Action, Proxy};
 use anyhow::Context;
 use anyhow::Result;
 use ipnet::Ipv4Net;
@@ -58,43 +58,54 @@ impl Proxy<'_> {
         None
     }
 
-    fn allowed_from_vm_ipv4(&self, ipv4_pkt: Ipv4Packet<&[u8]>) -> Option<()> {
-        // Have we learned the VM's IP from the DHCP snooping?
-        if let Some(lease) = &self.dhcp_snooper.lease() {
-            // If so, allow all global traffic
+    pub(crate) fn allowed_from_vm_ipv4(&self, ipv4_pkt: Ipv4Packet<&[u8]>) -> Option<()> {
+        // Is this packet coming from VM's IP address that we've learned from DHCP snooping?
+        if let Some(lease) = &self.dhcp_snooper.lease()
+            && lease.valid_ip_source(ipv4_pkt.src_addr())
+        {
             let dst_addr = ipv4_pkt.dst_addr();
-            let dst_is_global = ip_network::IpNetwork::from(dst_addr).is_global();
 
-            if lease.valid_ip_source(ipv4_pkt.src_addr()) && dst_is_global {
+            // Filter traffic based on user-specified rules first
+            if !self.rules.is_empty() {
+                let dst_net = Ipv4Net::from(dst_addr);
+
+                if let Some((_, action)) = self.rules.get_lpm(&dst_net) {
+                    return match action {
+                        Action::Allow => Some(()),
+                        Action::Block => None,
+                    };
+                }
+            }
+
+            // When no user-specified rules matched, simply allow all global traffic
+            if ip_network::IpNetwork::from(dst_addr).is_global() {
                 return Some(());
             }
 
-            // Also allow all traffic to the user-specified CIDRs
-            let dst_net = Ipv4Net::from(dst_addr);
-
-            // Use get_lpm() instead of get_spm() to work around prefix-trie
-            // not handling prefixes like 0.0.0.0/0 correctly[1]
-            //
-            // [1]: https://github.com/tiborschneider/prefix-trie/issues/8
-            if self.allow.get_lpm(&dst_net).is_some() {
+            // Additionally, allow communication with the host,
+            // otherwise things like SSH to a VM won't work
+            if ipv4_pkt.dst_addr() == self.host.gateway_ip {
                 return Some(());
             }
+
+            // Additionally, allow DNS requests to DNS-servers
+            // provided to a VM by the host's DHCP server
+            if ipv4_pkt.next_header() == IpProtocol::Udp {
+                let udp_pkt = UdpPacket::new_checked(ipv4_pkt.payload()).ok()?;
+
+                if udp_pkt.is_dns_request()
+                    && self.dhcp_snooper.valid_dns_target(&ipv4_pkt.dst_addr())
+                {
+                    return Some(());
+                }
+            }
         }
 
-        // Allow communication with host
-        if ipv4_pkt.dst_addr() == self.host.gateway_ip {
-            return Some(());
-        }
-
+        // Allow outgoing DHCP requests to broadcast addresses,
+        // otherwise DHCP snooper will never be populated
         if ipv4_pkt.next_header() == IpProtocol::Udp {
             let udp_pkt = UdpPacket::new_checked(ipv4_pkt.payload()).ok()?;
 
-            // Allow DNS communication with the DNS-servers provided by DHCP
-            if udp_pkt.is_dns_request() && self.dhcp_snooper.valid_dns_target(&ipv4_pkt.dst_addr())
-            {
-                return Some(());
-            }
-
             // Allow DHCP communication with the bootpd(8) on host via broadcast address
             if udp_pkt.is_dhcp_request() && ipv4_pkt.dst_addr().is_broadcast() {
                 return Some(());

src/main.rs

@@ -1,20 +1,22 @@
-use anyhow::{anyhow, Context};
+use anyhow::{Context, Error, anyhow};
 use clap::Parser;
 use ipnet::Ipv4Net;
 use log::LevelFilter;
-use nix::sys::signal::{signal, SigHandler, Signal};
+use nix::sys::signal::{SigHandler, Signal, signal};
 use oslog::OsLogger;
 use prefix_trie::PrefixSet;
 use privdrop::PrivDrop;
+use softnet::NetType;
 use softnet::proxy::ExposedPort;
 use softnet::proxy::Proxy;
-use softnet::NetType;
 use std::borrow::Cow;
 use std::env;
+use std::net::{Ipv4Addr, SocketAddr, ToSocketAddrs};
 use std::os::raw::c_int;
 use std::os::unix::io::RawFd;
 use std::os::unix::process::CommandExt;
 use std::process::{Command, ExitCode};
+use std::str::FromStr;
 use system_configuration::core_foundation::base::TCFType;
 use system_configuration::core_foundation::dictionary::CFDictionary;
 use system_configuration::core_foundation::number::CFNumber;
@@ -23,6 +25,67 @@ use system_configuration::preferences::SCPreferences;
 use system_configuration::sys::preferences::{SCPreferencesCommitChanges, SCPreferencesSetValue};
 use uzers::{get_current_groupname, get_current_username, get_effective_uid};
 
+#[derive(Debug, Clone)]
+enum AllowBlockEntry {
+    Net(Ipv4Net),
+    Domain(String),
+}
+
+impl FromStr for AllowBlockEntry {
+    type Err = Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let trimmed = s.trim();
+        if trimmed.is_empty() {
+            return Err(anyhow!("empty allow/block entry"));
+        }
+
+        if let Ok(net) = trimmed.parse::<Ipv4Net>() {
+            return Ok(AllowBlockEntry::Net(net));
+        }
+
+        if let Ok(addr) = trimmed.parse::<Ipv4Addr>() {
+            return Ok(AllowBlockEntry::Net(Ipv4Net::from(addr)));
+        }
+
+        Ok(AllowBlockEntry::Domain(trimmed.to_string()))
+    }
+}
+
+fn resolve_allow_block_entries(
+    kind: &str,
+    entries: Vec<AllowBlockEntry>,
+) -> anyhow::Result<Vec<Ipv4Net>> {
+    let mut nets = Vec::new();
+
+    for entry in entries {
+        match entry {
+            AllowBlockEntry::Net(net) => nets.push(net),
+            AllowBlockEntry::Domain(domain) => {
+                // A-record resolution happens here via ToSocketAddrs, then we keep only IPv4s.
+                let resolved: Vec<Ipv4Net> = (domain.as_str(), 0)
+                    .to_socket_addrs()
+                    .with_context(|| format!("failed to resolve {kind} entry {domain}"))?
+                    .filter_map(|addr| match addr {
+                        SocketAddr::V4(v4) => Some(Ipv4Net::from(*v4.ip())),
+                        SocketAddr::V6(_) => None,
+                    })
+                    .collect();
+
+                if resolved.is_empty() {
+                    return Err(anyhow!(
+                        "no IPv4 addresses found for {kind} entry {domain}"
+                    ));
+                }
+
+                nets.extend(resolved);
+            }
+        }
+    }
+
+    Ok(nets)
+}
+
 #[derive(Parser, Debug)]
 struct Args {
     #[clap(
@@ -52,12 +115,32 @@ struct Args {
 
     #[clap(
         long,
-        help = "comma-separated list of CIDRs to allow the traffic to (e.g. --allow=192.168.0.0/24)",
-        value_name = "comma-separated CIDRs",
+        help = "Comma-separated list of CIDRs, IPs, or domains to allow the traffic to \
+        (e.g. --allow=192.168.0.0/24 or --allow=example.com). Domains are resolved to A records \
+        at startup. \
+        When used with --block, the longest prefix match always wins. \
+        In case an identical prefix is both --allow'ed and --block'ed, \
+        blocking will take precedence. --allow=0.0.0.0/0 is a special case, \
+        it additionally disables bridge isolation (even when --block=0.0.0.0/0 is specified).",
+        value_name = "comma-separated CIDRs/IPs/domains",
         use_value_delimiter = true,
         action = clap::ArgAction::Set
     )]
-    allow: Vec<Ipv4Net>,
+    allow: Vec<AllowBlockEntry>,
+
+    #[clap(
+        long,
+        help = "Comma-separated list of CIDRs, IPs, or domains to block the traffic to \
+        (e.g. --block=0.0.0.0/0 may be used to establish a default deny policy \
+        that is further relaxed with --allow). Domains are resolved to A records at startup. \
+        When used with --allow, \
+        the longest prefix match always wins. In case the same prefix is both \
+        --allow'ed and --block'ed, blocking takes precedence.",
+        value_name = "comma-separated CIDRs/IPs/domains",
+        use_value_delimiter = true,
+        action = clap::ArgAction::Set
+    )]
+    block: Vec<AllowBlockEntry>,
 
     #[clap(
         long,
@@ -78,7 +161,9 @@ struct Args {
 fn main() -> ExitCode {
     // Enable backtraces by default
     if env::var("RUST_BACKTRACE").is_err() {
-        env::set_var("RUST_BACKTRACE", "full");
+        unsafe {
+            env::set_var("RUST_BACKTRACE", "full");
+        }
     }
 
     // Initialize Sentry
@@ -125,7 +210,7 @@ fn try_main() -> anyhow::Result<()> {
     // [2]: https://developer.apple.com/library/archive/documentation/System/Conceptual/ManPages_iPhoneOS/man2/kqueue.2.html
     unsafe { signal(Signal::SIGINT, SigHandler::SigIgn) }?;
 
-    let args: Args = Args::parse();
+    let mut args: Args = Args::parse();
 
     // No need to run anything, just return
     // so that the invoker process knows we
@@ -168,6 +253,9 @@ fn try_main() -> anyhow::Result<()> {
         ));
     }
 
+    let allow = resolve_allow_block_entries("allow", std::mem::take(&mut args.allow))?;
+    let block = resolve_allow_block_entries("block", std::mem::take(&mut args.block))?;
+
     // Set bootpd(8) min/max lease time while still having the root privileges
     set_bootpd_lease_time(args.bootpd_lease_time);
 
@@ -176,7 +264,8 @@ fn try_main() -> anyhow::Result<()> {
         args.vm_fd as RawFd,
         args.vm_mac_address,
         args.vm_net_type,
-        PrefixSet::from_iter(args.allow),
+        PrefixSet::from_iter(allow),
+        PrefixSet::from_iter(block),
         args.expose,
     )
     .context("failed to initialize proxy")?;
@@ -228,3 +317,24 @@ fn set_bootpd_lease_time(lease_time: u32) {
         SCPreferencesCommitChanges(prefs.as_concrete_TypeRef());
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{AllowBlockEntry, resolve_allow_block_entries};
+    use ipnet::Ipv4Net;
+    use std::net::Ipv4Addr;
+
+    #[test]
+    fn resolve_domain_to_ipv4_nets_dns_google() {
+        let nets = resolve_allow_block_entries(
+            "allow",
+            vec![AllowBlockEntry::Domain("dns.google".to_string())],
+        )
+        .unwrap();
+
+        assert!(
+            nets.contains(&Ipv4Net::from(Ipv4Addr::new(8, 8, 8, 8)))
+                || nets.contains(&Ipv4Net::from(Ipv4Addr::new(8, 8, 4, 4)))
+        );
+    }
+}