Add socket_addr_check support

william-stacken · william-stacken · commit 1533739cf4ca · 2026-04-30T17:26:41.000+02:00
diff --git a/ext/src/ruby_api/store.rs b/ext/src/ruby_api/store.rs
@@ -167,12 +167,23 @@ impl Store {
         let wasi_config = kw.optional.0;
         let wasi_p1_config = kw.optional.1;
 
-        let wasi = wasi_config
+        let (wasi, wasi_proc) = wasi_config
             .map(|wasi_config| wasi_config.build(&ruby))
-            .transpose()?;
-        let wasi_p1 = wasi_p1_config
+            .transpose()?
+            .unzip();
+        let (wasi_p1, wasi_p1_proc) = wasi_p1_config
             .map(|wasi_config| wasi_config.build_p1(&ruby))
-            .transpose()?;
+            .transpose()?
+            .unzip();
+
+        // Collect any Procs that need to be retained
+        let mut refs = Vec::new();
+        if let Some(proc) = wasi_proc.flatten() {
+            refs.push(proc);
+        }
+        if let Some(proc) = wasi_p1_proc.flatten() {
+            refs.push(proc);
+        }
 
         let limiter = match kw.optional.2 {
             None => StoreLimitsBuilder::new(),
@@ -186,7 +197,7 @@ impl Store {
             user_data,
             wasi_p1,
             wasi,
-            refs: Default::default(),
+            refs,
             last_error: Default::default(),
             store_limits: limiter,
             resource_table: Default::default(),
diff --git a/ext/src/ruby_api/wasi_config.rs b/ext/src/ruby_api/wasi_config.rs
@@ -4,19 +4,27 @@ use crate::helpers::OutputLimitedBuffer;
 use crate::ruby_api::convert::ToValType;
 use crate::{define_rb_intern, helpers::SymbolEnum};
 use lazy_static::lazy_static;
+use magnus::block::Proc;
+use magnus::value::ReprValue;
 use magnus::{
     class, function, gc::Marker, method, typed_data::Obj, value::Opaque, DataTypeFunctions, Error,
     IntoValue, Module, Object, RArray, RHash, RString, Ruby, Symbol, TryConvert, TypedData, Value,
 };
 use rb_sys::ruby_rarray_flags::RARRAY_EMBED_FLAG;
+use rb_sys::VALUE;
 use std::cell::RefCell;
 use std::convert::TryFrom;
 use std::fs;
+use std::future::Future;
+use std::net::SocketAddr;
 use std::path::Path;
+use std::pin::Pin;
+use std::sync::Arc;
 use std::{fs::File, path::PathBuf};
 use wasmtime_wasi::cli::{InputFile, OutputFile};
 use wasmtime_wasi::p1::WasiP1Ctx;
 use wasmtime_wasi::p2::pipe::MemoryInputPipe;
+use wasmtime_wasi::sockets::SocketAddrUse;
 use wasmtime_wasi::{DirPerms, FilePerms, WasiCtx, WasiCtxBuilder};
 
 define_rb_intern!(
@@ -96,6 +104,43 @@ impl MappedDirectory {
     }
 }
 
+struct SocketAddrProc {
+    proc: Proc,
+}
+
+impl SocketAddrProc {
+    fn call(&self, addr: SocketAddr, use_: SocketAddrUse) -> bool {
+        let ruby = Ruby::get().unwrap();
+
+        // Convert arguments to Ruby values
+        let addr_str = ruby.str_new(&addr.to_string());
+        let use_sym = socket_addr_use_to_symbol(&ruby, use_);
+
+        match self.proc.call::<_, Value>((addr_str, use_sym)) {
+            Ok(result) => bool::try_convert(result).unwrap_or(false),
+            Err(_) => {
+                // Exception in Ruby block, deny access
+                false
+            }
+        }
+    }
+}
+
+// SAFETY: We only access the Ruby proc when we have the GVL (during WASI operations).
+// The Proc is kept alive by the Store's refs field, which is marked during GC.
+unsafe impl Send for SocketAddrProc {}
+unsafe impl Sync for SocketAddrProc {}
+
+fn socket_addr_use_to_symbol(ruby: &Ruby, use_: SocketAddrUse) -> Symbol {
+    match use_ {
+        SocketAddrUse::TcpBind => ruby.to_symbol("tcp_bind"),
+        SocketAddrUse::TcpConnect => ruby.to_symbol("tcp_connect"),
+        SocketAddrUse::UdpBind => ruby.to_symbol("udp_bind"),
+        SocketAddrUse::UdpConnect => ruby.to_symbol("udp_connect"),
+        SocketAddrUse::UdpOutgoingDatagram => ruby.to_symbol("udp_outgoing_datagram"),
+    }
+}
+
 #[derive(Default)]
 struct WasiConfigInner {
     stdin: Option<ReadStream>,
@@ -109,6 +154,7 @@ struct WasiConfigInner {
     allow_tcp: Option<bool>,
     allow_udp: Option<bool>,
     allow_ip_name_lookup: Option<bool>,
+    socket_addr_check: Option<Opaque<Proc>>,
 }
 
 impl WasiConfigInner {
@@ -131,6 +177,9 @@ impl WasiConfigInner {
         for v in &self.mapped_directories {
             v.mark(marker);
         }
+        if let Some(v) = self.socket_addr_check.as_ref() {
+            marker.mark(*v);
+        }
     }
 }
 
@@ -384,21 +433,47 @@ impl WasiConfig {
         rb_self
     }
 
-    pub fn build_p1(&self, ruby: &Ruby) -> Result<WasiP1Ctx, Error> {
-        let mut builder = self.build_impl(ruby)?;
+    /// @yard
+    /// Set a custom check function for socket address access control.
+    /// The block will be called for each socket operation with the socket address (as a String)
+    /// and the operation type (as a Symbol: :tcp_bind, :tcp_connect, :udp_bind, :udp_connect,
+    /// :udp_outgoing_datagram).
+    /// The block should return true to allow the operation or false to deny it.
+    /// If the block raises an exception, the operation will be denied.
+    ///
+    /// Note: any network access happens while the Global VM Lock (GVL) is held, so other
+    /// threads will be blocked in the meantime.
+    ///
+    /// @yieldparam addr [String] The socket address (e.g., "127.0.0.1:8080")
+    /// @yieldparam use [Symbol] The type of socket operation
+    /// @yieldreturn [Boolean] true to allow the operation, false to deny it
+    /// @def socket_addr_check
+    /// @return [WasiConfig] +self+
+    pub fn socket_addr_check(ruby: &Ruby, rb_self: RbSelf) -> RbSelf {
+        if ruby.block_given() {
+            let proc = ruby.block_proc().unwrap();
+            let mut inner = rb_self.inner.borrow_mut();
+            inner.socket_addr_check = Some(proc.into());
+        }
+        rb_self
+    }
+
+    pub fn build_p1(&self, ruby: &Ruby) -> Result<(WasiP1Ctx, Option<Value>), Error> {
+        let (mut builder, proc_value) = self.build_impl(ruby)?;
         let ctx = builder.build_p1();
-        Ok(ctx)
+        Ok((ctx, proc_value))
     }
 
-    pub fn build(&self, ruby: &Ruby) -> Result<WasiCtx, Error> {
-        let mut builder = self.build_impl(ruby)?;
+    pub fn build(&self, ruby: &Ruby) -> Result<(WasiCtx, Option<Value>), Error> {
+        let (mut builder, proc_value) = self.build_impl(ruby)?;
         let ctx = builder.build();
-        Ok(ctx)
+        Ok((ctx, proc_value))
     }
 
-    fn build_impl(&self, ruby: &Ruby) -> Result<WasiCtxBuilder, Error> {
+    fn build_impl(&self, ruby: &Ruby) -> Result<(WasiCtxBuilder, Option<Value>), Error> {
         let mut builder = WasiCtxBuilder::new();
         let inner = self.inner.borrow();
+        let mut proc_to_retain = None;
 
         if let Some(stdin) = inner.stdin.as_ref() {
             match stdin {
@@ -487,6 +562,20 @@ impl WasiConfig {
             }
         }
 
+        if let Some(check_proc) = inner.socket_addr_check.as_ref() {
+            let proc = ruby.get_inner(*check_proc);
+            let socket_addr_proc = Arc::new(SocketAddrProc { proc });
+
+            builder.socket_addr_check(move |addr, use_| {
+                let socket_addr_proc = socket_addr_proc.clone();
+                Box::pin(async move { socket_addr_proc.call(addr, use_) })
+                    as Pin<Box<dyn Future<Output = bool> + Send + Sync>>
+            });
+
+            // Store the Proc as a Value so the Store can retain it
+            proc_to_retain = Some(proc.as_value());
+        }
+
         for mapped_dir in &inner.mapped_directories {
             let host_path = ruby.get_inner(mapped_dir.host_path).to_string()?;
             let guest_path = ruby.get_inner(mapped_dir.guest_path).to_string()?;
@@ -503,7 +592,7 @@ impl WasiConfig {
                 .map_err(|e| error!("{}", e))?;
         }
 
-        Ok(builder)
+        Ok((builder, proc_to_retain))
     }
 }
 
@@ -558,6 +647,10 @@ pub fn init(ruby: &Ruby) -> Result<(), Error> {
         "allow_ip_name_lookup",
         method!(WasiConfig::allow_ip_name_lookup, 1),
     )?;
+    class.define_method(
+        "socket_addr_check",
+        method!(WasiConfig::socket_addr_check, 0),
+    )?;
 
     Ok(())
 }
diff --git a/spec/unit/wasi_spec.rb b/spec/unit/wasi_spec.rb
@@ -560,6 +560,74 @@ module Wasmtime
         expect(result["test_type"]).to eq("dns")
         expect(result["success"]).to eq(false)
       end
+
+      it "allows selective access with socket_addr_check" do
+        port_file = tempfile_path("tcp_port_selective")
+        server_pid = spawn_tcp_server(port_file)
+        port = wait_for_port(port_file)
+
+        stdout_str = ""
+        wasi_config = WasiConfig.new
+          .set_argv(["wasi-network", "tcp", "127.0.0.1", port.to_s])
+          .set_stdout_buffer(stdout_str, 40000)
+          .socket_addr_check do |addr, use|
+            # Only allow connections to localhost on the specific port
+            addr.start_with?("127.0.0.1:") && use == :tcp_connect
+          end
+
+        run_wasi_component_network(wasi_config)
+
+        result = JSON.parse(stdout_str)
+        expect(result["test_type"]).to eq("tcp")
+        expect(result["message"]).to match(/and exchanged data/)
+        expect(result["success"]).to eq(true)
+      ensure
+        cleanup_server(server_pid)
+      end
+
+      it "blocks access when socket_addr_check returns false" do
+        port_file = tempfile_path("tcp_port_blocked")
+        server_pid = spawn_tcp_server(port_file)
+        port = wait_for_port(port_file)
+
+        stdout_str = ""
+        wasi_config = WasiConfig.new
+          .set_argv(["wasi-network", "tcp", "127.0.0.1", port.to_s])
+          .set_stdout_buffer(stdout_str, 40000)
+          .socket_addr_check do |_addr, _use|
+            false # Block all access
+          end
+
+        run_wasi_component_network(wasi_config)
+
+        result = JSON.parse(stdout_str)
+        expect(result["test_type"]).to eq("tcp")
+        expect(result["success"]).to eq(false)
+      ensure
+        cleanup_server(server_pid)
+      end
+
+      it "blocks access when socket_addr_check raises an exception" do
+        port_file = tempfile_path("tcp_port_exception")
+        server_pid = spawn_tcp_server(port_file)
+        port = wait_for_port(port_file)
+
+        stdout_str = ""
+        wasi_config = WasiConfig.new
+          .set_argv(["wasi-network", "tcp", "127.0.0.1", port.to_s])
+          .set_stdout_buffer(stdout_str, 40000)
+          .socket_addr_check do |_addr, _use|
+            raise "Intentional error for testing"
+          end
+
+        run_wasi_component_network(wasi_config)
+
+        result = JSON.parse(stdout_str)
+        expect(result["test_type"]).to eq("tcp")
+        expect(result["success"]).to eq(false)
+      ensure
+        cleanup_server(server_pid)
+      end
     end
 
     describe "WasiConfig preview 1" do
