tls-virt-wasmtime: run the sandbox address check on tunneled connects

rust/tls-virt-wasmtime/README.md

@@ -85,12 +85,12 @@ offered cipher suites are the profile's, verbatim.
 - **Delegation is trait-deep only.** wasmtime-wasi's trait impls are
   public and callable, but everything beneath them is `pub(crate)`:
   the `TcpSocket` state machine cannot be constructed or driven
-  externally, and `SocketAddrCheck` cannot be invoked directly. In
-  consequence the tunnel's own connect bypasses the sandbox's address
-  check (its transport is native tokio); only the inner name
-  resolution of opted-in names goes through wasmtime-wasi and its
-  `allow-ip-name-lookup` gate. A production wrapper would need its own
-  address policy for tunneled connects.
+  externally, and the installed address check is not readable back
+  from the `WasiCtx`. The check is the embedder's closure, though
+  (`WasiCtxBuilder::socket_addr_check`), so the host keeps its own
+  handle to it and runs it against a tunnel's real destination before
+  dialing (the transport is native tokio) — delegated sockets and
+  tunnels share one policy.
 
 ## Limits
 
@@ -99,9 +99,7 @@ Trust roots are the repository's baked test fixtures
 options on a tunneled socket reach the parked placeholder socket rather
 than the tunnel's transport (as does `get-address-family`), and TLS
 failures surface as `connection-reset`/stream closure with detail on
-stderr only. Tunnel connects on both paths bypass the sandbox address
-check (see the findings above). See issue #16 for the productionization
-gaps.
+stderr only. See issue #16 for the productionization gaps.
 
 ## Running
 

rust/tls-virt-wasmtime/src/main.rs

@@ -29,8 +29,10 @@
 //! ```
 //!
 //! Runs the component's `wasi:cli/run@0.3.0` export with stdio
-//! inherited, network inherited, and name lookup allowed. Limits are
-//! recorded in README.md.
+//! inherited, network inherited, and name lookup allowed. A tunnel
+//! connect resolves the destination, runs the embedder's sandbox
+//! address check against it, and only then dials. Limits are recorded
+//! in README.md.
 
 mod p2;
 
@@ -61,9 +63,22 @@ use wasmtime_wasi::p3::bindings::sockets::types::{
 };
 use wasmtime_wasi::p3::bindings::Command;
 use wasmtime_wasi::p3::sockets::{SocketError, SocketResult};
-use wasmtime_wasi::sockets::{WasiSockets, WasiSocketsCtxView};
+use wasmtime_wasi::sockets::{SocketAddrUse, WasiSockets, WasiSocketsCtxView};
 use wasmtime_wasi::{WasiCtx, WasiCtxBuilder, WasiCtxView, WasiView};
 
+/// The sandbox address policy: the same closure installed on the
+/// `WasiCtx` via `WasiCtxBuilder::socket_addr_check`. The host runs it
+/// against a tunnel's real destination before dialing, so delegated
+/// sockets and tunnels share one policy.
+pub(crate) type AddrCheck = Arc<
+    dyn Fn(
+            SocketAddr,
+            SocketAddrUse,
+        ) -> Pin<Box<dyn std::future::Future<Output = bool> + Send + Sync>>
+        + Send
+        + Sync,
+>;
+
 /// Trust anchor for tunneled connections (the repository's test CA;
 /// see README.md).
 const ROOT: &[u8] = include_bytes!("../../quic/tests/testdata/ca.der");
@@ -117,6 +132,8 @@ pub(crate) struct VirtCtx {
     /// Runtime handle for the close_notify shutdown task (spawned from
     /// a `Drop` impl, which cannot await).
     pub(crate) runtime: tokio::runtime::Handle,
+    /// The sandbox address policy; see `AddrCheck`.
+    pub(crate) addr_check: AddrCheck,
 }
 
 struct Tunnel {
@@ -131,7 +148,7 @@ struct Tunnel {
 }
 
 impl VirtCtx {
-    fn new() -> Result<Self> {
+    fn new(addr_check: AddrCheck) -> Result<Self> {
         let mut roots = rustls::RootCertStore::empty();
         roots
             .add(CertificateDer::from(ROOT.to_vec()))
@@ -145,6 +162,7 @@ impl VirtCtx {
             p2_resolves: HashMap::new(),
             connector: TlsConnector::from(Arc::new(config)),
             runtime: tokio::runtime::Handle::current(),
+            addr_check,
         })
     }
 
@@ -368,7 +386,15 @@ impl HostTcpSocketWithStore<Ctx> for VirtSockets {
         // The guest's socket resource stays in its unconnected state and
         // serves only as the handle the tunnel is keyed under.
         let addr = tls_virt_common::pick_addr(&addrs, port).ok_or(ErrorCode::RemoteUnreachable)?;
-        let connector = accessor.with(|mut a| a.get().virt.connector.clone());
+        let (connector, check) = accessor.with(|mut a| {
+            (
+                a.get().virt.connector.clone(),
+                a.get().virt.addr_check.clone(),
+            )
+        });
+        if !check(addr, SocketAddrUse::TcpConnect).await {
+            return Err(ErrorCode::AccessDenied.into());
+        }
 
         let stream = TcpStream::connect(addr).await.map_err(ErrorCode::from)?;
         let local = stream.local_addr().ok();
@@ -858,10 +884,14 @@ async fn main() -> Result<()> {
     types::add_to_linker::<Ctx, VirtSockets>(&mut linker, virt_view)?;
     ip_name_lookup::add_to_linker::<Ctx, VirtSockets>(&mut linker, virt_view)?;
 
+    let addr_check: AddrCheck = Arc::new(|_, _| Box::pin(async { true }));
     let mut wasi = WasiCtxBuilder::new();
     wasi.inherit_stdio()
         .args(&args[1..])
-        .inherit_network()
+        .socket_addr_check({
+            let check = addr_check.clone();
+            move |addr, use_| check(addr, use_)
+        })
         .allow_ip_name_lookup(true)
         .allow_tcp(true);
 
@@ -870,7 +900,7 @@ async fn main() -> Result<()> {
         Ctx {
             wasi: wasi.build(),
             table: ResourceTable::new(),
-            virt: VirtCtx::new()?,
+            virt: VirtCtx::new(addr_check)?,
         },
     );
     // The guest picks its world by its exports: a 0.3 command runs

rust/tls-virt-wasmtime/src/p2.rs

@@ -12,11 +12,12 @@
 //! The 0.2 shapes change the tunnel plumbing:
 //!
 //! - `start-connect`/`finish-connect` are a two-phase, poll-driven
-//!   pair: start spawns the TCP+TLS handshake as a background task;
-//!   finish reports would-block until it resolves; `subscribe` on a
-//!   tunnel socket returns a pollable over the handshake's completion
-//!   (a fresh owned table entry per call, so the pollable's lifetime
-//!   manages it).
+//!   pair: start spawns a background task that runs the sandbox
+//!   address check against the resolved destination, then the TCP+TLS
+//!   handshake; finish reports would-block until it resolves;
+//!   `subscribe` on a tunnel socket returns a pollable over the
+//!   handshake's completion (a fresh owned table entry per call, so
+//!   the pollable's lifetime manages it).
 //! - The data path is `wasi:io` streams, not component-model streams:
 //!   the returned input/output streams are this module's
 //!   [`TlsInputStream`]/[`TlsOutputStream`], byte buffers over the TLS
@@ -58,6 +59,7 @@ use wasmtime_wasi::p2::{
     subscribe, DynInputStream, DynOutputStream, DynPollable, InputStream, OutputStream, Pollable,
     SocketError, SocketResult, StreamError, StreamResult,
 };
+use wasmtime_wasi::sockets::SocketAddrUse;
 
 use tls_virt_common::Entry;
 
@@ -500,11 +502,15 @@ impl HostTcpSocket for VirtView<'_> {
             .ok_or(ErrorCode::RemoteUnreachable)?;
 
         let connector = self.virt.connector.clone();
+        let check = self.virt.addr_check.clone();
         let (done_tx, done_rx) = watch::channel(false);
         let result: Arc<Mutex<Option<Result<TlsParts, ErrorCode>>>> = Arc::new(Mutex::new(None));
         let slot = Arc::clone(&result);
         self.virt.runtime.spawn(async move {
             let outcome = async {
+                if !check(addr, SocketAddrUse::TcpConnect).await {
+                    return Err(ErrorCode::AccessDenied);
+                }
                 let stream = TcpStream::connect(addr).await.map_err(ErrorCode::from)?;
                 let local = stream.local_addr().ok();
                 let server_name = ServerName::try_from(hostname.clone())