Repository: bytecodealliance/wasmtime
Pull request #7648 (2023-12-06 17:14:40)

Ensure `remote-address` is allowed when creating UDP stream
Until now, the `remote-address` argument to `wasi:sockets/udp-socket#stream` was not checked against the `Pool` of allowed IP addresses. Now, we store the `Pool` from `network` into the `socket` object in `wasi:sockets/udp-socket#bind` and later use that `Pool` in `wasi:sockets/udp-socket#stream` to check that `remote-address` is allowed.

The check itself is a bit hacky since `Pool` doesn't seem to expose a way to just check an address so we use `PoolExt` to create a `UdpConnecter` which we will fail if we don't have permission to connect to the remote address.

I can add a test once I figure out the testing story in https://github.com/bytecodealliance/wasmtime/pull/7647

rylev: I've addressed the PR feedback:
* Pool is now wrapped in an `Arc` to make cloning cheaper.
* `HostOutgoingDatagramStream` will check the provided addr against the `Pool` before sending a datagram.

@alexcrichton I'm not sure what the testing story should look like here. I don't see where the wasi implementations are being tested against non-default `WasiCtx`s. I believe all tests are running with full network. I suppose whatever solution we decide on here should also work for #7647 as well.
alexcrichton: I'm gonna go ahead and queue this for merge, and @rylev says they're happy to follow-up on semantic changes in a follow-up.

Patch:
diff --git a/crates/wasi/src/preview2/ctx.rs b/crates/wasi/src/preview2/ctx.rs
index 5d527b241..d3d52408d 100644
--- a/crates/wasi/src/preview2/ctx.rs
+++ b/crates/wasi/src/preview2/ctx.rs
@@ -12,6 +12,7 @@ use cap_std::net::Pool;
 use cap_std::{ambient_authority, AmbientAuthority};
 use std::mem;
 use std::net::{Ipv4Addr, Ipv6Addr};
+use std::sync::Arc;
 use wasmtime::component::ResourceTable;
 
 pub struct WasiCtxBuilder {
@@ -303,7 +304,7 @@ impl WasiCtxBuilder {
             env,
             args,
             preopens,
-            pool,
+            pool: Arc::new(pool),
             random,
             insecure_random,
             insecure_random_seed,
@@ -333,7 +334,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: Pool,
+    pub(crate) pool: Arc<Pool>,
     pub(crate) allowed_network_uses: AllowedNetworkUses,
 }
 
diff --git a/crates/wasi/src/preview2/host/udp.rs b/crates/wasi/src/preview2/host/udp.rs
index 9b5955511..cd6dd62c3 100644
--- a/crates/wasi/src/preview2/host/udp.rs
+++ b/crates/wasi/src/preview2/host/udp.rs
@@ -35,20 +35,24 @@ impl<T: WasiView> udp::HostUdpSocket for T {
     ) -> SocketResult<()> {
         self.ctx().allowed_network_uses.check_allowed_udp()?;
         let table = self.table_mut();
-        let socket = table.get(&this)?;
-        let network = table.get(&network)?;
-        let local_address: SocketAddr = local_address.into();
 
-        match socket.udp_state {
+        match table.get(&this)?.udp_state {
             UdpState::Default => {}
             UdpState::BindStarted => return Err(ErrorCode::ConcurrencyConflict.into()),
             UdpState::Bound | UdpState::Connected => return Err(ErrorCode::InvalidState.into()),
         }
 
+        // Set the pool on the socket so later functions have access to it through the socket handle
+        let pool = table.get(&network)?.pool.clone();
+        table.get_mut(&this)?.pool.replace(pool.clone());
+
+        let socket = table.get(&this)?;
+        let local_address: SocketAddr = local_address.into();
+
         util::validate_address_family(&local_address, &socket.family)?;
 
         {
-            let binder = network.pool.udp_binder(local_address)?;
+            let binder = pool.udp_binder(local_address)?;
             let udp_socket = &*socket
                 .udp_socket()
                 .as_socketlike_view::<cap_std::net::UdpSocket>();
@@ -128,8 +132,13 @@ impl<T: WasiView> udp::HostUdpSocket for T {
 
         // Step #2: (Re)connect
         if let Some(connect_addr) = remote_address {
+            let Some(pool) = socket.pool.as_ref() else {
+                return Err(ErrorCode::InvalidState.into());
+            };
             util::validate_remote_address(&connect_addr)?;
             util::validate_address_family(&connect_addr, &socket.family)?;
+            // We don't actually use the connecter, we just use it to verify that `connect_addr` is allowed
+            let _ = pool.udp_connecter(connect_addr)?;
 
             rustix::net::connect(socket.udp_socket(), &connect_addr).map_err(
                 |error| match error {
@@ -153,6 +162,7 @@ impl<T: WasiView> udp::HostUdpSocket for T {
             remote_address,
             family: socket.family,
             send_state: SendState::Idle,
+            pool: socket.pool.clone(),
         };
 
         Ok((
@@ -345,7 +355,6 @@ impl<T: WasiView> udp::HostIncomingDatagramStream for T {
                 _ => {}
             }
 
-            // FIXME: check permission to receive from `received_addr`.
             Ok(Some(udp::IncomingDatagram {
                 data: buf[..size].into(),
                 remote_address: received_addr.into(),
@@ -448,7 +457,16 @@ impl<T: WasiView> udp::HostOutgoingDatagramStream for T {
 
             let provided_addr = datagram.remote_address.map(SocketAddr::from);
             let addr = match (stream.remote_address, provided_addr) {
-                (None, Some(addr)) => addr,
+                (None, Some(addr)) => {
+                    let Some(pool) = stream.pool.as_ref() else {
+                        return Err(ErrorCode::InvalidState.into());
+                    };
+                    // We don't actually use the connecter, we just use it to verify that `addr`
+                    // is allowed. We only need to check the provided addr as the stream's remote
+                    // address was checked when the stream was created.
+                    let _ = pool.udp_connecter(addr)?;
+                    addr
+                }
                 (Some(addr), None) => addr,
                 (Some(connected_addr), Some(provided_addr)) if connected_addr == provided_addr => {
                     connected_addr
@@ -459,7 +477,6 @@ impl<T: WasiView> udp::HostOutgoingDatagramStream for T {
             util::validate_remote_address(&addr)?;
             util::validate_address_family(&addr, &stream.family)?;
 
-            // FIXME: check permission to send to `addr`.
             if stream.remote_address == Some(addr) {
                 stream.inner.try_send(&datagram.data)?;
             } else {
diff --git a/crates/wasi/src/preview2/network.rs b/crates/wasi/src/preview2/network.rs
index 13575202b..6384b1803 100644
--- a/crates/wasi/src/preview2/network.rs
+++ b/crates/wasi/src/preview2/network.rs
@@ -2,9 +2,10 @@ use crate::preview2::bindings::sockets::network::{Ipv4Address, Ipv6Address};
 use crate::preview2::bindings::wasi::sockets::network::ErrorCode;
 use crate::preview2::TrappableError;
 use cap_std::net::Pool;
+use std::sync::Arc;
 
 pub struct Network {
-    pub pool: Pool,
+    pub pool: Arc<Pool>,
     pub allow_ip_name_lookup: bool,
 }
 
diff --git a/crates/wasi/src/preview2/udp.rs b/crates/wasi/src/preview2/udp.rs
index 16335d601..13194c1f5 100644
--- a/crates/wasi/src/preview2/udp.rs
+++ b/crates/wasi/src/preview2/udp.rs
@@ -2,6 +2,7 @@ use crate::preview2::poll::Subscribe;
 use crate::preview2::with_ambient_tokio_runtime;
 use async_trait::async_trait;
 use cap_net_ext::{AddressFamily, Blocking, UdpSocketExt};
+use cap_std::net::Pool;
 use io_lifetimes::raw::{FromRawSocketlike, IntoRawSocketlike};
 use std::io;
 use std::net::SocketAddr;
@@ -42,6 +43,9 @@ pub struct UdpSocket {
 
     /// Socket address family.
     pub(crate) family: SocketAddressFamily,
+
+    /// The pool of allowed addresses
+    pub(crate) pool: Option<Arc<Pool>>,
 }
 
 #[async_trait]
@@ -67,6 +71,7 @@ impl UdpSocket {
             inner: Arc::new(socket),
             udp_state: UdpState::Default,
             family: socket_address_family,
+            pool: None,
         })
     }
 
@@ -104,6 +109,9 @@ pub struct OutgoingDatagramStream {
     pub(crate) family: SocketAddressFamily,
 
     pub(crate) send_state: SendState,
+
+    /// The pool of allowed addresses
+    pub(crate) pool: Option<Arc<Pool>>,
 }
 
 pub(crate) enum SendState {
diff --git a/src/commands/run.rs b/src/commands/run.rs
index b5b465535..c16f2d3c6 100644
--- a/src/commands/run.rs
+++ b/src/commands/run.rs
@@ -782,7 +782,7 @@ impl RunCommand {
             builder.allow_tcp(enable);
         }
         if let Some(enable) = self.run.common.wasi.udp {
-            builder.allow_tcp(enable);
+            builder.allow_udp(enable);
         }
 
         store.data_mut().preview2_ctx = Some(Arc::new(builder.build()));


Test patch:
diff --git a/crates/test-programs/src/bin/cli_no_udp.rs b/crates/test-programs/src/bin/cli_no_udp.rs
new file mode 100644
index 000000000..1dd024061
--- /dev/null
+++ b/crates/test-programs/src/bin/cli_no_udp.rs
@@ -0,0 +1,16 @@
+//! This test assumes that it will be run without udp support enabled
+use test_programs::wasi::sockets::{
+    network::IpAddress,
+    udp::{ErrorCode, IpAddressFamily, IpSocketAddress, Network, UdpSocket},
+};
+
+fn main() {
+    let net = Network::default();
+    let family = IpAddressFamily::Ipv4;
+    let remote1 = IpSocketAddress::new(IpAddress::new_loopback(family), 4321);
+    let sock = UdpSocket::new(family).unwrap();
+
+    let bind = sock.blocking_bind(&net, remote1);
+    eprintln!("Result of binding: {bind:?}");
+    assert!(matches!(bind, Err(ErrorCode::AccessDenied)));
+}
diff --git a/tests/all/cli_tests.rs b/tests/all/cli_tests.rs
index 7fb807eb1..527572cba 100644
--- a/tests/all/cli_tests.rs
+++ b/tests/all/cli_tests.rs
@@ -1514,4 +1514,20 @@ mod test_programs {
         assert!(output.status.success());
         Ok(())
     }
+
+    #[test]
+    fn cli_no_udp() -> Result<()> {
+        let output = super::run_wasmtime_for_output(
+            &[
+                "-Wcomponent-model",
+                // Turn on network but turn off UDP
+                "-Sinherit-network,udp=no",
+                CLI_NO_UDP_COMPONENT,
+            ],
+            None,
+        )?;
+        println!("{}", String::from_utf8_lossy(&output.stderr));
+        assert!(output.status.success());
+        Ok(())
+    }
 }