[WASI] Dynamic SocketAddr check for outbound socket traffic

crates/wasi/src/preview2/ctx.rs

@@ -1,18 +1,17 @@
-use super::clocks::host::{monotonic_clock, wall_clock};
 use crate::preview2::{
-    clocks::{self, HostMonotonicClock, HostWallClock},
+    clocks::{
+        host::{monotonic_clock, wall_clock},
+        HostMonotonicClock, HostWallClock,
+    },
     filesystem::Dir,
+    network::{SocketAddrCheck, SocketAddrUse},
     pipe, random, stdio,
     stdio::{StdinStream, StdoutStream},
     DirPerms, FilePerms,
 };
 use cap_rand::{Rng, RngCore, SeedableRng};
-use cap_std::ipnet::{self, IpNet};
-use cap_std::net::Pool;
-use cap_std::{ambient_authority, AmbientAuthority};
-use std::mem;
-use std::net::{Ipv4Addr, Ipv6Addr};
 use std::sync::Arc;
+use std::{mem, net::SocketAddr};
 use wasmtime::component::ResourceTable;
 
 pub struct WasiCtxBuilder {
@@ -22,8 +21,7 @@ pub struct WasiCtxBuilder {
     env: Vec<(String, String)>,
     args: Vec<String>,
     preopens: Vec<(Dir, String)>,
-
-    pool: Pool,
+    socket_addr_check: SocketAddrCheck,
     random: Box<dyn RngCore + Send + Sync>,
     insecure_random: Box<dyn RngCore + Send + Sync>,
     insecure_random_seed: u128,
@@ -73,7 +71,7 @@ impl WasiCtxBuilder {
             env: Vec::new(),
             args: Vec::new(),
             preopens: Vec::new(),
-            pool: Pool::new(),
+            socket_addr_check: SocketAddrCheck::default(),
             random: random::thread_rng(),
             insecure_random,
             insecure_random_seed,
@@ -173,80 +171,36 @@ impl WasiCtxBuilder {
         self.insecure_random = Box::new(insecure_random);
         self
     }
+
     pub fn insecure_random_seed(&mut self, insecure_random_seed: u128) -> &mut Self {
         self.insecure_random_seed = insecure_random_seed;
         self
     }
 
-    pub fn wall_clock(&mut self, clock: impl clocks::HostWallClock + 'static) -> &mut Self {
+    pub fn wall_clock(&mut self, clock: impl HostWallClock + 'static) -> &mut Self {
         self.wall_clock = Box::new(clock);
         self
     }
 
-    pub fn monotonic_clock(
-        &mut self,
-        clock: impl clocks::HostMonotonicClock + 'static,
-    ) -> &mut Self {
+    pub fn monotonic_clock(&mut self, clock: impl HostMonotonicClock + 'static) -> &mut Self {
         self.monotonic_clock = Box::new(clock);
         self
     }
 
-    /// Add all network addresses accessable to the host to the pool.
-    pub fn inherit_network(&mut self, ambient_authority: AmbientAuthority) -> &mut Self {
-        self.pool.insert_ip_net_port_any(
-            IpNet::new(Ipv4Addr::UNSPECIFIED.into(), 0).unwrap(),
-            ambient_authority,
-        );
-        self.pool.insert_ip_net_port_any(
-            IpNet::new(Ipv6Addr::UNSPECIFIED.into(), 0).unwrap(),
-            ambient_authority,
-        );
-        self
-    }
-
-    /// Add network addresses to the pool.
-    pub fn insert_addr<A: cap_std::net::ToSocketAddrs>(
-        &mut self,
-        addrs: A,
-    ) -> std::io::Result<&mut Self> {
-        self.pool.insert(addrs, ambient_authority())?;
-        Ok(self)
-    }
-
-    /// Add a specific [`cap_std::net::SocketAddr`] to the pool.
-    pub fn insert_socket_addr(&mut self, addr: cap_std::net::SocketAddr) -> &mut Self {
-        self.pool.insert_socket_addr(addr, ambient_authority());
-        self
-    }
-
-    /// Add a range of network addresses, accepting any port, to the pool.
-    ///
-    /// Unlike `insert_ip_net`, this function grants access to any requested port.
-    pub fn insert_ip_net_port_any(&mut self, ip_net: ipnet::IpNet) -> &mut Self {
-        self.pool
-            .insert_ip_net_port_any(ip_net, ambient_authority());
-        self
+    /// Allow all network addresses accessible to the host
+    pub fn inherit_network(&mut self) -> &mut Self {
+        self.socket_addr_check(|_, _| true)
     }
 
-    /// Add a range of network addresses, accepting a range of ports, to
-    /// per-instance networks.
+    /// A check that will be called for each socket address that is used.
     ///
-    /// This grants access to the port range starting at `ports_start` and, if
-    /// `ports_end` is provided, ending before `ports_end`.
-    pub fn insert_ip_net_port_range(
-        &mut self,
-        ip_net: ipnet::IpNet,
-        ports_start: u16,
-        ports_end: Option<u16>,
-    ) -> &mut Self {
-        self.pool
-            .insert_ip_net_port_range(ip_net, ports_start, ports_end, ambient_authority());
-        self
-    }
-
-    /// Add a range of network addresses with a specific port to the pool.
-    pub fn insert_ip_net(&mut self, ip_net: ipnet::IpNet, port: u16) -> &mut Self {
-        self.pool.insert_ip_net(ip_net, port, ambient_authority());
+    /// Returning `true` will permit socket connections to the `SocketAddr`,
+    /// while returning `false` will reject the connection.
+    pub fn socket_addr_check<F>(&mut self, check: F) -> &mut Self
+    where
+        F: Fn(&SocketAddr, SocketAddrUse) -> bool + Send + Sync + 'static,
+    {
+        self.socket_addr_check = SocketAddrCheck(Arc::new(check));
         self
     }
 
@@ -286,7 +240,7 @@ impl WasiCtxBuilder {
             env,
             args,
             preopens,
-            pool,
+            socket_addr_check,
             random,
             insecure_random,
             insecure_random_seed,
@@ -304,7 +258,7 @@ impl WasiCtxBuilder {
             env,
             args,
             preopens,
-            pool: Arc::new(pool),
+            socket_addr_check,
             random,
             insecure_random,
             insecure_random_seed,
@@ -334,7 +288,7 @@ pub struct WasiCtx {
     pub(crate) stdin: Box<dyn StdinStream>,
     pub(crate) stdout: Box<dyn StdoutStream>,
     pub(crate) stderr: Box<dyn StdoutStream>,
-    pub(crate) pool: Arc<Pool>,
+    pub(crate) socket_addr_check: SocketAddrCheck,
     pub(crate) allowed_network_uses: AllowedNetworkUses,
 }
 
@@ -360,8 +314,7 @@ impl AllowedNetworkUses {
             return Err(std::io::Error::new(
                 std::io::ErrorKind::PermissionDenied,
                 "UDP is not allowed",
-            )
-            .into());
+            ));
         }
 
         Ok(())
@@ -372,8 +325,7 @@ impl AllowedNetworkUses {
             return Err(std::io::Error::new(
                 std::io::ErrorKind::PermissionDenied,
                 "TCP is not allowed",
-            )
-            .into());
+            ));
         }
 
         Ok(())

crates/wasi/src/preview2/host/instance_network.rs

@@ -6,7 +6,7 @@ use wasmtime::component::Resource;
 impl<T: WasiView> instance_network::Host for T {
     fn instance_network(&mut self) -> Result<Resource<Network>, anyhow::Error> {
         let network = Network {
-            pool: self.ctx().pool.clone(),
+            socket_addr_check: self.ctx().socket_addr_check.clone(),
             allow_ip_name_lookup: self.ctx().allowed_network_uses.ip_name_lookup,
         };
         let network = self.table_mut().push(network)?;

crates/wasi/src/preview2/host/network.rs

@@ -218,7 +218,7 @@ pub(crate) mod util {
     use crate::preview2::bindings::sockets::network::ErrorCode;
     use crate::preview2::network::SocketAddressFamily;
     use crate::preview2::SocketResult;
-    use cap_net_ext::{Blocking, TcpBinder, TcpConnecter, TcpListenerExt, UdpBinder};
+    use cap_net_ext::{Blocking, TcpListenerExt};
     use cap_std::net::{TcpListener, TcpStream, UdpSocket};
     use rustix::fd::AsFd;
     use rustix::io::Errno;
@@ -302,42 +302,38 @@ pub(crate) mod util {
      * Syscalls wrappers with (opinionated) portability fixes.
      */
 
-    pub fn tcp_bind(listener: &TcpListener, binder: &TcpBinder) -> std::io::Result<()> {
-        binder
-            .bind_existing_tcp_listener(listener)
-            .map_err(|error| match Errno::from_io_error(&error) {
-                // See: https://learn.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-bind#:~:text=WSAENOBUFS
-                // Windows returns WSAENOBUFS when the ephemeral ports have been exhausted.
-                #[cfg(windows)]
-                Some(Errno::NOBUFS) => Errno::ADDRINUSE.into(),
-                _ => error,
-            })
+    pub fn tcp_bind(listener: &TcpListener, addr: &SocketAddr) -> std::io::Result<()> {
+        rustix::net::bind(listener, addr).map_err(|error| match error {
+            // See: https://learn.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-bind#:~:text=WSAENOBUFS
+            // Windows returns WSAENOBUFS when the ephemeral ports have been exhausted.
+            #[cfg(windows)]
+            Errno::NOBUFS => Errno::ADDRINUSE.into(),
+            _ => error.into(),
+        })
     }
 
-    pub fn udp_bind(socket: &UdpSocket, binder: &UdpBinder) -> std::io::Result<()> {
-        binder.bind_existing_udp_socket(socket).map_err(|error| {
-            match Errno::from_io_error(&error) {
+    pub fn udp_bind(socket: &UdpSocket, addr: &SocketAddr) -> std::io::Result<()> {
+        rustix::net::bind(socket, addr).map_err(|error| {
+            match error {
                 // See: https://learn.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-bind#:~:text=WSAENOBUFS
                 // Windows returns WSAENOBUFS when the ephemeral ports have been exhausted.
                 #[cfg(windows)]
-                Some(Errno::NOBUFS) => Errno::ADDRINUSE.into(),
-                _ => error,
+                Errno::NOBUFS => Errno::ADDRINUSE.into(),
+                _ => error.into(),
             }
         })
     }
 
-    pub fn tcp_connect(listener: &TcpListener, connecter: &TcpConnecter) -> std::io::Result<()> {
-        connecter.connect_existing_tcp_listener(listener).map_err(
-            |error| match Errno::from_io_error(&error) {
-                // On POSIX, non-blocking `connect` returns `EINPROGRESS`.
-                // Windows returns `WSAEWOULDBLOCK`.
-                //
-                // This normalized error code is depended upon by: tcp.rs
-                #[cfg(windows)]
-                Some(Errno::WOULDBLOCK) => Errno::INPROGRESS.into(),
-                _ => error,
-            },
-        )
+    pub fn tcp_connect(listener: &TcpListener, addr: &SocketAddr) -> std::io::Result<()> {
+        rustix::net::connect(listener, addr).map_err(|error| match error {
+            // On POSIX, non-blocking `connect` returns `EINPROGRESS`.
+            // Windows returns `WSAEWOULDBLOCK`.
+            //
+            // This normalized error code is depended upon by: tcp.rs
+            #[cfg(windows)]
+            Errno::WOULDBLOCK => Errno::INPROGRESS.into(),
+            _ => error.into(),
+        })
     }
 
     pub fn tcp_listen(listener: &TcpListener, backlog: Option<i32>) -> std::io::Result<()> {

crates/wasi/src/preview2/host/tcp.rs

@@ -1,4 +1,5 @@
 use crate::preview2::host::network::util;
+use crate::preview2::network::SocketAddrUse;
 use crate::preview2::tcp::{TcpSocket, TcpState};
 use crate::preview2::{
     bindings::{
@@ -9,7 +10,7 @@ use crate::preview2::{
     network::SocketAddressFamily,
 };
 use crate::preview2::{Pollable, SocketResult, WasiView};
-use cap_net_ext::{Blocking, PoolExt};
+use cap_net_ext::Blocking;
 use cap_std::net::TcpListener;
 use io_lifetimes::AsSocketlike;
 use rustix::io::Errno;
@@ -44,11 +45,12 @@ impl<T: WasiView> crate::preview2::host::tcp::tcp::HostTcpSocket for T {
         util::validate_address_family(&local_address, &socket.family)?;
 
         {
-            let binder = network.pool.tcp_binder(local_address)?;
+            // Ensure that we're allowed to connect to this address.
+            network.check_socket_addr(&local_address, SocketAddrUse::TcpBind)?;
             let listener = &*socket.tcp_socket().as_socketlike_view::<TcpListener>();
 
             // Perform the OS bind call.
-            util::tcp_bind(listener, &binder).map_err(|error| {
+            util::tcp_bind(listener, &local_address).map_err(|error| {
                 match Errno::from_io_error(&error) {
                     // From https://pubs.opengroup.org/onlinepubs/9699919799/functions/bind.html:
                     // > [EAFNOSUPPORT] The specified address is not a valid address for the address family of the specified socket
@@ -112,11 +114,12 @@ impl<T: WasiView> crate::preview2::host::tcp::tcp::HostTcpSocket for T {
             util::validate_remote_address(&remote_address)?;
             util::validate_address_family(&remote_address, &socket.family)?;
 
-            let connecter = network.pool.tcp_connecter(remote_address)?;
+            // Ensure that we're allowed to connect to this address.
+            network.check_socket_addr(&remote_address, SocketAddrUse::TcpConnect)?;
             let listener = &*socket.tcp_socket().as_socketlike_view::<TcpListener>();
 
             // Do an OS `connect`. Our socket is non-blocking, so it'll either...
-            util::tcp_connect(listener, &connecter)
+            util::tcp_connect(listener, &remote_address)
         };
 
         match r {