From 89f7c13667349e6a3d9825266a9838cae79cb60c Mon Sep 17 00:00:00 2001 From: gngpp Date: Sat, 11 Jul 2026 12:03:21 +0800 Subject: [PATCH] fix(body): validate body sender lifecycle --- lib/wreq_ruby/body.rb | 38 ++++-- src/client/body.rs | 1 + src/client/body/stream.rs | 178 ++++++++++++++++++++++----- src/error.rs | 29 ++++- test/body_sender_test.rb | 246 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 447 insertions(+), 45 deletions(-) create mode 100644 test/body_sender_test.rb diff --git a/lib/wreq_ruby/body.rb b/lib/wreq_ruby/body.rb index c3489bf..06013b0 100644 --- a/lib/wreq_ruby/body.rb +++ b/lib/wreq_ruby/body.rb @@ -2,10 +2,10 @@ unless defined?(Wreq) module Wreq - # BodySender: for streaming request bodies, allowing thread-safe chunked data push. - # Backed by a Rust channel, avoids buffering the entire payload in memory. + # Streams a request body through a bounded, thread-safe channel. # - # Supports multi-threaded chunk push: you can safely call `push` from multiple threads. + # The channel applies backpressure instead of buffering the entire request body. + # Multiple threads may safely push chunks while a request drains the receiving side. # # Usage: # sender = Wreq::BodySender.new(8) @@ -15,22 +15,42 @@ module Wreq # end # resp = client.post(url, body: sender) # - # Note: - # - Sender is for request upload only, not for response reading. - # - Each BodySender instance can only be used once (single-use): - # after being consumed by a request, further push or reuse is not allowed. + # 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. class BodySender - # @param capacity [Integer] channel buffer size, default 8 + # Create a bounded request-body sender. + # + # @param capacity [Integer] positive number of chunks that may wait in the channel; + # defaults to 8 and must be greater than zero + # @raise [ArgumentError] if capacity is zero, negative, or too large + # @raise [TypeError] if capacity is not an Integer def self.new(capacity = 8) end + # Push one binary chunk, waiting while the channel is full. + # # @param data [String] binary chunk + # @return [nil] + # @raise [IOError] if the sender or receiving side is closed def push(data) end - # Close the sender, signaling end of data. + # Close the producer and signal EOF after all queued chunks are read. + # + # This operation is idempotent. + # + # @return [nil] def close end + + # Return whether the sender can no longer accept chunks. + # + # This becomes true after {#close} or when the request stops consuming + # the receiving side. + # + # @return [Boolean] + def closed? + end end end end diff --git a/src/client/body.rs b/src/client/body.rs index c5d770f..4f96455 100644 --- a/src/client/body.rs +++ b/src/client/body.rs @@ -53,5 +53,6 @@ pub fn include(ruby: &Ruby, gem_module: &RModule) -> Result<(), Error> { sender_class.define_singleton_method("new", function!(BodySender::new, -1))?; sender_class.define_method("push", method!(BodySender::push, 1))?; sender_class.define_method("close", magnus::method!(BodySender::close, 0))?; + sender_class.define_method("closed?", magnus::method!(BodySender::is_closed, 0))?; Ok(()) } diff --git a/src/client/body/stream.rs b/src/client/body/stream.rs index 1146c36..909931c 100644 --- a/src/client/body/stream.rs +++ b/src/client/body/stream.rs @@ -1,34 +1,57 @@ use std::{ + cell::{Ref, RefCell, RefMut}, + panic::catch_unwind, pin::Pin, - sync::RwLock, task::{Context, Poll}, }; use bytes::Bytes; use futures_util::{Stream, StreamExt}; -use magnus::{Error, RString, TryConvert, Value}; -use tokio::sync::{ - Mutex, - mpsc::{self}, -}; +use magnus::{Error, Integer, RString, Ruby, Value}; +use tokio::sync::{Mutex, Semaphore, mpsc}; use crate::{ - error::{memory_error, mpsc_send_error_to_magnus, wreq_error_to_magnus}, + error::{ + body_sender_borrow_error_to_magnus, body_sender_borrow_mut_error_to_magnus, + body_sender_send_error_to_magnus, closed_body_sender_error, memory_error, + wreq_error_to_magnus, + }, rt, }; +/// Number of chunks buffered when Ruby omits the channel capacity. +const DEFAULT_CHANNEL_CAPACITY: usize = 8; + /// A receiver for streaming HTTP response bodies. pub struct BodyReceiver(Mutex> + Send>>>); -/// A sender for streaming HTTP request bodies. +/// A bounded producer for a single streaming HTTP request body. +/// +/// The receiving side may be attached to one request. The producer remains +/// writable until [`BodySender::close`] is called or the request drops its +/// receiver. Ruby's GVL protects state access; no [`RefCell`] borrow is kept +/// while request backpressure waits without the GVL. #[magnus::wrap(class = "Wreq::BodySender", free_immediately, size)] -pub struct BodySender(RwLock); +pub struct BodySender(RefCell); +/// Mutable ownership state for both halves of the body channel. struct InnerBodySender { + /// Producing side, removed by [`BodySender::close`]. tx: Option>, + /// Receiving side, removed when the sender is attached to a request. rx: Option>, } +impl InnerBodySender { + /// Return whether the channel can no longer accept body chunks. + fn is_closed(&self) -> bool { + match &self.tx { + Some(tx) => tx.is_closed(), + None => true, + } + } +} + // ===== impl BodyReceiver ===== impl BodyReceiver { @@ -56,47 +79,136 @@ impl BodyReceiver { // ===== impl BodySender ===== impl BodySender { - /// Ruby: `Wreq::Sender.new(capacity = 8)` - pub fn new(args: &[Value]) -> Self { - let capacity: usize = if let Some(v) = args.first() { - usize::try_convert(*v).unwrap_or(8) - } else { - 8 - }; + /// Create a bounded request-body channel. + /// + /// Ruby: `Wreq::BodySender.new(capacity = 8)`. Capacity must be greater + /// than zero and no larger than [`Semaphore::MAX_PERMITS`]. + /// + /// # Errors + /// + /// Returns `TypeError` for a non-Integer capacity and `ArgumentError` for + /// an invalid range or argument count. + pub fn new(ruby: &Ruby, args: &[Value]) -> Result { + let capacity = parse_capacity(ruby, args)?; + + // Create the Tokio channel without allowing an unwind to cross the Ruby FFI boundary. + // + // Known panic conditions are rejected by [`parse_capacity`]. The unwind guard + // remains as a defensive fallback if Tokio adds another channel invariant. + let (tx, rx) = + catch_unwind(|| mpsc::channel(capacity)).map_err(|_| invalid_capacity_error(ruby))?; - let (tx, rx) = mpsc::channel(capacity); - BodySender(RwLock::new(InnerBodySender { + Ok(BodySender(RefCell::new(InnerBodySender { tx: Some(tx), rx: Some(rx), - })) + }))) } - /// Ruby: `push(data)` where data is String or bytes + /// Push a binary chunk, waiting for capacity when the channel is full. + /// + /// Ruby: `push(data)` where `data` is a String. + /// + /// # Errors + /// + /// Returns `IOError` after either channel side has closed. An interrupted + /// wait retains the existing `Wreq::InterruptError` behavior. pub fn push(rb_self: &Self, data: RString) -> Result<(), Error> { - let bytes = data.to_bytes(); - let inner = rb_self.0.read().unwrap(); - if let Some(ref tx) = inner.tx { - rt::try_block_on(tx.send(bytes), mpsc_send_error_to_magnus)?; - } - Ok(()) + // Clone during the shared borrow, then release it before waiting + // for capacity. Request attachment needs a mutable borrow. + let tx = match &rb_self.read_inner()?.tx { + Some(tx) if !tx.is_closed() => tx.clone(), + _ => return Err(closed_body_sender_error()), + }; + + rt::try_block_on(tx.send(data.to_bytes()), body_sender_send_error_to_magnus) } - /// Ruby: `close` to close the sender - pub fn close(&self) { - let mut inner = self.0.write().unwrap(); + /// Close the producing side while retaining the receiver and queued chunks. + /// + /// Calling this method more than once has no additional effect. + /// + /// # Errors + /// + /// Returns `Wreq::BodyError` if the internal state is already borrowed. + pub fn close(&self) -> Result<(), Error> { + let mut inner = self.write_inner()?; inner.tx.take(); - inner.rx.take(); + Ok(()) + } + + /// Return whether this sender can no longer accept body chunks. + /// + /// # Errors + /// + /// Returns `Wreq::BodyError` if the internal state is already borrowed. + pub fn is_closed(&self) -> Result { + self.read_inner().map(|r| r.is_closed()) + } + + /// Borrow the channel state without panicking on accidental re-entry. + fn read_inner(&self) -> Result, Error> { + self.0 + .try_borrow() + .map_err(body_sender_borrow_error_to_magnus) } + + /// Mutably borrow the channel state without panicking on accidental re-entry. + fn write_inner(&self) -> Result, Error> { + self.0 + .try_borrow_mut() + .map_err(body_sender_borrow_mut_error_to_magnus) + } +} + +/// Parse and validate the optional Ruby channel capacity. +/// +/// [`mpsc::channel`] panics for zero or values above +/// [`Semaphore::MAX_PERMITS`], so validation must finish before channel creation. +fn parse_capacity(ruby: &Ruby, args: &[Value]) -> Result { + let value = match args { + [] => return Ok(DEFAULT_CHANNEL_CAPACITY), + [value] => *value, + _ => { + return Err(Error::new( + ruby.exception_arg_error(), + format!( + "wrong number of arguments (given {}, expected 0..1)", + args.len() + ), + )); + } + }; + + let integer = Integer::from_value(value) + .ok_or_else(|| Error::new(ruby.exception_type_error(), "capacity must be an Integer"))?; + let capacity = integer + .to_i64() + .ok() + .and_then(|capacity| usize::try_from(capacity).ok()) + .filter(|capacity| (1..=Semaphore::MAX_PERMITS).contains(capacity)) + .ok_or_else(|| invalid_capacity_error(ruby))?; + + Ok(capacity) +} + +/// Build the synchronous Ruby error used for an invalid channel capacity. +fn invalid_capacity_error(ruby: &Ruby) -> Error { + Error::new( + ruby.exception_arg_error(), + format!("capacity must be between 1 and {}", Semaphore::MAX_PERMITS), + ) } +/// Move the receiving side into one request body. +/// +/// The sender remains available for concurrent producers until it is closed or +/// the returned stream is dropped. A second attachment returns `Wreq::MemoryError`. impl TryFrom<&BodySender> for ReceiverStream { type Error = magnus::Error; fn try_from(sender: &BodySender) -> Result { sender - .0 - .write() - .unwrap() + .write_inner()? .rx .take() .map(ReceiverStream::new) diff --git a/src/error.rs b/src/error.rs index 215c92a..94f01b7 100644 --- a/src/error.rs +++ b/src/error.rs @@ -1,3 +1,5 @@ +use std::cell::{BorrowError, BorrowMutError}; + use magnus::{ Error as MagnusError, RModule, Ruby, exception::ExceptionClass, prelude::*, value::Lazy, }; @@ -93,11 +95,32 @@ pub fn no_block_given_error() -> MagnusError { ) } -/// Map [`tokio::sync::mpsc::error::SendError`] to corresponding [`magnus::Error`] -pub fn mpsc_send_error_to_magnus(err: SendError) -> MagnusError { +/// Build an `IOError` for writes to a closed request-body sender. +pub fn closed_body_sender_error() -> MagnusError { + MagnusError::new(ruby!().exception_io_error(), "closed body sender") +} + +/// Map a failed body-channel send to `IOError`. +pub fn body_sender_send_error_to_magnus(err: SendError) -> MagnusError { + MagnusError::new( + ruby!().exception_io_error(), + format!("closed body sender: {err}"), + ) +} + +/// Map an immutable sender-state borrow failure to `Wreq::BodyError`. +pub fn body_sender_borrow_error_to_magnus(err: BorrowError) -> MagnusError { + MagnusError::new( + ruby!().get_inner(&BODY_ERROR), + format!("body sender state is unavailable: {err}"), + ) +} + +/// Map a mutable sender-state borrow failure to `Wreq::BodyError`. +pub fn body_sender_borrow_mut_error_to_magnus(err: BorrowMutError) -> MagnusError { MagnusError::new( ruby!().get_inner(&BODY_ERROR), - format!("failed to send body chunk: {}", err), + format!("body sender state is unavailable: {err}"), ) } diff --git a/test/body_sender_test.rb b/test/body_sender_test.rb new file mode 100644 index 0000000..0f6b875 --- /dev/null +++ b/test/body_sender_test.rb @@ -0,0 +1,246 @@ +require "test_helper" +require "open3" +require "rbconfig" +require "socket" + +class BodySenderTest < Minitest::Test + def test_valid_capacity_creates_open_sender + sender = Wreq::BodySender.new(2) + + refute_predicate sender, :closed? + ensure + sender&.close + end + + def test_omitted_capacity_defaults_to_eight + sender = Wreq::BodySender.new + progress = Queue.new + producer = Thread.new do + 9.times do |index| + sender.push(index.to_s) + progress << index + end + sender.close + end + producer.report_on_exception = false + + assert wait_until { progress.size >= 8 }, "the default channel should buffer eight chunks" + assert_equal 8, progress.size + assert_predicate producer, :alive?, "the ninth push should wait for channel capacity" + + with_request_body_server do |url, request_thread| + response = Wreq.post(url, body: sender) + + assert_equal 200, response.code + assert producer.join(5), "the ninth push should finish once the request drains" + assert_equal "012345678", request_thread.value + end + ensure + producer&.kill if producer&.alive? + producer&.join(5) + sender&.close + end + + def test_non_positive_and_excessive_capacities_raise_argument_error + [0, -1, 2**256].each do |capacity| + error = assert_raises(ArgumentError) { Wreq::BodySender.new(capacity) } + assert_match(/capacity/, error.message) + end + end + + def test_non_integer_capacities_raise_type_error + [1.0, "not an integer"].each do |capacity| + error = assert_raises(TypeError) { Wreq::BodySender.new(capacity) } + assert_match(/capacity/, error.message) + end + end + + def test_too_many_constructor_arguments_raise_argument_error + assert_raises(ArgumentError) { Wreq::BodySender.new(1, 2) } + end + + def test_close_is_idempotent_and_updates_closed_state + sender = Wreq::BodySender.new + + refute_predicate sender, :closed? + assert_nil sender.close + assert_predicate sender, :closed? + assert_nil sender.close + end + + def test_push_after_close_raises_io_error + sender = Wreq::BodySender.new + sender.close + + error = assert_raises(IOError) { sender.push("data") } + assert_equal "closed body sender", error.message + end + + def test_queued_chunks_survive_close_before_request_attachment + sender = Wreq::BodySender.new(2) + sender.push("queued-") + sender.push("body") + sender.close + + with_request_body_server do |url, request_thread| + response = Wreq.post(url, body: sender) + + assert_equal 200, response.code + assert_equal "queued-body", request_thread.value + end + end + + def test_valid_capacity_keeps_backpressure_until_request_drains + sender = Wreq::BodySender.new(1) + sender.push("first-") + producer = Thread.new do + sender.push("second") + sender.close + end + producer.report_on_exception = false + + sleep 0.05 + assert_predicate producer, :alive?, "the second push should wait for channel capacity" + + with_request_body_server do |url, request_thread| + response = Wreq.post(url, body: sender) + + assert_equal 200, response.code + assert producer.join(5), "the blocked producer should finish once the request drains" + assert_equal "first-second", request_thread.value + end + ensure + producer&.kill if producer&.alive? + producer&.join(5) + sender&.close + end + + def test_receiver_termination_closes_sender + sender = Wreq::BodySender.new(1) + sender.push("data") + + with_reset_server do |url| + assert_raises(StandardError) { Wreq.post(url, body: sender, timeout: 2) } + end + + assert wait_until { sender.closed? }, "sender should close after the request drops its receiver" + assert_raises(IOError) { sender.push("more") } + ensure + sender&.close + end + + def test_zero_capacity_regression_exits_subprocess_normally + lib_dir = File.expand_path("../lib", __dir__) + script = <<~RUBY + require "wreq" + + begin + Wreq::BodySender.new(0) + rescue ArgumentError => error + warn "\#{error.class}: \#{error.message}" + exit 0 + end + + warn "zero capacity did not raise" + exit 2 + RUBY + + _stdout, stderr, status = Open3.capture3(RbConfig.ruby, "-I", lib_dir, "-e", script) + + assert status.success?, "subprocess failed with #{status.inspect}: #{stderr}" + assert_match(/ArgumentError:.*capacity/, stderr) + refute_match(/panicked|mpsc bounded channel/i, stderr) + end + + private + + def wait_until(timeout: 2) + deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + timeout + loop do + return true if yield + return false if Process.clock_gettime(Process::CLOCK_MONOTONIC) >= deadline + + sleep 0.01 + end + end + + def with_request_body_server + server = TCPServer.new("127.0.0.1", 0) + port = server.addr[1] + request_thread = Thread.new do + socket = server.accept + begin + socket.gets + headers = read_headers(socket) + body = read_request_body(socket, headers) + socket.write "HTTP/1.1 200 OK\r\n" + socket.write "Content-Length: 0\r\n" + socket.write "Connection: close\r\n\r\n" + body + ensure + socket.close unless socket.closed? + end + ensure + server.close unless server.closed? + end + request_thread.report_on_exception = false + + yield "http://127.0.0.1:#{port}/", request_thread + ensure + server&.close unless server&.closed? + request_thread&.join(5) + end + + def with_reset_server + server = TCPServer.new("127.0.0.1", 0) + port = server.addr[1] + thread = Thread.new do + socket = server.accept + socket.close + ensure + server.close unless server.closed? + end + thread.report_on_exception = false + + yield "http://127.0.0.1:#{port}/" + ensure + server&.close unless server&.closed? + thread&.join(5) + end + + def read_headers(socket) + headers = {} + while (line = socket.gets) + break if line == "\r\n" + + name, value = line.split(":", 2) + headers[name.downcase] = value.strip + end + headers + end + + def read_request_body(socket, headers) + if headers.fetch("transfer-encoding", "").downcase.include?("chunked") + read_chunked_body(socket) + else + socket.read(headers.fetch("content-length", "0").to_i) + end + end + + def read_chunked_body(socket) + body = "".b + loop do + size_line = socket.gets or raise EOFError, "missing chunk size" + size = Integer(size_line.split(";", 2).first, 16) + break if size.zero? + + body << socket.read(size) + raise IOError, "missing chunk terminator" unless socket.read(2) == "\r\n" + end + + while (line = socket.gets) + break if line == "\r\n" + end + body + end +end