Move Frame#write -> Framer#write_frame and elide Frame@length.

Author: samuel-williams-shopify

lib/protocol/websocket/frame.rb

@@ -21,7 +21,6 @@ class Frame
 
 			OPCODE = 0
 
-			# @parameter length [Integer] The length of the payload, or nil if the header has not been read yet.
 			# @parameter mask [Boolean | String] An optional 4-byte string which is used to mask the payload.
 			def initialize(finished = true, payload = nil, flags: 0, opcode: self.class::OPCODE, mask: false)
 				if mask == true
@@ -32,7 +31,6 @@ def initialize(finished = true, payload = nil, flags: 0, opcode: self.class::OPC
 				@flags = flags
 				@opcode = opcode
 				@mask = mask
-				@length = payload&.bytesize
 				@payload = payload
 			end
 
@@ -51,9 +49,9 @@ def <=> other
 			end
 
 			# Convert this frame to an array of its fields for comparison or inspection.
-			# @returns [Array] An array of `[finished, flags, opcode, mask, length, payload]`.
+			# @returns [Array] An array of `[finished, flags, opcode, mask, payload]`.
 			def to_ary
-				[@finished, @flags, @opcode, @mask, @length, @payload]
+				[@finished, @flags, @opcode, @mask, @payload]
 			end
 
 			# Check whether this is a control frame (opcode has bit 3 set).
@@ -104,9 +102,14 @@ def continued?
 			attr_accessor :flags
 			attr_accessor :opcode
 			attr_accessor :mask
-			attr_accessor :length
 			attr_accessor :payload
 
+			# The byte length of the payload.
+			# @returns [Integer | nil]
+			def length
+				@payload&.bytesize
+			end
+			
 			if IO.const_defined?(:Buffer) && IO::Buffer.respond_to?(:for) && IO::Buffer.method_defined?(:xor!)
 				private def mask_xor(data, mask)
 					buffer = data.dup
@@ -135,18 +138,14 @@ def continued?
 			# @parameter data [String] The payload data to pack.
 			# @returns [Frame] Returns `self`.
 			def pack(data = "")
-				length = data.bytesize
-				
-				if length.bit_length > 63
-					raise ProtocolError, "Frame length #{@length} bigger than allowed maximum!"
+				if data.bytesize.bit_length > 63
+					raise ProtocolError, "Frame length #{data.bytesize} bigger than allowed maximum!"
 				end
 
 				if @mask
-					@payload = mask_xor(data, mask)
-					@length = length
+					@payload = mask_xor(data, @mask)
 				else
 					@payload = data
-					@length = length
 				end
 
 				return self
@@ -167,45 +166,6 @@ def unpack
 			def apply(connection)
 				connection.receive_frame(self)
 			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.
-			def write(stream)
-				buffer = String.new(encoding: Encoding::BINARY)
-				
-				if @payload&.bytesize != @length
-					raise ProtocolError, "Invalid payload length: #{@length} != #{@payload.bytesize} for #{self}!"
-				end
-				
-				if @mask and @mask.bytesize != 4
-					raise ProtocolError, "Invalid mask length!"
-				end
-				
-				if length <= 125
-					short_length = length
-				elsif length.bit_length <= 16
-					short_length = 126
-				else
-					short_length = 127
-				end
-				
-				buffer << [
-					(@finished ? 0b1000_0000 : 0) | (@flags << 4) | @opcode,
-					(@mask ? 0b1000_0000 : 0) | short_length,
-				].pack("CC")
-				
-				if short_length == 126
-					buffer << [@length].pack("n")
-				elsif short_length == 127
-					buffer << [@length].pack("Q>")
-				end
-				
-				buffer << @mask if @mask
-				
-				stream.write(buffer)
-				stream.write(@payload)
-			end
 		end
 	end
 end

lib/protocol/websocket/framer.rb

@@ -117,8 +117,38 @@ def read_frame(maximum_frame_size = MAXIMUM_ALLOWED_FRAME_SIZE)
 			end
 
 			# Write a frame to the underlying stream.
+			# @parameter frame [Frame] The frame to serialize and write.
+			# @raises [ProtocolError] If the frame has an invalid mask.
 			def write_frame(frame)
-				frame.write(@stream)
+				if frame.mask and frame.mask.bytesize != 4
+					raise ProtocolError, "Invalid mask length!"
+				end
+				
+				length = frame.length
+				
+				if length <= 125
+					short_length = length
+				elsif length.bit_length <= 16
+					short_length = 126
+				else
+					short_length = 127
+				end
+				
+				buffer = [
+					(frame.finished ? 0b1000_0000 : 0) | (frame.flags << 4) | frame.opcode,
+					(frame.mask ? 0b1000_0000 : 0) | short_length,
+				].pack("CC")
+				
+				if short_length == 126
+					buffer << [length].pack("n")
+				elsif short_length == 127
+					buffer << [length].pack("Q>")
+				end
+				
+				buffer << frame.mask if frame.mask
+				
+				@stream.write(buffer)
+				@stream.write(frame.payload)
 			end
 		end
 	end

test/protocol/websocket/binary_frame.rb

@@ -5,6 +5,7 @@
 
 require "protocol/websocket/a_websocket_frame"
 require "protocol/websocket/binary_frame"
+require "protocol/websocket/framer"
 
 describe Protocol::WebSocket::BinaryFrame do
 	let(:frame) {subject.new}
@@ -28,9 +29,7 @@
 
 		it "encodes binary representation" do
 			buffer = StringIO.new
-			
-			frame.write(buffer)
-			
+			Protocol::WebSocket::Framer.new(buffer).write_frame(frame)
 			expect(buffer.string).to be == "\x82\x8Babcd)\a\x0F\b\x0EB4\v\x13\x0E\a"
 		end
 	end
@@ -42,9 +41,7 @@
 
 		it "encodes binary representation" do
 			buffer = StringIO.new
-			
-			frame.write(buffer)
-			
+			Protocol::WebSocket::Framer.new(buffer).write_frame(frame)
 			expect(buffer.string).to be == "\x82\vHello World"
 		end
 	end

test/protocol/websocket/close_frame.rb

@@ -43,7 +43,6 @@
 
 		it "rejects a close frame with an invalid format" do
 			frame.payload = "1"
-			frame.length = 1
 
 			expect do
 				frame.unpack

test/protocol/websocket/frame.rb

@@ -29,26 +29,4 @@
 		end
 	end
 
-	with ".write" do
-		let(:stream) {StringIO.new}
-		
-		it "fails with invalid payload length" do
-			frame.length = 5
-			frame.payload = "1234"
-			
-			expect do
-				frame.write(stream)
-			end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Invalid payload length/)
-		end
-		
-		it "fails with invalid mask size" do
-			frame.length = 5
-			frame.payload = "12345"
-			frame.mask = "bad"
-			
-			expect do
-				frame.write(stream)
-			end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Invalid mask length/)
-		end
-	end
 end

test/protocol/websocket/framer.rb

@@ -9,6 +9,30 @@
 	let(:stream) {StringIO.new}
 	let(:framer) {subject.new(stream)}
 
+	with "#write_frame" do
+		it "fails with invalid mask size" do
+			frame = Protocol::WebSocket::Frame.new(true, "12345")
+			frame.mask = "bad"
+			
+			expect do
+				framer.write_frame(frame)
+			end.to raise_exception(Protocol::WebSocket::ProtocolError, message: be =~ /Invalid mask length/)
+		end
+		
+		it "writes a text frame and reads it back" do
+			output = StringIO.new
+			writer = Protocol::WebSocket::Framer.new(output)
+			
+			frame = Protocol::WebSocket::TextFrame.new(true, "Hello")
+			writer.write_frame(frame)
+			
+			reader = Protocol::WebSocket::Framer.new(StringIO.new(output.string))
+			received = reader.read_frame
+			expect(received).to be_a(Protocol::WebSocket::TextFrame)
+			expect(received.payload).to be == "Hello"
+		end
+	end
+	
 	with "#read_frame" do
 		it "fails if it can't read the frame header" do
 			expect do

test/protocol/websocket/text_frame.rb

@@ -5,6 +5,7 @@
 
 require "protocol/websocket/a_websocket_frame"
 require "protocol/websocket/text_frame"
+require "protocol/websocket/framer"
 
 describe Protocol::WebSocket::TextFrame do
 	let(:frame) {subject.new}
@@ -24,9 +25,7 @@
 
 		it "encodes binary representation" do
 			buffer = StringIO.new
-			
-			frame.write(buffer)
-			
+			Protocol::WebSocket::Framer.new(buffer).write_frame(frame)
 			expect(buffer.string).to be == "\x81\x8Babcd)\a\x0F\b\x0EB4\v\x13\x0E\a"
 		end
 	end
@@ -38,9 +37,7 @@
 
 		it "encodes binary representation" do
 			buffer = StringIO.new
-			
-			frame.write(buffer)
-			
+			Protocol::WebSocket::Framer.new(buffer).write_frame(frame)
 			expect(buffer.string).to be == "\x81\vHello World"
 		end
 	end