Optimize frame reading.

Co-authored-by: Claude <noreply@anthropic.com>

Author: 2 people

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.

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

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

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}

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

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