Retain error raised by ruby socket check proc

william-stacken · william-stacken · commit 45ebe6dc0c32 · 2026-05-05T15:05:12.000+02:00
diff --git a/ext/src/ruby_api/component/func.rs b/ext/src/ruby_api/component/func.rs
@@ -105,6 +105,11 @@ impl Func {
         func.call(store.context_mut(), &params, &mut results)
             .map_err(|e| store_context_value.handle_wasm_error(ruby, e))?;
 
+        // Check for any errors stored during execution (e.g., from socket checks)
+        if let Some(error) = store_context_value.take_last_error()? {
+            return Err(error);
+        }
+
         match results_ty.len() {
             0 => Ok(ruby.qnil().as_value()),
             1 => component_val_to_rb(
diff --git a/ext/src/ruby_api/component/wasi_command.rs b/ext/src/ruby_api/component/wasi_command.rs
@@ -9,6 +9,7 @@ use crate::{
     ruby_api::{
         component::{linker::Linker, Component},
         errors,
+        store::StoreContextValue,
     },
     Store,
 };
@@ -39,13 +40,21 @@ impl WasiCommand {
     /// @def call_run(store)
     /// @param store [Store]
     /// @return [nil]
-    pub fn call_run(_ruby: &Ruby, rb_self: Obj<Self>, store: &Store) -> Result<(), Error> {
+    pub fn call_run(_ruby: &Ruby, rb_self: Obj<Self>, store: Obj<Store>) -> Result<(), Error> {
+        let store_context_value = StoreContextValue::from(store);
         rb_self
             .command
             .wasi_cli_run()
             .call_run(store.context_mut())
             .map_err(|err| error!("{err}"))?
-            .map_err(|_| error!("Error running `run`"))
+            .map_err(|_| error!("Error running `run`"))?;
+
+        // Check for any errors stored during execution (e.g., from socket checks)
+        if let Some(error) = store_context_value.take_last_error()? {
+            return Err(error);
+        }
+
+        Ok(())
     }
 }
 
diff --git a/ext/src/ruby_api/func.rs b/ext/src/ruby_api/func.rs
@@ -221,6 +221,11 @@ impl<'a> Func<'a> {
         func.call(context, &params, &mut results)
             .map_err(|e| store.handle_wasm_error(ruby, e))?;
 
+        // Check for any errors stored during execution (e.g., from socket checks)
+        if let Some(error) = store.take_last_error()? {
+            return Err(error);
+        }
+
         match results.as_slice() {
             [] => Ok(().into_value_with(ruby)),
             [result] => result.to_ruby_value(ruby, store),
diff --git a/ext/src/ruby_api/store.rs b/ext/src/ruby_api/store.rs
@@ -1,14 +1,16 @@
 use super::errors::wasi_exit_error;
 use super::{caller::Caller, engine::Engine, root, trap::Trap};
 use crate::{define_rb_intern, error, WasiConfig};
+use magnus::value::ReprValue;
 use magnus::value::StaticSymbol;
 use magnus::{
     class, function,
     gc::{Compactor, Marker},
     method, scan_args,
     typed_data::Obj,
     value::Opaque,
-    DataTypeFunctions, Error, IntoValue, Module, Object, Ruby, TypedData, Value,
+    DataTypeFunctions, Error, ExceptionClass, IntoValue, Module, Object, Ruby, TryConvert,
+    TypedData, Value,
 };
 use magnus::{Class, RHash};
 use rb_sys::tracking_allocator::{ManuallyTracked, TrackingAllocator};
@@ -70,6 +72,33 @@ impl StoreData {
         self.last_error.take()
     }
 
+    pub fn check_socket_errors(&mut self) {
+        use crate::ruby_api::wasi_config::WasiRetainedData;
+
+        // Check all retained values for WasiRetainedData and extract error storages
+        for value in &self.refs {
+            let ruby = Ruby::get().unwrap();
+
+            // Try to convert to WasiRetainedData
+            if let Ok(retained) = Obj::<WasiRetainedData>::try_convert(*value) {
+                if let Some(storage) = retained.error_storage() {
+                    if let Ok(mut guard) = storage.lock() {
+                        if let Some((class_name, error_msg)) = guard.take() {
+                            // Look up the exception class by name and reconstruct the error
+                            let exception_class = ruby
+                                .class_object()
+                                .const_get::<_, ExceptionClass>(class_name)
+                                .unwrap_or_else(|_| ruby.exception_runtime_error());
+
+                            self.last_error = Some(Error::new(exception_class, error_msg));
+                            return;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     pub fn mark(&self, marker: &Marker) {
         marker.mark_movable(self.user_data);
 
@@ -167,23 +196,21 @@ impl Store {
         let wasi_config = kw.optional.0;
         let wasi_p1_config = kw.optional.1;
 
-        let (wasi, wasi_proc) = wasi_config
+        let (wasi, wasi_retained) = wasi_config
             .map(|wasi_config| wasi_config.build(&ruby))
             .transpose()?
             .unzip();
-        let (wasi_p1, wasi_p1_proc) = wasi_p1_config
+        let (wasi_p1, wasi_p1_retained) = wasi_p1_config
             .map(|wasi_config| wasi_config.build_p1(&ruby))
             .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);
-        }
+        // Collect any Values that need to be retained for GC
+        let refs: Vec<Value> = [wasi_retained, wasi_p1_retained]
+            .into_iter()
+            .flatten()
+            .flatten()
+            .collect();
 
         let limiter = match kw.optional.2 {
             None => StoreLimitsBuilder::new(),
@@ -374,8 +401,14 @@ impl StoreContextValue<'_> {
         Ok(())
     }
 
-    fn take_last_error(&self) -> Result<Option<Error>, Error> {
+    pub fn take_last_error(&self) -> Result<Option<Error>, Error> {
         let ruby = Ruby::get().unwrap();
+
+        // Check for socket errors first and merge into last_error
+        if let Ok(mut context) = self.context_mut() {
+            context.data_mut().check_socket_errors();
+        }
+
         match self {
             Self::Store(store) => Ok(ruby.get_inner(*store).take_last_error()),
             Self::Caller(caller) => Ok(ruby
diff --git a/ext/src/ruby_api/wasi_config.rs b/ext/src/ruby_api/wasi_config.rs
@@ -6,6 +6,7 @@ use crate::{define_rb_intern, helpers::SymbolEnum};
 use lazy_static::lazy_static;
 use magnus::block::Proc;
 use magnus::value::ReprValue;
+use magnus::Class;
 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,
@@ -19,14 +20,49 @@ use std::future::Future;
 use std::net::SocketAddr;
 use std::path::Path;
 use std::pin::Pin;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 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};
 
+/// Storage for errors that occur in socket_addr_check callbacks.
+/// We store (class_name, message) tuples since magnus::Error is not Send+Sync.
+/// Storing the class name allows us to reconstruct the original exception type.
+type SocketErrorStorage = Arc<Mutex<Option<(String, String)>>>;
+
+/// Container for data that needs to be retained by the Store for WASI functionality.
+/// This includes Ruby procs (for GC marking) and error storages (for error propagation).
+#[derive(TypedData)]
+#[magnus(class = "Wasmtime::WasiRetainedData", mark, free_immediately)]
+pub struct WasiRetainedData {
+    proc: Option<Opaque<Proc>>,
+    error_storage: Option<SocketErrorStorage>,
+}
+
+impl DataTypeFunctions for WasiRetainedData {
+    fn mark(&self, marker: &Marker) {
+        if let Some(proc) = self.proc {
+            marker.mark(proc);
+        }
+    }
+}
+
+impl WasiRetainedData {
+    pub fn new(proc: Option<Proc>, error_storage: Option<SocketErrorStorage>) -> Self {
+        Self {
+            proc: proc.map(|p| p.into()),
+            error_storage,
+        }
+    }
+
+    pub fn error_storage(&self) -> Option<&SocketErrorStorage> {
+        self.error_storage.as_ref()
+    }
+}
+
 define_rb_intern!(
     READ => "read",
     WRITE => "write",
@@ -106,6 +142,7 @@ impl MappedDirectory {
 
 struct SocketAddrProc {
     proc: Proc,
+    error_storage: SocketErrorStorage,
 }
 
 impl SocketAddrProc {
@@ -118,8 +155,18 @@ impl SocketAddrProc {
 
         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
+            Err(error) => {
+                // Store both class name and message for later retrieval
+                if let Ok(mut storage) = self.error_storage.lock() {
+                    // Get the exception class name from the error's value
+                    let class_name = if let Some(exception_value) = error.value() {
+                        unsafe { exception_value.class().name().into_owned() }
+                    } else {
+                        "RuntimeError".to_string()
+                    };
+                    *storage = Some((class_name, error.to_string()));
+                }
+                // Deny access when an exception occurs
                 false
             }
         }
@@ -459,21 +506,22 @@ impl WasiConfig {
     }
 
     pub fn build_p1(&self, ruby: &Ruby) -> Result<(WasiP1Ctx, Option<Value>), Error> {
-        let (mut builder, proc_value) = self.build_impl(ruby)?;
+        let (mut builder, retained_data) = self.build_impl(ruby)?;
         let ctx = builder.build_p1();
-        Ok((ctx, proc_value))
+        Ok((ctx, retained_data))
     }
 
     pub fn build(&self, ruby: &Ruby) -> Result<(WasiCtx, Option<Value>), Error> {
-        let (mut builder, proc_value) = self.build_impl(ruby)?;
+        let (mut builder, retained_data) = self.build_impl(ruby)?;
         let ctx = builder.build();
-        Ok((ctx, proc_value))
+        Ok((ctx, retained_data))
     }
 
     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;
+        let mut error_storage_to_retain = None;
 
         if let Some(stdin) = inner.stdin.as_ref() {
             match stdin {
@@ -550,16 +598,21 @@ 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 });
+            let error_storage = Arc::new(Mutex::new(None));
+            let socket_addr_proc = Arc::new(SocketAddrProc {
+                proc,
+                error_storage: error_storage.clone(),
+            });
 
             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());
+            // Store the Proc and error storage together
+            proc_to_retain = Some(proc);
+            error_storage_to_retain = Some(error_storage);
         }
 
         for mapped_dir in &inner.mapped_directories {
@@ -578,7 +631,15 @@ impl WasiConfig {
                 .map_err(|e| error!("{}", e))?;
         }
 
-        Ok((builder, proc_to_retain))
+        // Wrap retained data in a single Ruby object if we have anything to retain
+        let retained_value = if proc_to_retain.is_some() || error_storage_to_retain.is_some() {
+            let retained_data = WasiRetainedData::new(proc_to_retain, error_storage_to_retain);
+            Some(ruby.obj_wrap(retained_data).as_value())
+        } else {
+            None
+        };
+
+        Ok((builder, retained_value))
     }
 
     fn check_determinism(&self) -> Result<(), Error> {
@@ -589,7 +650,9 @@ impl WasiConfig {
             || inner.allow_ip_name_lookup == Some(true);
 
         if inner.deterministic && has_network_enabled {
-            Err(error!("Sources of indeterminism cannot be combined with determinism"))
+            Err(error!(
+                "Sources of indeterminism cannot be combined with determinism"
+            ))
         } else {
             Ok(())
         }
@@ -608,6 +671,9 @@ pub fn file_w(path: RString) -> Result<File, Error> {
 }
 
 pub fn init(ruby: &Ruby) -> Result<(), Error> {
+    // Register internal WasiRetainedData class (not exposed to Ruby API)
+    root().define_class("WasiRetainedData", ruby.class_object())?;
+
     let class = root().define_class("WasiConfig", ruby.class_object())?;
     class.define_singleton_method("new", function!(WasiConfig::new, 0))?;
 
diff --git a/spec/unit/wasi_spec.rb b/spec/unit/wasi_spec.rb
@@ -617,10 +617,12 @@ module Wasmtime
           .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"
+            raise ArgumentError, "Intentional error for testing"
           end
 
-        run_wasi_component_network(wasi_config)
+        expect {
+          run_wasi_component_network(wasi_config)
+        }.to raise_error(ArgumentError, /Intentional error for testing/)
 
         result = JSON.parse(stdout_str)
         expect(result["test_type"]).to eq("tcp")
