Optimize frame reading throughput

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

@@ -168,75 +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.unpack("C").first
-
-				finished = (byte & 0b1000_0000 != 0)
-				flags = (byte & 0b0111_0000) >> 4
-				opcode = byte & 0b0000_1111
-
-				if (0x3 .. 0x7).include?(opcode)
-					raise ProtocolError, "Non-control opcode = #{opcode} is reserved!"
-				elsif (0xB .. 0xF).include?(opcode)
-					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.
-			# @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)
-				buffer = stream.read(1) or raise EOFError, "Could not read header!"
-				byte = buffer.unpack("C").first
-
-				mask = (byte & 0b1000_0000 != 0)
-				length = 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
-					buffer = stream.read(2) or raise EOFError, "Could not read length!"
-					length = buffer.unpack("n").first
-				elsif length == 127
-					buffer = stream.read(8) or raise EOFError, "Could not read length!"
-					length = buffer.unpack("Q>").first
-				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
-					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.

lib/protocol/websocket/framer.rb

@@ -50,29 +50,76 @@ 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)
 
-				# Read the frame:
-				klass = @frames[opcode] || Frame
-				frame = klass.read(finished, flags, opcode, @stream, maximum_frame_size)
+				unless buffer and buffer.bytesize == 2
+					raise EOFError, "Could not read frame header!"
+				end
+
+				first_byte = buffer.getbyte(0)
+				second_byte = buffer.getbyte(1)
+
+				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
+
+				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
 
-				return frame
+				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.
 			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

@@ -5,6 +5,7 @@
 # Copyright, 2019, by Soumya.
 
 require "socket"
+require "securerandom"
 require "protocol/websocket/connection"
 
 describe Protocol::WebSocket::Connection do
@@ -301,6 +302,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

@@ -29,54 +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")
-
-			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
-	end
-
 	with ".write" do
 		let(:stream) {StringIO.new}
 

test/protocol/websocket/framer.rb

@@ -15,5 +15,74 @@
 				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
+
+		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