[rb] fix silent hang on oversized WebSocket frames

rb/lib/selenium/webdriver/common/websocket_connection.rb

@@ -21,6 +21,7 @@
 
 module Selenium
   module WebDriver
+    # @api private
     class WebSocketConnection
       CONNECTION_ERRORS = [
         Errno::ECONNRESET, # connection is aborted (browser process was killed)
@@ -35,6 +36,15 @@ class WebSocketConnection
 
       MAX_LOG_MESSAGE_SIZE = 9999
 
+      # websocket-ruby defaults to a 20MB limit and silently drops larger
+      # frames, which can stall the listener. CDP payloads (e.g. large data:
+      # URLs) can exceed that, so raise the ceiling for our connections.
+      # The gem only exposes the limit as process-global state, so it is
+      # raised (never lowered) and not restored on close - concurrent
+      # connections share the value. Other bindings rely on their own
+      # websocket clients' limits; this constant only affects websocket-ruby.
+      MAX_FRAME_SIZE = 100 * 1024 * 1024 # 100MB
+
       def initialize(url:)
         @callback_threads = ThreadGroup.new
 
@@ -46,26 +56,19 @@ def initialize(url:)
         @session_id = nil
         @url = url
 
+        apply_frame_size_limit
         process_handshake
         @socket_thread = attach_socket_listener
       end
 
+      # Idempotent: the listener may already have initiated shutdown (see
+      # #frame_dropped?), so always close the socket and join threads rather
+      # than short-circuiting on @closing.
       def close
-        @closing_mtx.synchronize do
-          return if @closing
-
-          @closing = true
-        end
-
-        begin
-          socket.close
-        rescue *CONNECTION_ERRORS => e
-          WebDriver.logger.debug "WebSocket listener closed: #{e.class}: #{e.message}", id: :ws
-          # already closed
-        end
+        close_socket
 
         # Let threads unwind instead of calling exit
-        @socket_thread&.join(0.5)
+        @socket_thread&.join(0.5) unless @socket_thread == Thread.current
         @callback_threads.list.each do |thread|
           thread.join(0.5)
         rescue StandardError => e
@@ -97,6 +100,9 @@ def remove_callback(event, id)
       end
 
       def send_cmd(**payload)
+        # IOError to match what writing to an already-closed socket raises
+        raise IOError, 'WebSocket connection is closed' if @closing
+
         id = next_id
         data = payload.merge(id: id)
         WebDriver.logger.debug "WebSocket -> #{data}"[...MAX_LOG_MESSAGE_SIZE], id: :ws
@@ -109,7 +115,11 @@ def send_cmd(**payload)
           raise e, "WebSocket is closed (#{e.class}: #{e.message})"
         end
 
-        wait.until { @messages_mtx.synchronize { messages.delete(id) } }
+        wait.until do
+          raise IOError, 'WebSocket connection closed while waiting for a response' if @closing
+
+          @messages_mtx.synchronize { messages.delete(id) }
+        end
       end
 
       private
@@ -132,22 +142,51 @@ def attach_socket_listener
 
             incoming_frame << socket.readpartial(1024)
 
-            while (frame = incoming_frame.next)
-              break if @closing
-
-              message = process_frame(frame)
-              next unless message['method']
-
-              @messages_mtx.synchronize { callbacks[message['method']].dup }.each do |callback|
-                @callback_threads.add(callback_thread(message['params'], &callback))
-              end
-            end
+            process_incoming_frames
+            break if frame_dropped?
           end
         rescue *CONNECTION_ERRORS, WebSocket::Error => e
           WebDriver.logger.debug "WebSocket listener closed: #{e.class}: #{e.message}", id: :ws
         end
       end
 
+      def close_socket
+        @closing_mtx.synchronize { @closing = true }
+        socket.close
+      rescue *CONNECTION_ERRORS => e
+        WebDriver.logger.debug "WebSocket socket closed: #{e.class}: #{e.message}", id: :ws
+      end
+
+      def process_incoming_frames
+        while (frame = incoming_frame.next)
+          break if @closing
+
+          message = process_frame(frame)
+          next unless message['method']
+
+          @messages_mtx.synchronize { callbacks[message['method']].dup }.each do |callback|
+            @callback_threads.add(callback_thread(message['params'], &callback))
+          end
+        end
+      end
+
+      # True when the buffered frame could not be decoded (e.g. exceeds MAX_FRAME_SIZE).
+      # websocket-ruby swallows the error and keeps returning nil, so surface it and close
+      # the connection here instead of leaving a dead listener on an open socket.
+      def frame_dropped?
+        return false unless incoming_frame.error?
+
+        WebDriver.logger.error("WebSocket frame dropped (#{incoming_frame.error}); if payloads can legitimately " \
+                               "exceed #{WebSocket.max_frame_size} bytes, set WebSocket.max_frame_size " \
+                               'to a higher value', id: :ws)
+        close_socket
+        true
+      end
+
+      def apply_frame_size_limit
+        WebSocket.max_frame_size = MAX_FRAME_SIZE if WebSocket.max_frame_size < MAX_FRAME_SIZE
+      end
+
       def incoming_frame
         @incoming_frame ||= WebSocket::Frame::Incoming::Client.new(version: ws.version)
       end
@@ -204,8 +243,7 @@ def ws
       end
 
       def next_id
-        @id ||= 0
-        @id += 1
+        @id = (@id || 0) + 1
       end
     end # BiDi
   end # WebDriver

rb/sig/gems/websocket/websocket.rbs

@@ -2,6 +2,9 @@ module WebSocket
   DEFAULT_VERSION: Integer
   ROOT: String
 
+  def self.max_frame_size: () -> Integer
+  def self.max_frame_size=: (Integer) -> Integer
+
   # Defining autoloaded classes as interface or class is dependent on their implementation,
   # which is not provided. Here, we'll declare them as modules, but they might need to be
   # updated according to their actual definitions.

rb/sig/lib/selenium/webdriver/common/websocket_connection.rbs

@@ -53,6 +53,8 @@ module Selenium
 
       MAX_LOG_MESSAGE_SIZE: Integer
 
+      MAX_FRAME_SIZE: Integer
+
       def initialize: (url: String) -> void
 
       def add_callback: (untyped event) { () -> void } -> Integer
@@ -77,6 +79,14 @@ module Selenium
 
       def attach_socket_listener: () -> untyped
 
+      def close_socket: () -> untyped
+
+      def process_incoming_frames: () -> untyped
+
+      def frame_dropped?: () -> bool
+
+      def apply_frame_size_limit: () -> void
+
       def incoming_frame: () -> untyped
 
       def process_frame: (untyped frame) -> (Hash[untyped, untyped] | untyped)

rb/spec/unit/selenium/webdriver/common/websocket_connection_spec.rb

@@ -0,0 +1,111 @@
+# frozen_string_literal: true
+
+# Licensed to the Software Freedom Conservancy (SFC) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The SFC licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+require File.expand_path('../spec_helper', __dir__)
+require 'stringio'
+
+module Selenium
+  module WebDriver
+    describe WebSocketConnection do
+      # Build an instance without opening a socket so the frame-handling logic
+      # can be exercised in isolation.
+      subject(:connection) { described_class.allocate }
+
+      around do |example|
+        original = WebSocket.max_frame_size
+        example.call
+      ensure
+        WebSocket.max_frame_size = original
+      end
+
+      describe '#apply_frame_size_limit' do
+        it 'raises the global limit when it is below the Selenium default' do
+          WebSocket.max_frame_size = 1
+
+          connection.send(:apply_frame_size_limit)
+
+          expect(WebSocket.max_frame_size).to eq(described_class::MAX_FRAME_SIZE)
+        end
+
+        it 'leaves a larger user-configured limit untouched' do
+          larger = described_class::MAX_FRAME_SIZE * 2
+          WebSocket.max_frame_size = larger
+
+          connection.send(:apply_frame_size_limit)
+
+          expect(WebSocket.max_frame_size).to eq(larger)
+        end
+      end
+
+      describe '#close' do
+        let(:socket) { StringIO.new }
+
+        before do
+          connection.instance_variable_set(:@closing_mtx, Mutex.new)
+          connection.instance_variable_set(:@socket, socket)
+          connection.instance_variable_set(:@callback_threads, ThreadGroup.new)
+        end
+
+        it 'still closes the socket when the listener already initiated shutdown' do
+          connection.instance_variable_set(:@closing, true)
+
+          connection.close
+
+          expect(socket).to be_closed
+        end
+      end
+
+      describe '#send_cmd' do
+        it 'fails fast when the connection is closing' do
+          connection.instance_variable_set(:@closing, true)
+
+          expect { connection.send_cmd(method: 'foo') }
+            .to raise_error(IOError, /closed/)
+        end
+      end
+
+      describe '#frame_dropped?' do
+        let(:incoming_frame) { WebSocket::Frame::Incoming::Client.new(version: 13) }
+        let(:socket) { StringIO.new }
+
+        before do
+          connection.instance_variable_set(:@incoming_frame, incoming_frame)
+          connection.instance_variable_set(:@closing_mtx, Mutex.new)
+          connection.instance_variable_set(:@socket, socket)
+        end
+
+        it 'is false when there is no decoding error' do
+          expect(connection.send(:frame_dropped?)).to be(false)
+        end
+
+        it 'is true and closes the connection when a frame exceeds the maximum size' do
+          WebSocket.max_frame_size = 1
+          raw = WebSocket::Frame::Outgoing::Server.new(version: 13, data: 'a' * 1024, type: 'text').to_s
+          incoming_frame << raw
+          incoming_frame.next
+
+          expect(connection.send(:frame_dropped?)).to be(true)
+          expect(incoming_frame.error?).to be(true)
+          expect(connection.instance_variable_get(:@closing)).to be(true)
+          expect(socket).to be_closed
+        end
+      end
+    end
+  end
+end