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..70b0f8d 100644 --- a/lib/protocol/websocket/frame.rb +++ b/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. diff --git a/lib/protocol/websocket/framer.rb b/lib/protocol/websocket/framer.rb index 9ddfdb7..3de357b 100644 --- a/lib/protocol/websocket/framer.rb +++ b/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 diff --git a/test/protocol/websocket/connection.rb b/test/protocol/websocket/connection.rb index e3830c0..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 @@ -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} diff --git a/test/protocol/websocket/frame.rb b/test/protocol/websocket/frame.rb index 44878f2..4f574c9 100644 --- a/test/protocol/websocket/frame.rb +++ 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} diff --git a/test/protocol/websocket/framer.rb b/test/protocol/websocket/framer.rb index cf025f6..22bee4d 100644 --- a/test/protocol/websocket/framer.rb +++ b/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