From f65d71e77f1a10c42b9ce70cb29cfbfbdd69907d Mon Sep 17 00:00:00 2001 From: gngpp Date: Thu, 16 Jul 2026 00:43:57 +0800 Subject: [PATCH 1/3] fix(runtime): reject use after fork --- lib/wreq.rb | 4 ++ lib/wreq_ruby/body.rb | 3 ++ lib/wreq_ruby/client.rb | 5 ++ lib/wreq_ruby/error.rb | 12 +++++ lib/wreq_ruby/response.rb | 3 ++ src/arch.rs | 50 ++++++++++++++++++++ src/client.rs | 7 ++- src/client/resp.rs | 2 + src/error.rs | 28 +++++++++++ src/lib.rs | 4 +- src/macros.rs | 1 + src/rt.rs | 57 +++++++++++++++++++---- test/fork_test.rb | 97 +++++++++++++++++++++++++++++++++++++++ 13 files changed, 261 insertions(+), 12 deletions(-) create mode 100644 test/fork_test.rb diff --git a/lib/wreq.rb b/lib/wreq.rb index 5e5ee21..41cd88f 100644 --- a/lib/wreq.rb +++ b/lib/wreq.rb @@ -27,6 +27,10 @@ module Wreq # raise ArgumentError. Known values retain the error class from their Ruby # or native conversion, such as TypeError or Wreq::BuilderError. Validation # finishes before network I/O. + # + # Requests made in a child process forked after wreq-ruby was loaded raise + # Wreq::ForkError. Load wreq-ruby inside each worker after it has been + # forked. # Send an HTTP request. # diff --git a/lib/wreq_ruby/body.rb b/lib/wreq_ruby/body.rb index 549b517..3189f91 100644 --- a/lib/wreq_ruby/body.rb +++ b/lib/wreq_ruby/body.rb @@ -17,6 +17,8 @@ module Wreq # # A sender can be attached to one request. Closing it prevents further writes but # retains queued chunks so a request attached afterward can still drain them. + # Calling {#push} raises Wreq::ForkError in a child process forked after + # wreq-ruby was loaded. class BodySender # Create a bounded request-body sender. # @@ -33,6 +35,7 @@ def self.new(capacity = 8) # @param data [String] binary chunk # @return [nil] # @raise [IOError] if the sender or receiving side is closed + # @raise [Wreq::ForkError] if the process inherited wreq-ruby through fork def push(data) end diff --git a/lib/wreq_ruby/client.rb b/lib/wreq_ruby/client.rb index c9bef7c..d9a4cfb 100644 --- a/lib/wreq_ruby/client.rb +++ b/lib/wreq_ruby/client.rb @@ -17,6 +17,10 @@ module Wreq # native conversion, such as TypeError or Wreq::BuilderError. Request # validation finishes before network I/O. # + # Clients cannot be created or used in a child process forked after wreq-ruby + # was loaded. These operations raise Wreq::ForkError. Load wreq-ruby inside + # each worker after it has been forked. + # # @example Basic usage # client = Wreq::Client.new # # Use client for HTTP requests @@ -165,6 +169,7 @@ class Client # value cannot be converted or validated. # @raise [Wreq::BuilderError, Wreq::TlsError] if the native client cannot # be initialized. + # @raise [Wreq::ForkError] if the process inherited wreq-ruby through fork. # # @example Minimal client # client = Wreq::Client.new diff --git a/lib/wreq_ruby/error.rb b/lib/wreq_ruby/error.rb index 65dd4dc..745433c 100644 --- a/lib/wreq_ruby/error.rb +++ b/lib/wreq_ruby/error.rb @@ -7,6 +7,18 @@ module Wreq # Memory allocation failed. class MemoryError < StandardError; end + # The native extension was inherited from a parent process. + # + # Raised when wreq-ruby is used in a child process forked after the + # extension was loaded. Its process-global native state cannot be reused + # safely in the child. + # + # @example + # Process.fork do + # Wreq::Client.new # Raises when the parent loaded wreq-ruby. + # end + class ForkError < RuntimeError; end + # Network connection errors # Connection to the server failed. diff --git a/lib/wreq_ruby/response.rb b/lib/wreq_ruby/response.rb index fc77eac..fe816fa 100644 --- a/lib/wreq_ruby/response.rb +++ b/lib/wreq_ruby/response.rb @@ -8,6 +8,9 @@ module Wreq # access to HTTP response data including status codes, headers, body # content, and streaming capabilities. # + # Reading or streaming the body raises Wreq::ForkError in a child process + # forked after wreq-ruby was loaded. + # # @example Basic response handling # response = client.get("https://api.example.com") # puts response.status.as_int # => 200 diff --git a/src/arch.rs b/src/arch.rs index cf427d5..a0ea5c2 100644 --- a/src/arch.rs +++ b/src/arch.rs @@ -27,6 +27,56 @@ pub(crate) const SUPPORTS_INTERFACE: bool = cfg!(any( target_os = "watchos", )); +#[cfg(unix)] +mod unix { + use std::{ + ffi::c_int, + io, process, + sync::atomic::{AtomicBool, AtomicU32, Ordering}, + }; + + static FORKED: AtomicBool = AtomicBool::new(false); + static OWNER_PID: AtomicU32 = AtomicU32::new(0); + + unsafe extern "C" { + fn pthread_atfork( + prepare: Option, + parent: Option, + child: Option, + ) -> c_int; + } + + /// Mark the copied extension state as unusable in the forked child. + unsafe extern "C" fn mark_forked() { + FORKED.store(true, Ordering::Relaxed); + } + + /// Register process fork tracking before the extension exposes its API. + pub(crate) fn initialize_fork_tracking() -> io::Result<()> { + OWNER_PID.store(process::id(), Ordering::Relaxed); + + // POSIX runs the child handler before fork returns. The callback has a + // static lifetime and only stores an atomic flag. + // https://pubs.opengroup.org/onlinepubs/9799919799/functions/pthread_atfork.html + let status = unsafe { pthread_atfork(None, None, Some(mark_forked)) }; + if status == 0 { + Ok(()) + } else { + Err(io::Error::from_raw_os_error(status)) + } + } + + /// Return process IDs only when this process inherited the extension. + pub(crate) fn forked_process_ids() -> Option<(u32, u32)> { + FORKED + .load(Ordering::Relaxed) + .then(|| (OWNER_PID.load(Ordering::Relaxed), process::id())) + } +} + +#[cfg(unix)] +pub(crate) use unix::{forked_process_ids, initialize_fork_tracking}; + #[cfg(all(target_os = "windows", target_env = "gnu"))] mod windows_gnu { //! Windows GNU support. diff --git a/src/client.rs b/src/client.rs index 3bdec58..fc4dd6e 100644 --- a/src/client.rs +++ b/src/client.rs @@ -21,6 +21,7 @@ use crate::{ header::{Headers, OrigHeaders, UserAgent}, http::Method, options::{NativeOption, Options}, + rt, }; /// A builder for `Client`. @@ -214,8 +215,12 @@ impl Client { /// /// # Errors /// - /// Maps native build failures only after the GVL has been reacquired. + /// Returns `Wreq::ForkError` before touching native client state when the + /// extension was inherited from a parent process. Maps native build + /// failures only after the GVL has been reacquired. fn build(ruby: &Ruby, mut params: Builder) -> Result { + rt::ensure_current(ruby)?; + let result = gvl::nogvl(|| { let mut builder = wreq::Client::builder(); diff --git a/src/client/resp.rs b/src/client/resp.rs index 9fbc08c..10b3811 100644 --- a/src/client/resp.rs +++ b/src/client/resp.rs @@ -65,6 +65,8 @@ impl Response { /// Internal method to get the wreq::Response, optionally streaming the body. fn response(&self, ruby: &Ruby, stream: bool) -> Result { + rt::ensure_current(ruby)?; + let build_response = |body: wreq::Body| -> wreq::Response { let mut response = HttpResponse::new(body); *response.version_mut() = self.version.into_ffi(); diff --git a/src/error.rs b/src/error.rs index 3549407..81a8ed7 100644 --- a/src/error.rs +++ b/src/error.rs @@ -56,6 +56,7 @@ macro_rules! map_wreq_error { // System-level and runtime errors define_exception!(MEMORY, "MemoryError", exception_runtime_error); +define_exception!(FORK_ERROR, "ForkError", exception_runtime_error); // Network connection errors define_exception!(CONNECTION_ERROR, "ConnectionError", exception_runtime_error); @@ -94,6 +95,26 @@ pub fn interrupt_error(ruby: &Ruby) -> MagnusError { MagnusError::new(ruby.exception_interrupt(), "request interrupted") } +/// Build `Wreq::ForkError` without touching inherited native state. +#[cfg(unix)] +pub fn fork_error(ruby: &Ruby, owner_pid: u32, current_pid: u32) -> MagnusError { + MagnusError::new( + ruby.get_inner(&FORK_ERROR), + format!( + "wreq loaded in process {owner_pid} cannot be used after fork in process {current_pid}" + ), + ) +} + +/// Map a failed process-fork handler registration to `Wreq::ForkError`. +#[cfg(unix)] +pub fn fork_handler_error(ruby: &Ruby, err: &std::io::Error) -> MagnusError { + MagnusError::new( + ruby.get_inner(&FORK_ERROR), + format!("failed to initialize process fork tracking: {err}"), + ) +} + /// Map a Tokio runtime initialization failure to `Wreq::BuilderError`. pub fn runtime_initialization_error(ruby: &Ruby, err: &std::io::Error) -> MagnusError { MagnusError::new( @@ -235,6 +256,13 @@ pub fn include(ruby: &Ruby, gem_module: &RModule) -> Result<(), MagnusError> { "MemoryError", exception_runtime_error ); + initialize_exception!( + ruby, + gem_module, + FORK_ERROR, + "ForkError", + exception_runtime_error + ); initialize_exception!( ruby, gem_module, diff --git a/src/lib.rs b/src/lib.rs index ae53f87..82d852d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -85,6 +85,7 @@ pub fn patch(ruby: &Ruby, args: &[Value]) -> Result { fn init(ruby: &Ruby) -> Result<(), Error> { let gem_module = ruby.define_module(RUBY_MODULE_NAME)?; gem_module.const_set("VERSION", VERSION)?; + error::include(ruby, &gem_module)?; gem_module.define_module_function("request", magnus::function!(request, -1))?; gem_module.define_module_function("get", magnus::function!(get, -1))?; gem_module.define_module_function("post", magnus::function!(post, -1))?; @@ -99,6 +100,7 @@ fn init(ruby: &Ruby) -> Result<(), Error> { cookie::include(ruby, &gem_module)?; client::include(ruby, &gem_module)?; emulate::include(ruby, &gem_module)?; - error::include(ruby, &gem_module)?; + #[cfg(unix)] + rt::initialize(ruby)?; Ok(()) } diff --git a/src/macros.rs b/src/macros.rs index 1418b02..f9724e9 100644 --- a/src/macros.rs +++ b/src/macros.rs @@ -134,6 +134,7 @@ macro_rules! extract_request { let args = magnus::scan_args::scan_args::<$required, (), (), (), magnus::RHash, ()>($args)?; let required = args.required; let ruby = magnus::Ruby::get_with(args.keywords); + crate::rt::ensure_current(&ruby)?; let request = crate::client::req::Request::new(&ruby, args.keywords)?; (required, request) }}; diff --git a/src/rt.rs b/src/rt.rs index 28feba9..c037401 100644 --- a/src/rt.rs +++ b/src/rt.rs @@ -1,26 +1,57 @@ -use std::sync::LazyLock; +use std::{io, sync::OnceLock}; use magnus::Ruby; -use tokio::runtime::{Builder, Runtime}; +use tokio::runtime::{Builder, Runtime as TokioRuntime}; use crate::{ error::{interrupt_error, runtime_initialization_error}, gvl, }; -/// Initialize the global runtime lazily and preserve failures for Ruby. -static RUNTIME: LazyLock> = LazyLock::new(|| { - let mut builder = Builder::new_multi_thread(); +#[cfg(unix)] +use crate::{ + arch, + error::{fork_error, fork_handler_error}, +}; - builder.enable_all().build() -}); +/// Initialize the global runtime lazily and preserve failures for Ruby. +static RUNTIME: OnceLock> = OnceLock::new(); enum BlockOnError { Interrupted, Future(E), } -/// Block on a future to completion on the global Tokio runtime. +/// Register fork tracking while the native extension is being loaded. +/// +/// # Errors +/// +/// Returns `Wreq::ForkError` if the platform cannot install its child-process +/// callback. +#[cfg(unix)] +pub fn initialize(ruby: &Ruby) -> Result<(), magnus::Error> { + arch::initialize_fork_tracking().map_err(|err| fork_handler_error(ruby, &err)) +} + +/// Reject a child process that inherited the loaded native extension. +/// +/// # Errors +/// +/// Returns `Wreq::ForkError` when the extension was loaded before the current +/// process was forked. +pub fn ensure_current(ruby: &Ruby) -> Result<(), magnus::Error> { + #[cfg(unix)] + if let Some((owner_pid, current_pid)) = arch::forked_process_ids() { + return Err(fork_error(ruby, owner_pid, current_pid)); + } + + #[cfg(not(unix))] + let _ = ruby; + + Ok(()) +} + +/// Block on a future to completion on the current process's global Tokio runtime. /// /// The future runs without Ruby's GVL, so it must not construct Ruby objects or /// Ruby exceptions. Convert Rust errors back into Ruby errors after the GVL has @@ -28,15 +59,21 @@ enum BlockOnError { /// /// # Errors /// -/// Returns `Wreq::BuilderError` if the Tokio runtime cannot be initialized, -/// Ruby's standard `Interrupt` if Ruby interrupts the request, or the error produced +/// Returns `Wreq::ForkError` if the extension belongs to a parent process, +/// `Wreq::BuilderError` if the Tokio runtime cannot be initialized, Ruby's +/// standard `Interrupt` if Ruby interrupts the request, or the error produced /// by `map_err` if the future fails. pub fn try_block_on(ruby: &Ruby, future: F, map_err: M) -> Result where F: Future>, M: FnOnce(&Ruby, E) -> magnus::Error, { + ensure_current(ruby)?; let runtime = RUNTIME + .get_or_init(|| { + let mut builder = Builder::new_multi_thread(); + builder.enable_all().build() + }) .as_ref() .map_err(|err| runtime_initialization_error(ruby, err))?; let result = gvl::nogvl_cancellable(|flag| { diff --git a/test/fork_test.rb b/test/fork_test.rb new file mode 100644 index 0000000..1a6a259 --- /dev/null +++ b/test/fork_test.rb @@ -0,0 +1,97 @@ +# frozen_string_literal: true + +require "test_helper" +require "open3" +require "rbconfig" +require "timeout" + +class ForkTest < Minitest::Test + def test_fork_error_is_a_runtime_error + assert_operator Wreq::ForkError, :<, RuntimeError + end + + def test_loaded_extension_is_rejected_after_fork + skip "fork is not supported on this platform" unless Process.respond_to?(:fork) + + lib_dir = File.expand_path("../lib", __dir__) + script = <<~'RUBY' + require "socket" + require "timeout" + require "wreq" + + $stdout.sync = true + $stderr.sync = true + + def expect_fork_error(label) + child_pid = fork do + begin + Timeout.timeout(5) { yield } + rescue Wreq::ForkError => error + warn "#{label}=#{error.class}: #{error.message}" + exit! 0 + rescue Exception => error + warn "#{label}=unexpected #{error.class}: #{error.message}" + exit! 2 + end + + warn "#{label}=missing Wreq::ForkError" + exit! 3 + end + + _, status = Process.wait2(child_pid) + abort "#{label} child failed with #{status.inspect}" unless status.success? + end + + expect_fork_error("before_runtime") { Wreq::Client.new } + + server = TCPServer.new("127.0.0.1", 0) + port = server.addr[1] + server_pid = fork do + 2.times do + ready = IO.select([server], nil, nil, 10) + exit! 4 unless ready + + socket = server.accept + begin + while (line = socket.gets) + break if line == "\r\n" + end + socket.write("HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok") + ensure + socket.close + end + end + exit! 0 + ensure + server.close + end + server.close + + url = "http://127.0.0.1:#{port}/" + client = Wreq::Client.new + abort "parent warm-up failed" unless client.get(url).bytes == "ok" + + expect_fork_error("fresh_client") { Wreq::Client.new } + expect_fork_error("inherited_client") { client.get(url) } + + abort "parent request after fork failed" unless client.get(url).bytes == "ok" + _, server_status = Process.wait2(server_pid) + abort "server failed with #{server_status.inspect}" unless server_status.success? + + puts "ok" + RUBY + + stdout, stderr, status = Timeout.timeout(20) do + Open3.capture3(RbConfig.ruby, "-I", lib_dir, "-e", script) + end + + assert status.success?, "subprocess failed with #{status.inspect}: #{stderr}" + assert_equal "ok\n", stdout + assert_match(/before_runtime=Wreq::ForkError:.*cannot be used after fork/, stderr) + assert_match(/fresh_client=Wreq::ForkError:.*cannot be used after fork/, stderr) + assert_match(/inherited_client=Wreq::ForkError:.*cannot be used after fork/, stderr) + refute_match(/\[BUG\]|segmentation fault|panicked/i, stderr) + rescue Timeout::Error + flunk "fork safety subprocess timed out" + end +end From bd16c7f380f043797a369c9dd1f722d08139a164 Mon Sep 17 00:00:00 2001 From: gngpp Date: Thu, 16 Jul 2026 09:08:16 +0800 Subject: [PATCH 2/3] test(runtime): extract fork safety script --- test/fork_test.rb | 69 ++----------------------------------- test/scripts/fork_safety.rb | 66 +++++++++++++++++++++++++++++++++++ 2 files changed, 68 insertions(+), 67 deletions(-) create mode 100644 test/scripts/fork_safety.rb diff --git a/test/fork_test.rb b/test/fork_test.rb index 1a6a259..9330695 100644 --- a/test/fork_test.rb +++ b/test/fork_test.rb @@ -14,75 +14,10 @@ def test_loaded_extension_is_rejected_after_fork skip "fork is not supported on this platform" unless Process.respond_to?(:fork) lib_dir = File.expand_path("../lib", __dir__) - script = <<~'RUBY' - require "socket" - require "timeout" - require "wreq" - - $stdout.sync = true - $stderr.sync = true - - def expect_fork_error(label) - child_pid = fork do - begin - Timeout.timeout(5) { yield } - rescue Wreq::ForkError => error - warn "#{label}=#{error.class}: #{error.message}" - exit! 0 - rescue Exception => error - warn "#{label}=unexpected #{error.class}: #{error.message}" - exit! 2 - end - - warn "#{label}=missing Wreq::ForkError" - exit! 3 - end - - _, status = Process.wait2(child_pid) - abort "#{label} child failed with #{status.inspect}" unless status.success? - end - - expect_fork_error("before_runtime") { Wreq::Client.new } - - server = TCPServer.new("127.0.0.1", 0) - port = server.addr[1] - server_pid = fork do - 2.times do - ready = IO.select([server], nil, nil, 10) - exit! 4 unless ready - - socket = server.accept - begin - while (line = socket.gets) - break if line == "\r\n" - end - socket.write("HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok") - ensure - socket.close - end - end - exit! 0 - ensure - server.close - end - server.close - - url = "http://127.0.0.1:#{port}/" - client = Wreq::Client.new - abort "parent warm-up failed" unless client.get(url).bytes == "ok" - - expect_fork_error("fresh_client") { Wreq::Client.new } - expect_fork_error("inherited_client") { client.get(url) } - - abort "parent request after fork failed" unless client.get(url).bytes == "ok" - _, server_status = Process.wait2(server_pid) - abort "server failed with #{server_status.inspect}" unless server_status.success? - - puts "ok" - RUBY + script = File.expand_path("scripts/fork_safety.rb", __dir__) stdout, stderr, status = Timeout.timeout(20) do - Open3.capture3(RbConfig.ruby, "-I", lib_dir, "-e", script) + Open3.capture3(RbConfig.ruby, "-I", lib_dir, script) end assert status.success?, "subprocess failed with #{status.inspect}: #{stderr}" diff --git a/test/scripts/fork_safety.rb b/test/scripts/fork_safety.rb new file mode 100644 index 0000000..1093bdb --- /dev/null +++ b/test/scripts/fork_safety.rb @@ -0,0 +1,66 @@ +# frozen_string_literal: true + +require "socket" +require "timeout" +require "wreq" + +$stdout.sync = true +$stderr.sync = true + +def expect_fork_error(label) + child_pid = fork do + begin + Timeout.timeout(5) { yield } + rescue Wreq::ForkError => error + warn "#{label}=#{error.class}: #{error.message}" + exit! 0 + rescue => error + warn "#{label}=unexpected #{error.class}: #{error.message}" + exit! 2 + end + + warn "#{label}=missing Wreq::ForkError" + exit! 3 + end + + _, status = Process.wait2(child_pid) + abort "#{label} child failed with #{status.inspect}" unless status.success? +end + +expect_fork_error("before_runtime") { Wreq::Client.new } + +server = TCPServer.new("127.0.0.1", 0) +port = server.addr[1] +server_pid = fork do + 2.times do + ready = IO.select([server], nil, nil, 10) + exit! 4 unless ready + + socket = server.accept + begin + while (line = socket.gets) + break if line == "\r\n" + end + socket.write("HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok") + ensure + socket.close + end + end + exit! 0 +ensure + server.close +end +server.close + +url = "http://127.0.0.1:#{port}/" +client = Wreq::Client.new +abort "parent warm-up failed" unless client.get(url).bytes == "ok" + +expect_fork_error("fresh_client") { Wreq::Client.new } +expect_fork_error("inherited_client") { client.get(url) } + +abort "parent request after fork failed" unless client.get(url).bytes == "ok" +_, server_status = Process.wait2(server_pid) +abort "server failed with #{server_status.inspect}" unless server_status.success? + +puts "ok" From 9e85866402efb7069e62df11944dbb3cf59f716e Mon Sep 17 00:00:00 2001 From: gngpp Date: Thu, 16 Jul 2026 09:53:05 +0800 Subject: [PATCH 3/3] refactor(runtime): use forkguard for fork detection --- Cargo.lock | 10 ++++++ Cargo.toml | 3 ++ src/arch.rs | 68 +++++++++++++++++++++---------------- test/fork_test.rb | 7 ++-- test/scripts/fork_safety.rb | 26 ++++++++------ 5 files changed, 71 insertions(+), 43 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7c8e22c..3c90168 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -349,6 +349,15 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b" +[[package]] +name = "forkguard" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "723f03dfba28f1949cae68c6ac89a1e77f753c38c0acfceeefec77578a1c357a" +dependencies = [ + "libc", +] + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -1547,6 +1556,7 @@ dependencies = [ "arc-swap", "bytes", "cookie", + "forkguard", "futures-util", "http", "http-body-util", diff --git a/Cargo.toml b/Cargo.toml index bad747a..db49e7c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -47,6 +47,9 @@ http = "1.4.1" http-body-util = "0.1.3" futures-util = { version = "0.3.32", default-features = false } +[target.'cfg(unix)'.dependencies] +forkguard = { version = "0.1.3", features = ["atfork"] } + [dev-dependencies] magnus = { version = "0.8.2", features = ["embed"] } diff --git a/src/arch.rs b/src/arch.rs index a0ea5c2..18cb2a4 100644 --- a/src/arch.rs +++ b/src/arch.rs @@ -29,48 +29,56 @@ pub(crate) const SUPPORTS_INTERFACE: bool = cfg!(any( #[cfg(unix)] mod unix { - use std::{ - ffi::c_int, - io, process, - sync::atomic::{AtomicBool, AtomicU32, Ordering}, - }; + use std::{io, process, sync::OnceLock}; - static FORKED: AtomicBool = AtomicBool::new(false); - static OWNER_PID: AtomicU32 = AtomicU32::new(0); - - unsafe extern "C" { - fn pthread_atfork( - prepare: Option, - parent: Option, - child: Option, - ) -> c_int; + /// Process state captured when the extension initializes. + /// + /// The atfork guard uses a POSIX child handler to advance an atomic fork + /// generation without running Ruby code. + /// https://pubs.opengroup.org/onlinepubs/9799919799/functions/pthread_atfork.html + struct ForkGuard { + detector: forkguard::Guard, + owner_pid: u32, } - /// Mark the copied extension state as unusable in the forked child. - unsafe extern "C" fn mark_forked() { - FORKED.store(true, Ordering::Relaxed); + impl ForkGuard { + /// Create a guard and register fork detection with the process. + fn new() -> io::Result { + forkguard::Guard::try_new() + .map(|detector| Self { + detector, + owner_pid: process::id(), + }) + .map_err(|error| io::Error::from_raw_os_error(error.code().get())) + } + + /// Return process IDs when this guard was inherited through a fork. + fn forked_process_ids(&self) -> Option<(u32, u32)> { + // Keep the stored generation unchanged so every runtime access in + // the child remains rejected. Cloning the detector copies one usize. + let mut detector = self.detector.clone(); + detector + .detected_fork() + .then(|| (self.owner_pid, process::id())) + } } + static FORK_GUARD: OnceLock = OnceLock::new(); + /// Register process fork tracking before the extension exposes its API. pub(crate) fn initialize_fork_tracking() -> io::Result<()> { - OWNER_PID.store(process::id(), Ordering::Relaxed); - - // POSIX runs the child handler before fork returns. The callback has a - // static lifetime and only stores an atomic flag. - // https://pubs.opengroup.org/onlinepubs/9799919799/functions/pthread_atfork.html - let status = unsafe { pthread_atfork(None, None, Some(mark_forked)) }; - if status == 0 { - Ok(()) - } else { - Err(io::Error::from_raw_os_error(status)) + if FORK_GUARD.get().is_some() { + return Ok(()); } + + let guard = ForkGuard::new()?; + let _ = FORK_GUARD.set(guard); + Ok(()) } /// Return process IDs only when this process inherited the extension. pub(crate) fn forked_process_ids() -> Option<(u32, u32)> { - FORKED - .load(Ordering::Relaxed) - .then(|| (OWNER_PID.load(Ordering::Relaxed), process::id())) + FORK_GUARD.get().and_then(ForkGuard::forked_process_ids) } } diff --git a/test/fork_test.rb b/test/fork_test.rb index 9330695..60da69a 100644 --- a/test/fork_test.rb +++ b/test/fork_test.rb @@ -22,9 +22,10 @@ def test_loaded_extension_is_rejected_after_fork assert status.success?, "subprocess failed with #{status.inspect}: #{stderr}" assert_equal "ok\n", stdout - assert_match(/before_runtime=Wreq::ForkError:.*cannot be used after fork/, stderr) - assert_match(/fresh_client=Wreq::ForkError:.*cannot be used after fork/, stderr) - assert_match(/inherited_client=Wreq::ForkError:.*cannot be used after fork/, stderr) + %w[before_runtime fresh_client inherited_client].each do |label| + assert_match(/#{label}=Wreq::ForkError:.*cannot be used after fork/, stderr) + assert_match(/#{label}_retry=Wreq::ForkError:.*cannot be used after fork/, stderr) + end refute_match(/\[BUG\]|segmentation fault|panicked/i, stderr) rescue Timeout::Error flunk "fork safety subprocess timed out" diff --git a/test/scripts/fork_safety.rb b/test/scripts/fork_safety.rb index 1093bdb..b50b4a5 100644 --- a/test/scripts/fork_safety.rb +++ b/test/scripts/fork_safety.rb @@ -9,18 +9,24 @@ def expect_fork_error(label) child_pid = fork do - begin - Timeout.timeout(5) { yield } - rescue Wreq::ForkError => error - warn "#{label}=#{error.class}: #{error.message}" - exit! 0 - rescue => error - warn "#{label}=unexpected #{error.class}: #{error.message}" - exit! 2 + 2.times do |attempt| + attempt_label = attempt.zero? ? label : "#{label}_retry" + + begin + Timeout.timeout(5) { yield } + rescue Wreq::ForkError => error + warn "#{attempt_label}=#{error.class}: #{error.message}" + next + rescue => error + warn "#{attempt_label}=unexpected #{error.class}: #{error.message}" + exit! 2 + end + + warn "#{attempt_label}=missing Wreq::ForkError" + exit! 3 end - warn "#{label}=missing Wreq::ForkError" - exit! 3 + exit! 0 end _, status = Process.wait2(child_pid)