em-websocket 0.4.0 → 0.5.3

data/lib/em-websocket/handshake.rb

@@ -1,4 +1,5 @@
 require "http/parser"
+require "uri"
 
 module EventMachine
   module WebSocket
@@ -24,7 +25,7 @@ module EventMachine
       def receive_data(data)
         @parser << data
 
-        if @headers
+        if defined? @headers
           process(@headers, @parser.upgrade_data)
         end
       rescue HTTP::Parser::Error => e
@@ -48,12 +49,12 @@ module EventMachine
       # Returns the request path (excluding any query params)
       #
       def path
-        @parser.request_path
+        @path
       end
 
       # Returns the query params as a string foo=bar&baz=...
       def query_string
-        @parser.query_string
+        @query_string
       end
 
       def query
@@ -66,6 +67,10 @@ module EventMachine
         @headers["origin"] || @headers["sec-websocket-origin"] || nil
       end
 
+      def secure?
+        @secure
+      end
+
       private
 
       def process(headers, remains)
@@ -73,13 +78,27 @@ module EventMachine
           raise HandshakeError, "Must be GET request"
         end
 
+        # Validate request path
+        #
+        # According to http://tools.ietf.org/search/rfc2616#section-5.1.2, an
+        # invalid Request-URI should result in a 400 status code, but
+        # HandshakeError's currently result in a WebSocket abort. It's not
+        # clear which should take precedence, but an abort will do just fine.
+        begin
+          uri = URI.parse(@parser.request_url)
+          @path = uri.path
+          @query_string = uri.query || ""
+        rescue URI::InvalidURIError
+          raise HandshakeError, "Invalid request URI: #{@parser.request_url}"
+        end
+
         # Validate Upgrade
         unless @parser.upgrade?
           raise HandshakeError, "Not an upgrade request"
         end
         upgrade = @headers['upgrade']
         unless upgrade.kind_of?(String) && upgrade.downcase == 'websocket'
-          raise HandshakeError, "Invalid upgrade header: #{upgrade}"
+          raise HandshakeError, "Invalid upgrade header: #{upgrade.inspect}"
         end
 
         # Determine version heuristically

data/lib/em-websocket/handshake04.rb

@@ -10,11 +10,6 @@ module EventMachine
           raise HandshakeError, "sec-websocket-key header is required"
         end
 
-        # Optional
-        origin = headers['sec-websocket-origin']
-        protocols = headers['sec-websocket-protocol']
-        extensions = headers['sec-websocket-extensions']
-
         string_to_sign = "#{key}258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
         signature = Base64.encode64(Digest::SHA1.digest(string_to_sign)).chomp
 
@@ -22,12 +17,21 @@ module EventMachine
         upgrade << "Upgrade: websocket"
         upgrade << "Connection: Upgrade"
         upgrade << "Sec-WebSocket-Accept: #{signature}"
+        if protocol = headers['sec-websocket-protocol']
+          validate_protocol!(protocol)
+          upgrade << "Sec-WebSocket-Protocol: #{protocol}"
+        end
 
-        # TODO: Support sec-websocket-protocol
+        # TODO: Support sec-websocket-protocol selection
         # TODO: sec-websocket-extensions
 
         return upgrade.join("\r\n") + "\r\n\r\n"
       end
+
+      def self.validate_protocol!(protocol)
+        raise HandshakeError, "Invalid WebSocket-Protocol: empty" if protocol.empty?
+        # TODO: Validate characters
+      end
     end
   end
 end

data/lib/em-websocket/handshake75.rb

@@ -9,10 +9,20 @@ module EventMachine
         upgrade << "Upgrade: WebSocket\r\n"
         upgrade << "Connection: Upgrade\r\n"
         upgrade << "WebSocket-Origin: #{headers['origin']}\r\n"
-        upgrade << "WebSocket-Location: #{location}\r\n\r\n"
+        upgrade << "WebSocket-Location: #{location}\r\n"
+        if protocol = headers['sec-websocket-protocol']
+          validate_protocol!(protocol)
+          upgrade << "Sec-WebSocket-Protocol: #{protocol}\r\n"
+        end
+        upgrade << "\r\n"
 
         return upgrade
       end
+
+      def self.validate_protocol!(protocol)
+        raise HandshakeError, "Invalid WebSocket-Protocol: empty" if protocol.empty?
+        # TODO: Validate characters
+      end
     end
   end
 end

data/lib/em-websocket/handshake76.rb

@@ -39,6 +39,10 @@ module EventMachine::WebSocket
       end
 
       def numbers_over_spaces(string)
+        unless string
+          raise HandshakeError, "WebSocket key1 or key2 is missing"
+        end
+
         numbers = string.scan(/[0-9]/).join.to_i
 
         spaces = string.scan(/ /).size

data/lib/em-websocket/masking04.rb

@@ -15,12 +15,8 @@ module EventMachine
         @masking_key = nil
       end
 
-      def slice_mask
-        slice!(0, 4)
-      end
-
       def getbyte(index)
-        if @masking_key
+        if defined?(@masking_key) && @masking_key
           masked_char = super
           masked_char ? masked_char ^ @masking_key.getbyte(index % 4) : nil
         else

data/lib/em-websocket/message_processor_03.rb

@@ -6,17 +6,20 @@ module EventMachine
       def message(message_type, extension_data, application_data)
         case message_type
         when :close
+          @close_info = {
+            :code => 1005,
+            :reason => "",
+            :was_clean => true,
+          }
           if @state == :closing
             # TODO: Check that message body matches sent data
             # We can close connection immediately since there is no more data
             # is allowed to be sent or received on this connection
             @connection.close_connection
-            @state = :closed
           else
             # Acknowlege close
             # The connection is considered closed
             send_frame(:close, application_data)
-            @state = :closed
             @connection.close_connection_after_writing
           end
         when :ping
@@ -31,7 +34,7 @@ module EventMachine
           end
           @connection.trigger_on_message(application_data)
         when :binary
-          @connection.trigger_on_message(application_data)
+          @connection.trigger_on_binary(application_data)
         end
       end
 

data/lib/em-websocket/message_processor_06.rb

@@ -19,32 +19,53 @@ module EventMachine
           
           debug [:close_frame_received, status_code, application_data]
           
+          @close_info = {
+            :code => status_code || 1005,
+            :reason => application_data,
+            :was_clean => true,
+          }
+
           if @state == :closing
             # We can close connection immediately since no more data may be
             # sent or received on this connection
             @connection.close_connection
-            @state = :closed
-          else
-            # Acknowlege close
+          elsif @state == :connected
+            # Acknowlege close & echo status back to client
             # The connection is considered closed
-            send_frame(:close, '')
-            @state = :closed
+            close_data = [status_code || 1000].pack('n')
+            send_frame(:close, close_data)
             @connection.close_connection_after_writing
-            # TODO: Send close status code and body to app code
           end
         when :ping
-          # Pong back the same data
-          send_frame(:pong, application_data)
+          # There are a couple of protections here against malicious/broken WebSocket abusing ping frames.
+          #
+          # 1. Delay 200ms before replying. This reduces the number of pings from WebSocket clients behaving as
+          #    `for (;;) { send_ping(conn); rcv_pong(conn); }`. The spec says we "SHOULD respond with Pong frame as soon
+          #    as is practical".
+          # 2. Reply at most every 200ms. This reduces the number of pong frames sent to WebSocket clients behaving as
+          #    `for (;;) { send_ping(conn); }`. The spec says "If an endpoint receives a Ping frame and has not yet sent
+          #    Pong frame(s) in response to previous Ping frame(s), the endpoint MAY elect to send a Pong frame for only
+          #    the most recently processed Ping frame."
+          @most_recent_pong_application_data = application_data
+          if @pong_timer == nil then
+            @pong_timer = EventMachine.add_timer(0.2) do
+              @pong_timer = nil
+              send_frame(:pong, @most_recent_pong_application_data)
+            end
+          end
           @connection.trigger_on_ping(application_data)
         when :pong
           @connection.trigger_on_pong(application_data)
         when :text
           if application_data.respond_to?(:force_encoding)
             application_data.force_encoding("UTF-8")
+            unless application_data.valid_encoding?
+              raise InvalidDataError, "Invalid UTF8 data"
+            end
           end
           @connection.trigger_on_message(application_data)
         when :binary
-          @connection.trigger_on_message(application_data)
+          @connection.trigger_on_binary(application_data)
         end
       end
 

data/lib/em-websocket/version.rb

@@ -1,5 +1,5 @@
 module EventMachine
   module Websocket
-    VERSION = "0.4.0"
+    VERSION = "0.5.3"
   end
 end

data/lib/em-websocket/websocket.rb

@@ -2,8 +2,11 @@ module EventMachine
   module WebSocket
     class << self
       attr_accessor :max_frame_size
+      attr_accessor :close_timeout
     end
     @max_frame_size = 10 * 1024 * 1024 # 10MB
+    # Connections are given 60s to close after being sent a close handshake
+    @close_timeout = 60
 
     # All errors raised by em-websocket should descend from this class
     class WebSocketError < RuntimeError; end
@@ -17,6 +20,10 @@ module EventMachine
       def code; 1002; end
     end
 
+    class InvalidDataError < WSProtocolError
+      def code; 1007; end
+    end
+
     # 1009: Message too big to process
     class WSMessageTooBigError < WSProtocolError
       def code; 1009; end

data/spec/helper.rb

@@ -1,10 +1,15 @@
+# encoding: BINARY
+
 require 'rubygems'
 require 'rspec'
 require 'em-spec/rspec'
-require 'pp'
 require 'em-http'
 
 require 'em-websocket'
+require 'em-websocket-client'
+
+require 'integration/shared_examples'
+require 'integration/gte_03_examples'
 
 RSpec.configure do |c|
   c.mock_with :rspec
@@ -27,75 +32,80 @@ class FakeWebSocketClient < EM::Connection
     # puts "RECEIVE DATA #{data}"
     if @state == :new
       @handshake_response = data
-      @onopen.call if @onopen
+      @onopen.call if defined? @onopen
       @state = :open
     else
-      @onmessage.call(data) if @onmessage
+      @onmessage.call(data) if defined? @onmessage
       @packets << data
     end
   end
 
-  def send(data)
-    send_data("\x00#{data}\xff")
+  def send(application_data)
+    send_frame(:text, application_data)
+  end
+
+  def send_frame(type, application_data)
+    send_data construct_frame(type, application_data)
   end
 
   def unbind
-    @onclose.call if @onclose
+    @onclose.call if defined? @onclose
+  end
+
+  private
+
+  def construct_frame(type, data)
+    "\x00#{data}\xff"
   end
 end
 
 class Draft03FakeWebSocketClient < FakeWebSocketClient
-  def send(application_data)
-    frame = ''
-    opcode = 4 # fake only supports text frames
-    byte1 = opcode # since more, rsv1-3 are 0
-    frame << byte1
+  private
 
-    length = application_data.size
+  def construct_frame(type, data)
+    frame = ""
+    frame << EM::WebSocket::Framing03::FRAME_TYPES[type]
+    frame << encoded_length(data.size)
+    frame << data
+  end
+
+  def encoded_length(length)
     if length <= 125
-      byte2 = length # since rsv4 is 0
-      frame << byte2
+      [length].pack('C') # since rsv4 is 0
     elsif length < 65536 # write 2 byte length
-      frame << 126
-      frame << [length].pack('n')
+      "\126#{[length].pack('n')}"
     else # write 8 byte length
-      frame << 127
-      frame << [length >> 32, length & 0xFFFFFFFF].pack("NN")
+      "\127#{[length >> 32, length & 0xFFFFFFFF].pack("NN")}"
     end
-
-    frame << application_data
-
-    send_data(frame)
   end
 end
 
-class Draft07FakeWebSocketClient < FakeWebSocketClient
-  def send(application_data)
-    frame = ''
-    opcode = 1 # fake only supports text frames
-    byte1 = opcode | 0b10000000 # since more, rsv1-3 are 0
-    frame << byte1
+class Draft05FakeWebSocketClient < Draft03FakeWebSocketClient
+  private
 
-    length = application_data.size
-    if length <= 125
-      byte2 = length # since rsv4 is 0
-      frame << byte2
-    elsif length < 65536 # write 2 byte length
-      frame << 126
-      frame << [length].pack('n')
-    else # write 8 byte length
-      frame << 127
-      frame << [length >> 32, length & 0xFFFFFFFF].pack("NN")
-    end
+  def construct_frame(type, data)
+    frame = ""
+    frame << "\x00\x00\x00\x00" # Mask with nothing for simplicity
+    frame << (EM::WebSocket::Framing05::FRAME_TYPES[type] | 0b10000000)
+    frame << encoded_length(data.size)
+    frame << data
+  end
+end
 
-    frame << application_data
+class Draft07FakeWebSocketClient < Draft05FakeWebSocketClient
+  private
 
-    send_data(frame)
+  def construct_frame(type, data)
+    frame = ""
+    frame << (EM::WebSocket::Framing07::FRAME_TYPES[type] | 0b10000000)
+    # Should probably mask the data, but I get away without bothering since
+    # the server doesn't enforce that incoming frames are masked
+    frame << encoded_length(data.size)
+    frame << data
   end
 end
 
-
-# Wrap EM:HttpRequest in a websocket like interface so that it can be used in the specs with the same interface as FakeWebSocketClient
+# Wrapper around em-websocket-client
 class Draft75WebSocketClient
   def onopen(&blk);     @onopen = blk;    end
   def onclose(&blk);    @onclose = blk;   end
@@ -103,18 +113,17 @@ class Draft75WebSocketClient
   def onmessage(&blk);  @onmessage = blk; end
 
   def initialize
-    @ws = EventMachine::HttpRequest.new('ws://127.0.0.1:12345/').get({
-      :timeout => 0,
-      :origin => 'http://example.com',
-    })
-    @ws.errback { @onerror.call if @onerror }
-    @ws.callback { @onopen.call if @onopen }
-    @ws.stream { |msg| @onmessage.call(msg) if @onmessage }
-    @ws.disconnect { @onclose.call if @onclose }
+    @ws = EventMachine::WebSocketClient.connect('ws://127.0.0.1:12345/',
+                                                :version => 75,
+                                                :origin => 'http://example.com')
+    @ws.errback { |err| @onerror.call if defined? @onerror }
+    @ws.callback { @onopen.call if defined? @onopen }
+    @ws.stream { |msg| @onmessage.call(msg) if defined? @onmessage }
+    @ws.disconnect { @onclose.call if defined? @onclose }
   end
 
   def send(message)
-    @ws.send(message)
+    @ws.send_msg(message)
   end
 
   def close_connection
@@ -122,6 +131,12 @@ class Draft75WebSocketClient
   end
 end
 
+def start_server(opts = {})
+  EM::WebSocket.run({:host => "0.0.0.0", :port => 12345}.merge(opts)) { |ws|
+    yield ws if block_given?
+  }
+end
+
 def format_request(r)
   data = "#{r[:method]} #{r[:path]} HTTP/1.1\r\n"
   header_lines = r[:headers].map { |k,v| "#{k}: #{v}" }

data/spec/integration/common_spec.rb

@@ -14,32 +14,29 @@ describe "WebSocket server" do
         http.callback { fail }
       end
 
-      EM::WebSocket.start(:host => "0.0.0.0", :port => 12345) {}
+      start_server
     }
   end
 
   it "should expose the WebSocket request headers, path and query params" do
     em {
       EM.add_timer(0.1) do
-        http = EM::HttpRequest.new('ws://127.0.0.1:12345/').get :timeout => 0
-        http.errback { fail }
-        http.callback {
-          http.response_header.status.should == 101
-          http.close_connection
-        }
-        http.stream { |msg| }
+        ws = EventMachine::WebSocketClient.connect('ws://127.0.0.1:12345/',
+                                                   :origin => 'http://example.com')
+        ws.errback { fail }
+        ws.callback { ws.close_connection }
+        ws.stream { |msg| }
       end
 
-      EM::WebSocket.start(:host => "0.0.0.0", :port => 12345) do |ws|
+      start_server do |ws|
         ws.onopen { |handshake|
           headers = handshake.headers
-          headers["User-Agent"].should == "EventMachine HttpClient"
           headers["Connection"].should == "Upgrade"
-          headers["Upgrade"].should == "WebSocket"
+          headers["Upgrade"].should == "websocket"
           headers["Host"].to_s.should == "127.0.0.1:12345"
           handshake.path.should == "/"
           handshake.query.should == {}
-          handshake.origin.should == "127.0.0.1"
+          handshake.origin.should == 'http://example.com'
         }
         ws.onclose {
           ws.state.should == :closed
@@ -52,19 +49,15 @@ describe "WebSocket server" do
   it "should expose the WebSocket path and query params when nonempty" do
     em {
       EM.add_timer(0.1) do
-        http = EM::HttpRequest.new('ws://127.0.0.1:12345/hello').get({
-          :query => {'foo' => 'bar', 'baz' => 'qux'},
-          :timeout => 0
-        })
-        http.errback { fail }
-        http.callback {
-          http.response_header.status.should == 101
-          http.close_connection
+        ws = EventMachine::WebSocketClient.connect('ws://127.0.0.1:12345/hello?foo=bar&baz=qux')
+        ws.errback { fail }
+        ws.callback {
+          ws.close_connection
         }
-        http.stream { |msg| }
+        ws.stream { |msg| }
       end
 
-      EM::WebSocket.start(:host => "0.0.0.0", :port => 12345) do |ws|
+      start_server do |ws|
         ws.onopen { |handshake|
           handshake.path.should == '/hello'
           handshake.query_string.split('&').sort.
@@ -82,7 +75,7 @@ describe "WebSocket server" do
   it "should raise an exception if frame sent before handshake complete" do
     em {
       # 1. Start WebSocket server
-      EM::WebSocket.start(:host => "0.0.0.0", :port => 12345) { |ws|
+      start_server { |ws|
         # 3. Try to send a message to the socket
         lambda {
           ws.send('early message')
@@ -98,18 +91,15 @@ describe "WebSocket server" do
   it "should allow the server to be started inside an existing EM" do
     em {
       EM.add_timer(0.1) do
-        http = EM::HttpRequest.new('ws://127.0.0.1:12345/').get :timeout => 0
-        http.errback { fail }
-        http.callback {
-          http.response_header.status.should == 101
-          http.close_connection
-        }
-        http.stream { |msg| }
+        http = EM::HttpRequest.new('http://127.0.0.1:12345/').get :timeout => 0
+        http.errback { |e| done }
+        http.callback { fail }
       end
 
-      EM::WebSocket.run(:host => "0.0.0.0", :port => 12345) do |ws|
+      start_server do |ws|
         ws.onopen { |handshake|
-          handshake.headers["User-Agent"].should == "EventMachine HttpClient"
+          headers = handshake.headers
+          headers["Host"].to_s.should == "127.0.0.1:12345"
         }
         ws.onclose {
           ws.state.should == :closed
@@ -118,4 +108,31 @@ describe "WebSocket server" do
       end
     }
   end
+
+  context "outbound limit set" do
+    it "should close the connection if the limit is reached" do
+      em {
+        start_server(:outbound_limit => 150) do |ws|
+          # Increase the message size by one on each loop
+          ws.onmessage{|msg| ws.send(msg + "x") }
+          ws.onclose{|status|
+            status[:code].should == 1006 # Unclean
+            status[:was_clean].should be false
+          }
+        end
+
+        EM.add_timer(0.1) do
+          ws = EventMachine::WebSocketClient.connect('ws://127.0.0.1:12345/')
+          ws.callback { ws.send_msg "hello" }
+          ws.disconnect { done } # Server closed the connection
+          ws.stream { |msg|
+            # minus frame size ? (getting 146 max here)
+            msg.data.size.should <= 150
+            # Return back the message
+            ws.send_msg(msg.data)
+          }
+        end
+      }
+    end
+  end
 end