Move all read logic to Framer.

Author: samuel-williams-shopify

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.

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.

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}
 

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