Allow rules to match TCP and UDP ports

lib/proxy/control.rs

@@ -106,7 +106,7 @@ fn parse_rules(rules: Vec<String>) -> std::result::Result<Vec<Rule>, ErrorObject
         let parsed_rule = rule.parse().map_err(|_| {
             rpc_error(
                 INVALID_PARAMS,
-                format!("invalid rule {rule:?}: expected TARGET, \"in TARGET\", or \"out TARGET\""),
+                format!("invalid rule {rule:?}: expected TARGET[:PORTS], \"in TARGET[:PORTS]\", or \"out TARGET[:PORTS]\""),
             )
         })?;
         parsed.push(parsed_rule);
@@ -645,13 +645,13 @@ mod tests {
         assert_eq!(
             policy
                 .rules
-                .policy_decision(Ipv4Address::new(10, 0, 0, 1), Direction::Out),
+                .policy_decision(Ipv4Address::new(10, 0, 0, 1), None, Direction::Out),
             Some(PolicyDecision::Block)
         );
         assert_eq!(
             policy
                 .rules
-                .policy_decision(Ipv4Address::new(192, 168, 64, 1), Direction::Out),
+                .policy_decision(Ipv4Address::new(192, 168, 64, 1), None, Direction::Out),
             Some(PolicyDecision::Block)
         );
     }
@@ -707,11 +707,11 @@ mod tests {
         assert_eq!(response["result"]["ruleCount"], 2);
         let target = Ipv4Address::new(10, 1, 2, 3);
         assert_eq!(
-            policy.rules.policy_decision(target, Direction::In),
+            policy.rules.policy_decision(target, None, Direction::In),
             Some(PolicyDecision::Block)
         );
         assert_eq!(
-            policy.rules.policy_decision(target, Direction::Out),
+            policy.rules.policy_decision(target, None, Direction::Out),
             Some(PolicyDecision::AllowStateful)
         );
     }

lib/proxy/host.rs

@@ -1,5 +1,5 @@
 use crate::proxy::flows::{FlowDirection, FlowMatch};
-use crate::proxy::{Direction, PolicyDecision, Proxy};
+use crate::proxy::{Direction, PolicyDecision, Proxy, transport_ports};
 use anyhow::{Context, Result};
 use dhcproto::v4::{CLIENT_PORT, SERVER_PORT};
 use smoltcp::phy::ChecksumCapabilities;
@@ -84,9 +84,11 @@ impl Proxy<'_> {
         };
 
         // The flow is either pending or untracked, evaluate it against inbound policy
+        let src_port = transport_ports(ipv4_pkt).map(|(src_port, _)| src_port);
+
         match self
             .rules
-            .policy_decision(ipv4_pkt.src_addr(), Direction::In)
+            .policy_decision(ipv4_pkt.src_addr(), src_port, Direction::In)
         {
             // Return traffic was handled above; enforce explicit inbound blocks here
             Some(PolicyDecision::Block) => None,
@@ -99,9 +101,12 @@ impl Proxy<'_> {
 
             // No inbound rule matched, so allow by default. Track the flow when needed
             // so its reply is not treated as a new outbound flow
-            None => {
-                self.admit_with_tracking_if_stateful(pending, ipv4_pkt.src_addr(), Direction::Out)
-            }
+            None => self.admit_with_tracking_if_stateful(
+                pending,
+                ipv4_pkt.src_addr(),
+                src_port,
+                Direction::Out,
+            ),
         }
     }
 }

lib/proxy/mod.rs

@@ -17,9 +17,9 @@ pub use exposed_port::ExposedPort;
 use flows::{FlowTable, PendingFlow};
 use ipnet::Ipv4Net;
 use mac_address::MacAddress;
-pub use rule::{Direction, Rule, Target};
+pub use rule::{Direction, Ports, Rule, Target};
 pub(crate) use rules::{PolicyDecision, Rules};
-use smoltcp::wire::{EthernetAddress, EthernetFrame, Ipv4Address};
+use smoltcp::wire::{EthernetAddress, EthernetFrame, IpProtocol, Ipv4Address, Ipv4Packet};
 use std::io::ErrorKind;
 use std::os::unix::io::{AsRawFd, RawFd};
 use std::time::Duration;
@@ -60,7 +60,7 @@ impl Proxy<'_> {
         let mut host = Host::new(
             vm_net_type,
             vm_mac_address,
-            !allow.contains(&Rule::Stateless(Target::Prefix(Ipv4Net::default()))),
+            !allow.contains(&Rule::Stateless(Target::Prefix(Ipv4Net::default()), None)),
         )?;
         let poller_timeout = Duration::from_millis(100);
         let control = control_fd
@@ -264,20 +264,42 @@ impl Proxy<'_> {
         &mut self,
         pending: Option<PendingFlow>,
         peer_addr: Ipv4Address,
+        peer_port: Option<u16>,
         return_direction: Direction,
     ) -> Option<()> {
-        if self.rules.is_stateful(peer_addr, return_direction) {
+        if self
+            .rules
+            .is_stateful(peer_addr, peer_port, return_direction)
+        {
             self.admit_with_tracking_if_trackable(pending)
         } else {
             Some(())
         }
     }
 }
 
+/// Returns the TCP or UDP source and destination ports, which only the first fragment carries.
+pub(crate) fn transport_ports(ipv4_pkt: &Ipv4Packet<&[u8]>) -> Option<(u16, u16)> {
+    if !matches!(ipv4_pkt.next_header(), IpProtocol::Tcp | IpProtocol::Udp)
+        || ipv4_pkt.frag_offset() != 0
+    {
+        return None;
+    }
+
+    // TCP and UDP both begin with the source port, then the destination port
+    match ipv4_pkt.payload() {
+        [src_hi, src_lo, dst_hi, dst_lo, ..] => Some((
+            u16::from_be_bytes([*src_hi, *src_lo]),
+            u16::from_be_bytes([*dst_hi, *dst_lo]),
+        )),
+        _ => None,
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::NetType;
-    use crate::proxy::Proxy;
+    use crate::proxy::{Proxy, transport_ports};
     use mac_address::MacAddress;
     use nix::sys::socket::{AddressFamily, SockFlag, SockType, socketpair};
     use serial_test::serial;
@@ -382,6 +404,31 @@ mod tests {
         assert!(proxy.allowed_from_host_ipv4(&reply).is_some());
     }
 
+    #[test]
+    fn test_transport_ports_only_in_first_fragment() {
+        let src = Ipv4Address::new(10, 0, 0, 2);
+        let dst = Ipv4Address::new(192, 168, 0, 10);
+        let mut bytes = udp_packet(src, 50000, dst, 8444);
+
+        assert_eq!(
+            transport_ports(&Ipv4Packet::new_unchecked(bytes.as_slice())),
+            Some((50000, 8444))
+        );
+
+        Ipv4Packet::new_unchecked(bytes.as_mut_slice()).set_frag_offset(8);
+        assert_eq!(
+            transport_ports(&Ipv4Packet::new_unchecked(bytes.as_slice())),
+            None
+        );
+
+        let mut bytes = udp_packet(src, 50000, dst, 8444);
+        Ipv4Packet::new_unchecked(bytes.as_mut_slice()).set_next_header(IpProtocol::Icmp);
+        assert_eq!(
+            transport_ports(&Ipv4Packet::new_unchecked(bytes.as_slice())),
+            None
+        );
+    }
+
     fn create_proxy<'test>(allow: Vec<&str>, block: Vec<&str>) -> Proxy<'test> {
         let (vm_fd, _) = socketpair(
             AddressFamily::Unix,