diff --git a/rb/lib/selenium/webdriver/common/websocket_connection.rb b/rb/lib/selenium/webdriver/common/websocket_connection.rb index 0774a46918009..74eb868c075e6 100644 --- a/rb/lib/selenium/webdriver/common/websocket_connection.rb +++ b/rb/lib/selenium/webdriver/common/websocket_connection.rb @@ -35,6 +35,11 @@ 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. + MAX_FRAME_SIZE = 100 * 1024 * 1024 # 100MB + def initialize(url:) @callback_threads = ThreadGroup.new @@ -46,6 +51,7 @@ def initialize(url:) @session_id = nil @url = url + apply_frame_size_limit process_handshake @socket_thread = attach_socket_listener end @@ -57,12 +63,7 @@ def close @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) @@ -132,22 +133,50 @@ 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}): exceeds max_frame_size " \ + "(#{WebSocket.max_frame_size} bytes). Raise WebSocket.max_frame_size=", 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 diff --git a/rb/sig/gems/websocket/websocket.rbs b/rb/sig/gems/websocket/websocket.rbs index 03cbfaecf3ef4..ef61442212160 100644 --- a/rb/sig/gems/websocket/websocket.rbs +++ b/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. diff --git a/rb/sig/lib/selenium/webdriver/common/websocket_connection.rbs b/rb/sig/lib/selenium/webdriver/common/websocket_connection.rbs index 5d763fc23542c..ae44a109ed7bd 100644 --- a/rb/sig/lib/selenium/webdriver/common/websocket_connection.rbs +++ b/rb/sig/lib/selenium/webdriver/common/websocket_connection.rbs @@ -35,6 +35,8 @@ module Selenium MAX_LOG_MESSAGE_SIZE: Integer + MAX_FRAME_SIZE: Integer + def initialize: (url: String) -> void def add_callback: (untyped event) { () -> void } -> Integer @@ -59,6 +61,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) diff --git a/rb/spec/unit/selenium/webdriver/common/websocket_connection_spec.rb b/rb/spec/unit/selenium/webdriver/common/websocket_connection_spec.rb new file mode 100644 index 0000000000000..7ace3f6ff7b46 --- /dev/null +++ b/rb/spec/unit/selenium/webdriver/common/websocket_connection_spec.rb @@ -0,0 +1,84 @@ +# 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 '#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