Merge commit from fork

Co-authored-by: Alida W <alida.weatherbee@workos.com>

Author: gjtorikian

lib/workos/base_client.rb

@@ -137,7 +137,20 @@ def execute_request(request:, request_options: nil)
         loggable_path = redact_path(request.path)
         log(:debug, "request start", method: request.method, path: loggable_path, attempt: attempt + 1)
         http = connection_for(base, timeout)
-        response = http.request(request)
+        request_completed = false
+        begin
+          response = http.request(request)
+          request_completed = true
+        ensure
+          # Any exit from #request other than a returned response can leave
+          # the socket mid-stream: a connection error the rescue below knows
+          # about, one it doesn't (OpenSSL::SSL::SSLError,
+          # Net::HTTPBadResponse), or a non-local exit such as an
+          # application-level Timeout.timeout or Thread#kill. A half-read
+          # socket handed back to the pool desyncs the *next* request on this
+          # thread, so drop it here rather than in the rescue.
+          evict_connection(base) unless request_completed
+        end
         return response if response.is_a?(Net::HTTPSuccess)
 
         if attempt < retries && retryable?(response)

test/workos/test_base_client.rb

@@ -82,6 +82,24 @@ def finish
     end
   end
 
+  # A pooled connection that either returns a canned response or raises when
+  # driven, so execute_request can be exercised without real TCP+TLS.
+  class StubConnection < FakeConnection
+    attr_accessor :read_timeout, :open_timeout
+
+    def initialize(response: nil, error: nil, **kwargs)
+      super(**kwargs)
+      @response = response
+      @error = error
+    end
+
+    def request(_request)
+      raise @error if @error
+
+      @response
+    end
+  end
+
   def setup
     @client = WorkOS::BaseClient.new(api_key: "sk_test_123", max_retries: 1)
   end
@@ -206,6 +224,143 @@ def test_evict_connection_removes_matching_pooled_connections
     refute keep.finished
   end
 
+  # An exception execute_request's rescue clause doesn't list still leaves the
+  # socket mid-stream, so the connection must not survive in the pool for the
+  # next request on this thread to pick up.
+  def test_unlisted_request_error_evicts_the_pooled_connection
+    conn = StubConnection.new(error: Net::HTTPBadResponse.new("wrong version"))
+    cache = @client.send(:thread_connections)
+    cache["https:api.workos.com:443:30"] = conn
+
+    assert_raises(Net::HTTPBadResponse) do
+      @client.execute_request(request: Net::HTTP::Get.new("/things"))
+    end
+
+    refute cache.key?("https:api.workos.com:443:30"),
+      "a connection whose request did not complete must not stay pooled"
+    assert conn.finished
+  end
+
+  def test_completed_request_keeps_the_connection_pooled
+    conn = StubConnection.new(response: Net::HTTPOK.new("1.1", "200", "OK"))
+    cache = @client.send(:thread_connections)
+    cache["https:api.workos.com:443:30"] = conn
+
+    @client.execute_request(request: Net::HTTP::Get.new("/things"))
+
+    assert cache.key?("https:api.workos.com:443:30")
+    refute conn.finished
+  end
+
+  # Raised asynchronously into the worker thread to stand in for a caller-side
+  # abort: Timeout.timeout, Thread#kill, a signal handler unwinding the stack.
+  # It descends from Exception rather than StandardError so that nothing in
+  # execute_request can catch it — the `ensure` is the only cleanup that runs.
+  class AbortSignal < Exception; end # standard:disable Lint/InheritException
+
+  # End-to-end regression test over a real keep-alive socket, for the failure
+  # the eviction actually prevents: request one is abandoned mid-flight, the
+  # server then writes response one onto that socket, and request two on the
+  # same thread reads those stale bytes as if they were its own response.
+  #
+  # Before the fix, request two sees marker "one" (or a mangled response) off
+  # the pooled socket. After it, the abandoned socket is closed and the server
+  # accepts a fresh connection for request two.
+  def test_aborted_request_does_not_leak_its_response_to_the_next_request
+    WebMock.disable!
+    server = TCPServer.new("127.0.0.1", 0)
+    port = server.addr[1]
+    client = WorkOS::BaseClient.new(
+      api_key: "sk_test_123",
+      base_url: "http://127.0.0.1:#{port}",
+      timeout: 5,
+      max_retries: 0
+    )
+
+    request_one_read = Queue.new
+    release_response_one = Queue.new
+    response_one_written = Queue.new
+    aborted = Queue.new
+    connections = Queue.new
+
+    server_thread = Thread.new do
+      # Connection one: read the request, then hold the response back until
+      # the client has been aborted and has unwound.
+      first = server.accept
+      connections << first
+      read_http_request(first)
+      request_one_read << true
+      release_response_one.pop
+      begin
+        write_http_response(first, {marker: "one"})
+      rescue Errno::EPIPE, Errno::ECONNRESET, IOError
+        # Expected once the fix closes the abandoned socket.
+      end
+      response_one_written << true
+
+      # Connection two: a fresh accept, which only completes because the
+      # client did not reuse the socket above.
+      second = server.accept
+      connections << second
+      read_http_request(second)
+      write_http_response(second, {marker: "two"})
+    end
+
+    worker = Thread.new do
+      begin
+        client.execute_request(request: Net::HTTP::Get.new("/one"))
+      rescue AbortSignal
+        # The caller unwinds but the thread survives and goes on to serve more
+        # work, the way a Puma or Solid Queue worker does. execute_request's
+        # `ensure` has already run by the time this body executes.
+        aborted << true
+        response_one_written.pop
+      end
+      response = client.execute_request(request: Net::HTTP::Get.new("/two"))
+      JSON.parse(response.body)["marker"]
+    end
+
+    marker = Timeout.timeout(15) do
+      request_one_read.pop
+      worker.raise(AbortSignal)
+      aborted.pop
+      release_response_one << true
+      worker.value
+    end
+
+    assert_equal "two", marker,
+      "request two read the abandoned socket's response instead of its own"
+  ensure
+    server_thread&.kill
+    worker&.kill
+    until connections.nil? || connections.empty?
+      socket = connections.pop
+      socket.close unless socket.closed?
+    end
+    server&.close
+    WebMock.enable!
+  end
+
+  def read_http_request(socket)
+    socket.gets # request line
+    loop do
+      line = socket.gets
+      break if line.nil? || line == "\r\n"
+    end
+  end
+
+  def write_http_response(socket, body)
+    payload = JSON.generate(body)
+    socket.write(
+      "HTTP/1.1 200 OK\r\n" \
+      "Content-Type: application/json\r\n" \
+      "Content-Length: #{payload.bytesize}\r\n" \
+      "Connection: keep-alive\r\n" \
+      "\r\n#{payload}"
+    )
+    socket.flush
+  end
+
   # Regression test for https://github.com/workos/workos-ruby/issues/496
   #
   # The connection cache must be isolated per thread. A previous version