From 5e52f03d01d54afdc54b66d9be740fa775322069 Mon Sep 17 00:00:00 2001 From: "William T. Nelson" <35801+wtn@users.noreply.github.com> Date: Wed, 15 Apr 2026 07:21:41 -0500 Subject: [PATCH 1/4] Optimize frame reading. Co-authored-by: Claude --- lib/protocol/websocket/connection.rb | 6 +++- lib/protocol/websocket/frame.rb | 44 +++++++++++++++++---------- lib/protocol/websocket/framer.rb | 31 +++++++++++-------- test/protocol/websocket/connection.rb | 31 +++++++++++++++++++ test/protocol/websocket/frame.rb | 9 ++++++ test/protocol/websocket/framer.rb | 18 +++++++++++ 6 files changed, 109 insertions(+), 30 deletions(-) diff --git a/lib/protocol/websocket/connection.rb b/lib/protocol/websocket/connection.rb index 45bc054..0632c90 100644 --- a/lib/protocol/websocket/connection.rb +++ b/lib/protocol/websocket/connection.rb @@ -281,7 +281,11 @@ def write(message, **options) # The default implementation for reading a message buffer. This is used by the {#reader} interface. def unpack_frames(frames) - frames.map(&:unpack).join("") + if frames.size == 1 + frames[0].unpack + else + frames.map(&:unpack).join("") + end end # Read a message from the connection. If an error occurs while reading the message, the connection will be closed. diff --git a/lib/protocol/websocket/frame.rb b/lib/protocol/websocket/frame.rb index e9340d0..8858505 100644 --- a/lib/protocol/websocket/frame.rb +++ b/lib/protocol/websocket/frame.rb @@ -173,15 +173,15 @@ def apply(connection) # @returns [Array] A tuple of `[finished, flags, opcode]`. # @raises [ProtocolError] If the opcode is a reserved non-control or control opcode. def self.parse_header(buffer) - byte = buffer.unpack("C").first + byte = buffer.getbyte(0) finished = (byte & 0b1000_0000 != 0) flags = (byte & 0b0111_0000) >> 4 opcode = byte & 0b0000_1111 - if (0x3 .. 0x7).include?(opcode) + if opcode >= 0x3 && opcode <= 0x7 raise ProtocolError, "Non-control opcode = #{opcode} is reserved!" - elsif (0xB .. 0xF).include?(opcode) + elsif opcode >= 0xB raise ProtocolError, "Control opcode = #{opcode} is reserved!" end @@ -197,12 +197,14 @@ def self.parse_header(buffer) # @returns [Frame] The fully read and populated frame. # @raises [ProtocolError] If the frame violates protocol constraints. # @raises [EOFError] If the stream ends unexpectedly. - def self.read(finished, flags, opcode, stream, maximum_frame_size) - buffer = stream.read(1) or raise EOFError, "Could not read header!" - byte = buffer.unpack("C").first + def self.read(finished, flags, opcode, stream, maximum_frame_size, second_byte = nil) + unless second_byte + buffer = stream.read(1) or raise EOFError, "Could not read header!" + second_byte = buffer.getbyte(0) + end - mask = (byte & 0b1000_0000 != 0) - length = byte & 0b0111_1111 + mask = (second_byte & 0b1000_0000 != 0) + length = second_byte & 0b0111_1111 if opcode & 0x8 != 0 if length > 125 @@ -213,21 +215,31 @@ def self.read(finished, flags, opcode, stream, maximum_frame_size) end if length == 126 - buffer = stream.read(2) or raise EOFError, "Could not read length!" - length = buffer.unpack("n").first + if mask + buffer = stream.read(6) or raise EOFError, "Could not read length and mask!" + length = buffer.unpack1("n") + mask = buffer.byteslice(2, 4) + else + buffer = stream.read(2) or raise EOFError, "Could not read length!" + length = buffer.unpack1("n") + end elsif length == 127 - buffer = stream.read(8) or raise EOFError, "Could not read length!" - length = buffer.unpack("Q>").first + if mask + buffer = stream.read(12) or raise EOFError, "Could not read length and mask!" + length = buffer.unpack1("Q>") + mask = buffer.byteslice(8, 4) + else + buffer = stream.read(8) or raise EOFError, "Could not read length!" + length = buffer.unpack1("Q>") + end + elsif mask + mask = stream.read(4) or raise EOFError, "Could not read mask!" end if length > maximum_frame_size raise ProtocolError, "Invalid payload length: #{length} > #{maximum_frame_size}!" end - if mask - mask = stream.read(4) or raise EOFError, "Could not read mask!" - end - payload = stream.read(length) or raise EOFError, "Could not read payload!" if payload.bytesize != length diff --git a/lib/protocol/websocket/framer.rb b/lib/protocol/websocket/framer.rb index 9ddfdb7..ab760e2 100644 --- a/lib/protocol/websocket/framer.rb +++ b/lib/protocol/websocket/framer.rb @@ -50,12 +50,26 @@ def flush # Read a frame from the underlying stream. # @returns [Frame] the frame read from the stream. def read_frame(maximum_frame_size = MAXIMUM_ALLOWED_FRAME_SIZE) - # Read the header: - finished, flags, opcode = read_header + buffer = @stream.read(2) + + unless buffer and buffer.bytesize == 2 + raise EOFError, "Could not read frame header!" + end + + first_byte = buffer.getbyte(0) + + finished = (first_byte & 0b1000_0000 != 0) + flags = (first_byte & 0b0111_0000) >> 4 + opcode = first_byte & 0b0000_1111 + + if opcode >= 0x3 && opcode <= 0x7 + raise ProtocolError, "Non-control opcode = #{opcode} is reserved!" + elsif opcode >= 0xB + raise ProtocolError, "Control opcode = #{opcode} is reserved!" + end - # Read the frame: klass = @frames[opcode] || Frame - frame = klass.read(finished, flags, opcode, @stream, maximum_frame_size) + frame = klass.read(finished, flags, opcode, @stream, maximum_frame_size, buffer.getbyte(1)) return frame end @@ -64,15 +78,6 @@ def read_frame(maximum_frame_size = MAXIMUM_ALLOWED_FRAME_SIZE) def write_frame(frame) frame.write(@stream) end - - # Read the header of the frame. - def read_header - if buffer = @stream.read(1) and buffer.bytesize == 1 - return Frame.parse_header(buffer) - end - - raise EOFError, "Could not read frame header!" - end end end end diff --git a/test/protocol/websocket/connection.rb b/test/protocol/websocket/connection.rb index e3830c0..5e19e6d 100644 --- a/test/protocol/websocket/connection.rb +++ b/test/protocol/websocket/connection.rb @@ -301,6 +301,37 @@ end end + with "masked frames with extended lengths" do + let(:connection) {subject.new(server)} + + it "can handle a masked medium message (length=126 encoding)" do + thread = Thread.new do + frame = Protocol::WebSocket::TextFrame.new(true, mask: true) + frame.pack("a" * 200) + client.write_frame(frame) + end + + message = connection.read + expect(message.size).to be == 200 + expect(message).to be == ("a" * 200) + + thread.join + end + + it "can handle a masked large message (length=127 encoding)" do + thread = Thread.new do + frame = Protocol::WebSocket::TextFrame.new(true, mask: true) + frame.pack("a" * 70_000) + client.write_frame(frame) + end + + message = connection.read + expect(message.size).to be == 70_000 + + thread.join + end + end + with "invalid unicode text message in 3 fragments" do let(:payload1) {"\xce\xba\xe1\xbd\xb9\xcf\x83\xce\xbc\xce\xb5".b} let(:payload2) {"\xf4\x90\x80\x80".b} diff --git a/test/protocol/websocket/frame.rb b/test/protocol/websocket/frame.rb index 44878f2..c6268bc 100644 --- a/test/protocol/websocket/frame.rb +++ b/test/protocol/websocket/frame.rb @@ -75,6 +75,15 @@ subject.read(false, 0, 0, stream, 128) end.to raise_exception(EOFError, message: be =~ /Incorrect payload length: \d+ != \d+!/) end + + it "accepts a pre-read second byte" do + stream = StringIO.new("Hello") + second_byte = 0x05 + + frame = subject.read(true, 0, 0x1, stream, 128, second_byte) + expect(frame.payload).to be == "Hello" + expect(frame.mask).to be == false + end end with ".write" do diff --git a/test/protocol/websocket/framer.rb b/test/protocol/websocket/framer.rb index cf025f6..e4a9549 100644 --- a/test/protocol/websocket/framer.rb +++ b/test/protocol/websocket/framer.rb @@ -15,5 +15,23 @@ framer.read_frame end.to raise_exception(EOFError, message: be =~ /Could not read frame header/) end + + it "rejects reserved non-control opcodes" do + stream.string = "\x83\x00" + stream.rewind + + expect do + framer.read_frame + end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Non-control opcode.*reserved/) + end + + it "rejects reserved control opcodes" do + stream.string = "\x8B\x00" + stream.rewind + + expect do + framer.read_frame + end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Control opcode.*reserved/) + end end end From 362809c0403a62cef14215e93df0bda871dad358 Mon Sep 17 00:00:00 2001 From: Samuel Williams Date: Thu, 16 Apr 2026 15:33:02 +0900 Subject: [PATCH 2/4] Add missing require. --- test/protocol/websocket/connection.rb | 1 + 1 file changed, 1 insertion(+) diff --git a/test/protocol/websocket/connection.rb b/test/protocol/websocket/connection.rb index 5e19e6d..cb83ba4 100644 --- a/test/protocol/websocket/connection.rb +++ b/test/protocol/websocket/connection.rb @@ -5,6 +5,7 @@ # Copyright, 2019, by Soumya. require "socket" +require "securerandom" require "protocol/websocket/connection" describe Protocol::WebSocket::Connection do From 3cc82e1cc1fe988eb45abed3ce6cc078c1ceb0da Mon Sep 17 00:00:00 2001 From: Samuel Williams Date: Thu, 16 Apr 2026 15:42:30 +0900 Subject: [PATCH 3/4] Remove redundant `Frame.parse_header`. --- lib/protocol/websocket/frame.rb | 20 -------------------- test/protocol/websocket/frame.rb | 14 -------------- 2 files changed, 34 deletions(-) diff --git a/lib/protocol/websocket/frame.rb b/lib/protocol/websocket/frame.rb index 8858505..7711044 100644 --- a/lib/protocol/websocket/frame.rb +++ b/lib/protocol/websocket/frame.rb @@ -168,26 +168,6 @@ def apply(connection) connection.receive_frame(self) end - # Parse the first byte of a frame header to extract FIN, RSV flags, and opcode. - # @parameter buffer [String] A 1-byte binary string. - # @returns [Array] A tuple of `[finished, flags, opcode]`. - # @raises [ProtocolError] If the opcode is a reserved non-control or control opcode. - def self.parse_header(buffer) - byte = buffer.getbyte(0) - - finished = (byte & 0b1000_0000 != 0) - flags = (byte & 0b0111_0000) >> 4 - opcode = byte & 0b0000_1111 - - if opcode >= 0x3 && opcode <= 0x7 - raise ProtocolError, "Non-control opcode = #{opcode} is reserved!" - elsif opcode >= 0xB - raise ProtocolError, "Control opcode = #{opcode} is reserved!" - end - - return finished, flags, opcode - end - # Read a full frame from the stream given pre-parsed header fields. # @parameter finished [Boolean] Whether the FIN bit was set. # @parameter flags [Integer] The RSV flag bits. diff --git a/test/protocol/websocket/frame.rb b/test/protocol/websocket/frame.rb index c6268bc..f8fc352 100644 --- a/test/protocol/websocket/frame.rb +++ b/test/protocol/websocket/frame.rb @@ -29,20 +29,6 @@ end end - with ".parse_header" do - it "rejects reserved non-control opcodes" do - expect do - subject.parse_header("\x03\x00") - end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Non-control opcode.*reserved/) - end - - it "rejects reserved control opcodes" do - expect do - subject.parse_header("\x0F\x00") - end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Control opcode.*reserved/) - end - end - with ".read" do it "rejects invalid control frame payload length" do stream = StringIO.new("\xFF") From f9fb1e2a72f5fb31cd99c982988a709d6f38a1f3 Mon Sep 17 00:00:00 2001 From: Samuel Williams Date: Thu, 16 Apr 2026 17:36:41 +0900 Subject: [PATCH 4/4] Move all read logic to `Framer`. --- lib/protocol/websocket/frame.rb | 61 ------------------------------- lib/protocol/websocket/framer.rb | 48 ++++++++++++++++++++++-- test/protocol/websocket/frame.rb | 43 ---------------------- test/protocol/websocket/framer.rb | 51 ++++++++++++++++++++++++++ 4 files changed, 96 insertions(+), 107 deletions(-) diff --git a/lib/protocol/websocket/frame.rb b/lib/protocol/websocket/frame.rb index 7711044..70b0f8d 100644 --- a/lib/protocol/websocket/frame.rb +++ b/lib/protocol/websocket/frame.rb @@ -168,67 +168,6 @@ def apply(connection) connection.receive_frame(self) end - # Read a full frame from the stream given pre-parsed header fields. - # @parameter finished [Boolean] Whether the FIN bit was set. - # @parameter flags [Integer] The RSV flag bits. - # @parameter opcode [Integer] The frame opcode. - # @parameter stream [IO] The stream to read from. - # @parameter maximum_frame_size [Integer] The maximum allowed payload size in bytes. - # @returns [Frame] The fully read and populated frame. - # @raises [ProtocolError] If the frame violates protocol constraints. - # @raises [EOFError] If the stream ends unexpectedly. - def self.read(finished, flags, opcode, stream, maximum_frame_size, second_byte = nil) - unless second_byte - buffer = stream.read(1) or raise EOFError, "Could not read header!" - second_byte = buffer.getbyte(0) - end - - mask = (second_byte & 0b1000_0000 != 0) - length = second_byte & 0b0111_1111 - - if opcode & 0x8 != 0 - if length > 125 - raise ProtocolError, "Invalid control frame payload length: #{length} > 125!" - elsif !finished - raise ProtocolError, "Fragmented control frame!" - end - end - - if length == 126 - if mask - buffer = stream.read(6) or raise EOFError, "Could not read length and mask!" - length = buffer.unpack1("n") - mask = buffer.byteslice(2, 4) - else - buffer = stream.read(2) or raise EOFError, "Could not read length!" - length = buffer.unpack1("n") - end - elsif length == 127 - if mask - buffer = stream.read(12) or raise EOFError, "Could not read length and mask!" - length = buffer.unpack1("Q>") - mask = buffer.byteslice(8, 4) - else - buffer = stream.read(8) or raise EOFError, "Could not read length!" - length = buffer.unpack1("Q>") - end - elsif mask - mask = stream.read(4) or raise EOFError, "Could not read mask!" - end - - if length > maximum_frame_size - raise ProtocolError, "Invalid payload length: #{length} > #{maximum_frame_size}!" - end - - payload = stream.read(length) or raise EOFError, "Could not read payload!" - - if payload.bytesize != length - raise EOFError, "Incorrect payload length: #{length} != #{payload.bytesize}!" - end - - return self.new(finished, payload, flags: flags, opcode: opcode, mask: mask) - end - # Write this frame to the given stream. # @parameter stream [IO] The stream to write the serialized frame to. # @raises [ProtocolError] If the frame has invalid length or mask. diff --git a/lib/protocol/websocket/framer.rb b/lib/protocol/websocket/framer.rb index ab760e2..3de357b 100644 --- a/lib/protocol/websocket/framer.rb +++ b/lib/protocol/websocket/framer.rb @@ -57,6 +57,7 @@ def read_frame(maximum_frame_size = MAXIMUM_ALLOWED_FRAME_SIZE) end first_byte = buffer.getbyte(0) + second_byte = buffer.getbyte(1) finished = (first_byte & 0b1000_0000 != 0) flags = (first_byte & 0b0111_0000) >> 4 @@ -68,10 +69,51 @@ def read_frame(maximum_frame_size = MAXIMUM_ALLOWED_FRAME_SIZE) raise ProtocolError, "Control opcode = #{opcode} is reserved!" end - klass = @frames[opcode] || Frame - frame = klass.read(finished, flags, opcode, @stream, maximum_frame_size, buffer.getbyte(1)) + mask = (second_byte & 0b1000_0000 != 0) + length = second_byte & 0b0111_1111 + + if opcode & 0x8 != 0 + if length > 125 + raise ProtocolError, "Invalid control frame payload length: #{length} > 125!" + elsif !finished + raise ProtocolError, "Fragmented control frame!" + end + end - return frame + if length == 126 + if mask + buffer = @stream.read(6) or raise EOFError, "Could not read length and mask!" + length = buffer.unpack1("n") + mask = buffer.byteslice(2, 4) + else + buffer = @stream.read(2) or raise EOFError, "Could not read length!" + length = buffer.unpack1("n") + end + elsif length == 127 + if mask + buffer = @stream.read(12) or raise EOFError, "Could not read length and mask!" + length = buffer.unpack1("Q>") + mask = buffer.byteslice(8, 4) + else + buffer = @stream.read(8) or raise EOFError, "Could not read length!" + length = buffer.unpack1("Q>") + end + elsif mask + mask = @stream.read(4) or raise EOFError, "Could not read mask!" + end + + if length > maximum_frame_size + raise ProtocolError, "Invalid payload length: #{length} > #{maximum_frame_size}!" + end + + payload = @stream.read(length) or raise EOFError, "Could not read payload!" + + if payload.bytesize != length + raise EOFError, "Incorrect payload length: #{length} != #{payload.bytesize}!" + end + + klass = @frames[opcode] || Frame + return klass.new(finished, payload, flags: flags, opcode: opcode, mask: mask) end # Write a frame to the underlying stream. diff --git a/test/protocol/websocket/frame.rb b/test/protocol/websocket/frame.rb index f8fc352..4f574c9 100644 --- a/test/protocol/websocket/frame.rb +++ b/test/protocol/websocket/frame.rb @@ -29,49 +29,6 @@ end end - with ".read" do - it "rejects invalid control frame payload length" do - stream = StringIO.new("\xFF") - - expect do - subject.read(true, 0, 0x8, stream, 128) - end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Invalid control frame payload length/) - end - - it "rejects fragmented control frames" do - stream = StringIO.new("\x0F") - - expect do - subject.read(false, 0, 0x8, stream, 128) - end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Fragmented control frame/) - end - - it "rejects frames bigger than the maximum frame size" do - stream = StringIO.new("\x7D") - - expect do - subject.read(false, 0, 0, stream, 124) - end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Invalid payload length: \d+ > \d*!/) - end - - it "rejects frames with truncated payload" do - stream = StringIO.new("\x051234") - - expect do - subject.read(false, 0, 0, stream, 128) - end.to raise_exception(EOFError, message: be =~ /Incorrect payload length: \d+ != \d+!/) - end - - it "accepts a pre-read second byte" do - stream = StringIO.new("Hello") - second_byte = 0x05 - - frame = subject.read(true, 0, 0x1, stream, 128, second_byte) - expect(frame.payload).to be == "Hello" - expect(frame.mask).to be == false - end - end - with ".write" do let(:stream) {StringIO.new} diff --git a/test/protocol/websocket/framer.rb b/test/protocol/websocket/framer.rb index e4a9549..22bee4d 100644 --- a/test/protocol/websocket/framer.rb +++ b/test/protocol/websocket/framer.rb @@ -33,5 +33,56 @@ framer.read_frame end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Control opcode.*reserved/) end + + it "rejects invalid control frame payload length" do + # FIN=1, opcode=0x8 (close), MASK=1, length=127 → violates max 125 for control frames + stream.string = "\x88\xFF" + stream.rewind + + expect do + framer.read_frame + end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Invalid control frame payload length/) + end + + it "rejects fragmented control frames" do + # FIN=0, opcode=0x8 (close), MASK=0, length=15 + stream.string = "\x08\x0F" + stream.rewind + + expect do + framer.read_frame + end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Fragmented control frame/) + end + + it "rejects frames bigger than the maximum frame size" do + # FIN=1, opcode=0x2 (binary), MASK=0, length=125 + stream.string = "\x82\x7D" + stream.rewind + + expect do + framer.read_frame(124) + end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Invalid payload length: \d+ > \d+!/) + end + + it "rejects frames with truncated payload" do + # FIN=1, opcode=0x2 (binary), MASK=0, length=5, only 4 bytes of payload + stream.string = "\x82\x051234" + stream.rewind + + expect do + framer.read_frame + end.to raise_exception(EOFError, message: be =~ /Incorrect payload length: \d+ != \d+!/) + end + + it "reads a text frame" do + # FIN=1, opcode=0x1 (text), MASK=0, length=5, payload="Hello" + stream.string = "\x81\x05Hello" + stream.rewind + + frame = framer.read_frame + expect(frame).to be_a(Protocol::WebSocket::TextFrame) + expect(frame.payload).to be == "Hello" + expect(frame.mask).to be == false + end end end