Relax required host capabilities in wasip3 tcp/udp bindings (#12085)

Don't require `async | store` when it's not necessary as this'll soon
have ramifications on the type of the function being bound in the
component model.

Author: alexcrichton

crates/wasi/src/p3/bindings.rs

@@ -83,11 +83,11 @@ mod generated {
             "wasi:cli/stdout": store | tracing | trappable,
             "wasi:cli/stderr": store | tracing | trappable,
             "wasi:filesystem/types.[method]descriptor.read-via-stream": store | tracing | trappable,
-            "wasi:sockets/types.[method]tcp-socket.bind": async | store | tracing | trappable,
-            "wasi:sockets/types.[method]tcp-socket.listen": async | store | tracing | trappable,
-            "wasi:sockets/types.[method]tcp-socket.receive": async | store | tracing | trappable,
-            "wasi:sockets/types.[method]udp-socket.bind": async | store | tracing | trappable,
-            "wasi:sockets/types.[method]udp-socket.connect": async | store | tracing | trappable,
+            "wasi:sockets/types.[method]tcp-socket.bind": async | tracing | trappable,
+            "wasi:sockets/types.[method]tcp-socket.listen":  store | tracing | trappable,
+            "wasi:sockets/types.[method]tcp-socket.receive": store | tracing | trappable,
+            "wasi:sockets/types.[method]udp-socket.bind": async | tracing | trappable,
+            "wasi:sockets/types.[method]udp-socket.connect": async | tracing | trappable,
             default: tracing | trappable,
         },
         exports: { default: async | store },

crates/wasi/src/p3/sockets/host/types/tcp.rs

@@ -18,7 +18,7 @@ use std::sync::Arc;
 use tokio::net::{TcpListener, TcpStream};
 use tokio::sync::oneshot;
 use wasmtime::component::{
-    Accessor, Destination, FutureReader, Resource, ResourceTable, Source, StreamConsumer,
+    Access, Accessor, Destination, FutureReader, Resource, ResourceTable, Source, StreamConsumer,
     StreamProducer, StreamReader, StreamResult,
 };
 use wasmtime::{AsContextMut as _, StoreContextMut};
@@ -236,23 +236,6 @@ impl<D> StreamConsumer<D> for SendStreamConsumer {
 }
 
 impl HostTcpSocketWithStore for WasiSockets {
-    async fn bind<T>(
-        store: &Accessor<T, Self>,
-        socket: Resource<TcpSocket>,
-        local_address: IpSocketAddress,
-    ) -> SocketResult<()> {
-        let local_address = SocketAddr::from(local_address);
-        if !is_addr_allowed(store, local_address, SocketAddrUse::TcpBind).await {
-            return Err(ErrorCode::AccessDenied.into());
-        }
-        store.with(|mut store| {
-            let socket = get_socket_mut(store.get().table, &socket)?;
-            socket.start_bind(local_address)?;
-            socket.finish_bind()?;
-            Ok(())
-        })
-    }
-
     async fn connect<T>(
         store: &Accessor<T, Self>,
         socket: Resource<TcpSocket>,
@@ -278,28 +261,26 @@ impl HostTcpSocketWithStore for WasiSockets {
         })
     }
 
-    async fn listen<T: 'static>(
-        store: &Accessor<T, Self>,
+    fn listen<T: 'static>(
+        mut store: Access<'_, T, Self>,
         socket: Resource<TcpSocket>,
     ) -> SocketResult<StreamReader<Resource<TcpSocket>>> {
         let getter = store.getter();
-        store.with(|mut store| {
-            let socket = get_socket_mut(store.get().table, &socket)?;
-            socket.start_listen()?;
-            socket.finish_listen()?;
-            let listener = socket.tcp_listener_arc().unwrap().clone();
-            let family = socket.address_family();
-            let options = socket.non_inherited_options().clone();
-            Ok(StreamReader::new(
-                &mut store,
-                ListenStreamProducer {
-                    listener,
-                    family,
-                    options,
-                    getter,
-                },
-            ))
-        })
+        let socket = get_socket_mut(store.get().table, &socket)?;
+        socket.start_listen()?;
+        socket.finish_listen()?;
+        let listener = socket.tcp_listener_arc().unwrap().clone();
+        let family = socket.address_family();
+        let options = socket.non_inherited_options().clone();
+        Ok(StreamReader::new(
+            &mut store,
+            ListenStreamProducer {
+                listener,
+                family,
+                options,
+                getter,
+            },
+        ))
     }
 
     async fn send<T: 'static>(
@@ -328,39 +309,52 @@ impl HostTcpSocketWithStore for WasiSockets {
         Ok(())
     }
 
-    async fn receive<T: 'static>(
-        store: &Accessor<T, Self>,
+    fn receive<T: 'static>(
+        mut store: Access<T, Self>,
         socket: Resource<TcpSocket>,
     ) -> wasmtime::Result<(StreamReader<u8>, FutureReader<Result<(), ErrorCode>>)> {
-        store.with(|mut store| {
-            let socket = get_socket_mut(store.get().table, &socket)?;
-            match socket.start_receive() {
-                Some(stream) => {
-                    let stream = Arc::clone(stream);
-                    let (result_tx, result_rx) = oneshot::channel();
-                    Ok((
-                        StreamReader::new(
-                            &mut store,
-                            ReceiveStreamProducer {
-                                stream,
-                                result: Some(result_tx),
-                            },
-                        ),
-                        FutureReader::new(&mut store, result_rx),
-                    ))
-                }
-                None => Ok((
-                    StreamReader::new(&mut store, iter::empty()),
-                    FutureReader::new(&mut store, async {
-                        anyhow::Ok(Err(ErrorCode::InvalidState))
-                    }),
-                )),
+        let socket = get_socket_mut(store.get().table, &socket)?;
+        match socket.start_receive() {
+            Some(stream) => {
+                let stream = Arc::clone(stream);
+                let (result_tx, result_rx) = oneshot::channel();
+                Ok((
+                    StreamReader::new(
+                        &mut store,
+                        ReceiveStreamProducer {
+                            stream,
+                            result: Some(result_tx),
+                        },
+                    ),
+                    FutureReader::new(&mut store, result_rx),
+                ))
             }
-        })
+            None => Ok((
+                StreamReader::new(&mut store, iter::empty()),
+                FutureReader::new(&mut store, async {
+                    anyhow::Ok(Err(ErrorCode::InvalidState))
+                }),
+            )),
+        }
     }
 }
 
 impl HostTcpSocket for WasiSocketsCtxView<'_> {
+    async fn bind(
+        &mut self,
+        socket: Resource<TcpSocket>,
+        local_address: IpSocketAddress,
+    ) -> SocketResult<()> {
+        let local_address = SocketAddr::from(local_address);
+        if !(self.ctx.socket_addr_check)(local_address, SocketAddrUse::TcpBind).await {
+            return Err(ErrorCode::AccessDenied.into());
+        }
+        let socket = get_socket_mut(self.table, &socket)?;
+        socket.start_bind(local_address)?;
+        socket.finish_bind()?;
+        Ok(())
+    }
+
     fn create(&mut self, address_family: IpAddressFamily) -> SocketResult<Resource<TcpSocket>> {
         let family = address_family.into();
         let socket = TcpSocket::new(self.ctx, family)?;

crates/wasi/src/p3/sockets/host/types/udp.rs

@@ -30,39 +30,6 @@ fn get_socket_mut<'a>(
 }
 
 impl HostUdpSocketWithStore for WasiSockets {
-    async fn bind<T>(
-        store: &Accessor<T, Self>,
-        socket: Resource<UdpSocket>,
-        local_address: IpSocketAddress,
-    ) -> SocketResult<()> {
-        let local_address = SocketAddr::from(local_address);
-        if !is_addr_allowed(store, local_address, SocketAddrUse::UdpBind).await {
-            return Err(ErrorCode::AccessDenied.into());
-        }
-        store.with(|mut view| {
-            let socket = get_socket_mut(view.get().table, &socket)?;
-            socket.bind(local_address)?;
-            socket.finish_bind()?;
-            Ok(())
-        })
-    }
-
-    async fn connect<T>(
-        store: &Accessor<T, Self>,
-        socket: Resource<UdpSocket>,
-        remote_address: IpSocketAddress,
-    ) -> SocketResult<()> {
-        let remote_address = SocketAddr::from(remote_address);
-        if !is_addr_allowed(store, remote_address, SocketAddrUse::UdpConnect).await {
-            return Err(ErrorCode::AccessDenied.into());
-        }
-        store.with(|mut view| {
-            let socket = get_socket_mut(view.get().table, &socket)?;
-            socket.connect(remote_address)?;
-            Ok(())
-        })
-    }
-
     async fn send<T>(
         store: &Accessor<T, Self>,
         socket: Resource<UdpSocket>,
@@ -103,6 +70,35 @@ impl HostUdpSocketWithStore for WasiSockets {
 }
 
 impl HostUdpSocket for WasiSocketsCtxView<'_> {
+    async fn bind(
+        &mut self,
+        socket: Resource<UdpSocket>,
+        local_address: IpSocketAddress,
+    ) -> SocketResult<()> {
+        let local_address = SocketAddr::from(local_address);
+        if !(self.ctx.socket_addr_check)(local_address, SocketAddrUse::UdpBind).await {
+            return Err(ErrorCode::AccessDenied.into());
+        }
+        let socket = get_socket_mut(self.table, &socket)?;
+        socket.bind(local_address)?;
+        socket.finish_bind()?;
+        Ok(())
+    }
+
+    async fn connect(
+        &mut self,
+        socket: Resource<UdpSocket>,
+        remote_address: IpSocketAddress,
+    ) -> SocketResult<()> {
+        let remote_address = SocketAddr::from(remote_address);
+        if !(self.ctx.socket_addr_check)(remote_address, SocketAddrUse::UdpConnect).await {
+            return Err(ErrorCode::AccessDenied.into());
+        }
+        let socket = get_socket_mut(self.table, &socket)?;
+        socket.connect(remote_address)?;
+        Ok(())
+    }
+
     fn create(&mut self, address_family: IpAddressFamily) -> SocketResult<Resource<UdpSocket>> {
         let socket = UdpSocket::new(self.ctx, address_family.into())?;
         self.table

crates/wit-bindgen/src/lib.rs

@@ -2508,7 +2508,11 @@ impl<'a> InterfaceGenerator<'a> {
             } else {
                 uwriteln!(
                     self.src,
-                    "let host = {wt}::component::Access::new(caller, host_getter);"
+                    "let access_cx = {wt}::AsContextMut::as_context_mut(&mut caller);"
+                );
+                uwriteln!(
+                    self.src,
+                    "let host = {wt}::component::Access::new(access_cx, host_getter);"
                 );
             }
         } else {
@@ -2581,9 +2585,9 @@ impl<'a> InterfaceGenerator<'a> {
             let convert = format!("{}::convert_{}", convert_trait, err_name.to_snake_case());
             let convert = if flags.contains(FunctionFlags::STORE) {
                 if flags.contains(FunctionFlags::ASYNC) {
-                    format!("host.with(|mut host| {convert}(&mut host.get(), e))?")
+                    format!("caller.with(|mut host| {convert}(&mut host_getter(host.get()), e))?")
                 } else {
-                    format!("{convert}(&mut host.get(), e)?")
+                    format!("{convert}(&mut host_getter(caller.data_mut()), e)?")
                 }
             } else {
                 format!("{convert}(host, e)?")