[wasi-sockets] disallow sending unpermitted datagrams (#12641)

In #12629, I removed the enforcement of `outgoing-datagram-stream.check-send`
permits, which turned out to be incorrect according to [the spec].  This
restores the enforcement to conform to the spec while still fixing the issue
addressed by the original PR.

[the spec]: https://github.com/WebAssembly/WASI/blob/41d0de898ec7568e032129156a3bc567d7ebea5d/proposals/sockets/wit/udp.wit#L256C1-L257C111

Author: dicej

crates/test-programs/src/bin/p2_udp_send_too_much.rs

@@ -0,0 +1,34 @@
+use test_programs::wasi::sockets::network::{IpAddress, IpAddressFamily, IpSocketAddress, Network};
+use test_programs::wasi::sockets::udp::{OutgoingDatagram, UdpSocket};
+
+/// `outgoing-datagram-stream.send` should trap we attempt to send more
+/// datagrams than `check-send` gave permission for.
+fn main() {
+    let net = Network::default();
+    let family = IpAddressFamily::Ipv4;
+    let unspecified_port = IpSocketAddress::new(IpAddress::new_loopback(family), 0);
+    let remote = IpSocketAddress::new(IpAddress::new_loopback(family), 4320);
+
+    let client = UdpSocket::new(family).unwrap();
+    client.blocking_bind(&net, unspecified_port).unwrap();
+
+    let (_, tx) = client.stream(Some(remote)).unwrap();
+    assert_eq!(client.remote_address(), Ok(remote));
+
+    tx.subscribe().block();
+
+    let permits = tx.check_send().unwrap();
+    assert!(permits > 0);
+    // This should trap according to the `wasi-sockets` spec since we're trying
+    // to send more than `check_send` gave permission for:
+    tx.send(
+        &std::iter::repeat_with(|| OutgoingDatagram {
+            data: b"hello".into(),
+            remote_address: None,
+        })
+        .take(usize::try_from(permits + 1).unwrap())
+        .collect::<Vec<_>>(),
+    )
+    .unwrap();
+    unreachable!("attempt to send excess datagrams should have trapped");
+}

crates/wasi/src/p2/host/udp.rs

@@ -85,6 +85,7 @@ impl udp::HostUdpSocket for WasiSocketsCtxView<'_> {
             remote_address,
             family: socket.address_family(),
             socket_addr_check: socket.socket_addr_check().cloned(),
+            check_send_permit_count: 0,
         };
 
         Ok((
@@ -264,20 +265,22 @@ impl udp::HostOutgoingDatagramStream for WasiSocketsCtxView<'_> {
     fn check_send(&mut self, this: Resource<udp::OutgoingDatagramStream>) -> SocketResult<u64> {
         let stream = self.table.get_mut(&this)?;
 
-        Ok(
-            if let Poll::Ready(_) =
-                pin!(stream.inner.ready(Interest::WRITABLE.add(Interest::ERROR)))
-                    .poll(&mut Context::from_waker(Waker::noop()))
-            {
-                // We don't know how many Tokio will accept, so we make up a
-                // reasonable number here.  If we're wrong and `send` returns
-                // `Ok(0)`, the guest will just have to deal with that, e.g. by
-                // looping or returning `EWOULDBLOCK`.
-                16
-            } else {
-                0
-            },
-        )
+        let count = if let Poll::Ready(_) =
+            pin!(stream.inner.ready(Interest::WRITABLE.add(Interest::ERROR)))
+                .poll(&mut Context::from_waker(Waker::noop()))
+        {
+            // We don't know how many Tokio will accept, so we make up a
+            // reasonable number here.  If we're wrong and `send` returns
+            // `Ok(0)`, the guest will just have to deal with that, e.g. by
+            // looping or returning `EWOULDBLOCK`.
+            16
+        } else {
+            0
+        };
+
+        stream.check_send_permit_count = count;
+
+        Ok(count.try_into().unwrap())
     }
 
     async fn send(
@@ -331,6 +334,14 @@ impl udp::HostOutgoingDatagramStream for WasiSocketsCtxView<'_> {
             return Ok(0);
         }
 
+        if datagrams.len() > stream.check_send_permit_count {
+            return Err(SocketError::trap(wasmtime::format_err!(
+                "unpermitted: argument exceeds permitted size"
+            )));
+        }
+
+        stream.check_send_permit_count -= datagrams.len();
+
         let mut count = 0;
 
         for datagram in datagrams {

crates/wasi/src/p2/udp.rs

@@ -20,4 +20,8 @@ pub struct OutgoingDatagramStream {
 
     /// The check of allowed addresses
     pub(crate) socket_addr_check: Option<SocketAddrCheck>,
+
+    /// Remaining number of datagrams permitted by most recent `check-send`
+    /// call.
+    pub(crate) check_send_permit_count: usize,
 }

crates/wasi/tests/all/p2/async_.rs

@@ -361,6 +361,16 @@ async fn p2_adapter_badfd() {
 async fn p2_file_read_write() {
     run(P2_FILE_READ_WRITE_COMPONENT, false).await.unwrap()
 }
+#[test_log::test(tokio::test(flavor = "multi_thread"))]
+async fn p2_udp_send_too_much() {
+    let e = run(P2_UDP_SEND_TOO_MUCH_COMPONENT, false)
+        .await
+        .unwrap_err();
+    assert_eq!(
+        format!("{}", e.source().expect("trap source")),
+        "unpermitted: argument exceeds permitted size"
+    )
+}
 #[expect(
     dead_code,
     reason = "tested in the wasi-cli crate, satisfying foreach_api! macro"

crates/wasi/tests/all/p2/sync.rs

@@ -343,6 +343,14 @@ fn p2_adapter_badfd() {
 fn p2_file_read_write() {
     run(P2_FILE_READ_WRITE_COMPONENT, false).unwrap()
 }
+#[test_log::test]
+fn p2_udp_send_too_much() {
+    let e = run(P2_UDP_SEND_TOO_MUCH_COMPONENT, false).unwrap_err();
+    assert_eq!(
+        format!("{}", e.source().expect("trap source")),
+        "unpermitted: argument exceeds permitted size"
+    )
+}
 #[expect(
     dead_code,
     reason = "tested in the wasi-cli crate, satisfying foreach_api! macro"