em-websocket 0.1.4 → 0.2.0

data/lib/em-websocket/framing76.rb

@@ -0,0 +1,96 @@
+module EventMachine
+  module WebSocket
+    module Framing76
+      
+      # Set the max frame lenth to very high value (10MB) until there is a
+      # limit specified in the spec to protect against malicious attacks
+      MAXIMUM_FRAME_LENGTH = 10 * 1024 * 1024
+      
+      def initialize_framing
+        @data = ''
+      end
+      
+      def process_data
+        debug [:message, @data]
+
+        # This algorithm comes straight from the spec
+        # http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76#section-4.2
+
+        error = false
+
+        while !error
+          pointer = 0
+          frame_type = @data[pointer].to_i
+          pointer += 1
+
+          if (frame_type & 0x80) == 0x80
+            # If the high-order bit of the /frame type/ byte is set
+            length = 0
+
+            loop do
+              return false if !@data[pointer]
+              b = @data[pointer].to_i
+              pointer += 1
+              b_v = b & 0x7F
+              length = length * 128 + b_v
+              break unless (b & 0x80) == 0x80
+            end
+
+            # Addition to the spec to protect against malicious requests
+            if length > MAXIMUM_FRAME_LENGTH
+              @connection.close_with_error(DataError.new("Frame length too long (#{length} bytes)"))
+              return false
+            end
+
+            if @data[pointer+length-1] == nil
+              debug [:buffer_incomplete, @data.inspect]
+              # Incomplete data - leave @data to accumulate
+              error = true
+            else
+              # Straight from spec - I'm sure this isn't crazy...
+              # 6. Read /length/ bytes.
+              # 7. Discard the read bytes.
+              @data = @data[(pointer+length)..-1]
+
+              # If the /frame type/ is 0xFF and the /length/ was 0, then close
+              if length == 0
+                @connection.send_data("\xff\x00")
+                @state = :closing
+                @connection.close_connection_after_writing
+              else
+                error = true
+              end
+            end
+          else
+            # If the high-order bit of the /frame type/ byte is _not_ set
+            msg = @data.slice!(/\A\x00([^\xff]*)\xff/)
+            if msg
+              msg.gsub!(/\A\x00|\xff\z/, '')
+              if @state == :closing
+                debug [:ignored_message, msg]
+              else
+                msg.force_encoding('UTF-8') if msg.respond_to?(:force_encoding)
+                @connection.trigger_on_message(msg)
+              end
+            else
+              error = true
+            end
+          end
+        end
+
+        false
+      end
+      
+      # frames need to start with 0x00-0x7f byte and end with
+      # an 0xFF byte. Per spec, we can also set the first
+      # byte to a value betweent 0x80 and 0xFF, followed by
+      # a leading length indicator
+      def send_text_frame(data)
+        ary = ["\x00", data, "\xff"]
+        ary.collect{ |s| s.force_encoding('UTF-8') if s.respond_to?(:force_encoding) }
+        @connection.send_data(ary.join)
+      end
+
+    end
+  end
+end

data/lib/em-websocket/handler.rb

@@ -3,17 +3,41 @@ module EventMachine
     class Handler
       include Debugger
 
-      attr_reader :request
+      attr_reader :request, :state
 
-      def initialize(request, response, debug = false)
-        @request = request
-        @response = response
+      def initialize(connection, request, debug = false)
+        @connection, @request = connection, request
         @debug = debug
+        @state = :handshake
+        initialize_framing
       end
 
+      def run
+        @connection.send_data handshake
+        @state = :connected
+        @connection.trigger_on_open
+      end
+
+      # Handshake response
       def handshake
         # Implemented in subclass
       end
+
+      def receive_data(data)
+        @data << data
+        process_data
+      end
+
+      def close_websocket
+        # Unless redefined in a subclass, just close the connection
+        @state = :closed
+        @connection.close_connection_after_writing
+      end
+
+      def unbind
+        @state = :closed
+        @connection.trigger_on_close
+      end
     end
   end
 end

data/lib/em-websocket/handler03.rb

@@ -0,0 +1,14 @@
+module EventMachine
+  module WebSocket
+    class Handler03 < Handler
+      include Handshake76
+      include Framing03
+
+      def close_websocket
+        # TODO: Should we send data and check the response matches?
+        send_frame(:close, '')
+        @state = :closing
+      end
+    end
+  end
+end

data/lib/em-websocket/handler75.rb

@@ -1,21 +1,8 @@
 module EventMachine
   module WebSocket
     class Handler75 < Handler
-      def handshake
-        location  = "#{@request['Host'].scheme}://#{@request['Host'].host}"
-        location << ":#{@request['Host'].port}" if @request['Host'].port
-        location << @request['Path']
-
-        upgrade =  "HTTP/1.1 101 Web Socket Protocol Handshake\r\n"
-        upgrade << "Upgrade: WebSocket\r\n"
-        upgrade << "Connection: Upgrade\r\n"
-        upgrade << "WebSocket-Origin: #{@request['Origin']}\r\n"
-        upgrade << "WebSocket-Location: #{location}\r\n\r\n"
-
-        debug [:upgrade_headers, upgrade]
-
-        return upgrade
-      end
+      include Handshake75
+      include Framing76
     end
   end
 end

data/lib/em-websocket/handler76.rb

@@ -1,64 +1,11 @@
-require 'digest/md5'
-
 module EventMachine
   module WebSocket
     class Handler76 < Handler
+      include Handshake76
+      include Framing76
+
       # "\377\000" is octet version and "\xff\x00" is hex version
       TERMINATE_STRING = "\xff\x00"
-
-      def handshake
-        challenge_response = solve_challenge(
-          @request['Sec-WebSocket-Key1'],
-          @request['Sec-WebSocket-Key2'],
-          @request['Third-Key']
-        )
-
-        location  = "#{@request['Host'].scheme}://#{@request['Host'].host}"
-        location << ":#{@request['Host'].port}" if @request['Host'].port
-        location << @request['Path']
-
-        upgrade =  "HTTP/1.1 101 WebSocket Protocol Handshake\r\n"
-        upgrade << "Upgrade: WebSocket\r\n"
-        upgrade << "Connection: Upgrade\r\n"
-        upgrade << "Sec-WebSocket-Location: #{location}\r\n"
-        upgrade << "Sec-WebSocket-Origin: #{@request['Origin']}\r\n"
-        if protocol = @request['Sec-WebSocket-Protocol']
-          validate_protocol!(protocol)
-          upgrade << "Sec-WebSocket-Protocol: #{protocol}\r\n"
-        end
-        upgrade << "\r\n"
-        upgrade << challenge_response
-
-        debug [:upgrade_headers, upgrade]
-
-        return upgrade
-      end
-
-      private
-
-      def solve_challenge(first, second, third)
-        # Refer to 5.2 4-9 of the draft 76
-        sum = [(extract_nums(first) / count_spaces(first))].pack("N*") +
-          [(extract_nums(second) / count_spaces(second))].pack("N*") +
-          third
-        Digest::MD5.digest(sum)
-      end
-
-      def extract_nums(string)
-        string.scan(/[0-9]/).join.to_i
-      end
-
-      def count_spaces(string)
-        spaces = string.scan(/ /).size
-        # As per 5.2.5, abort the connection if spaces are zero.
-        raise HandshakeError, "Websocket Key1 or Key2 does not contain spaces - this is a symptom of a cross-protocol attack" if spaces == 0
-        return spaces
-      end
-
-      def validate_protocol!(protocol)
-        raise HandshakeError, "Invalid WebSocket-Protocol: empty" if protocol.empty?
-        # TODO: Validate characters
-      end
     end
   end
 end

data/lib/em-websocket/handler_factory.rb

@@ -4,11 +4,16 @@ module EventMachine
       PATH   = /^(\w+) (\/[^\s]*) HTTP\/1\.1$/
       HEADER = /^([^:]+):\s*(.+)$/
 
-      def self.build(data, secure = false, debug = false)
+      def self.build(connection, data, secure = false, debug = false)
+        (header, remains) = data.split("\r\n\r\n", 2)
+        unless remains
+          # The whole header has not been received yet.
+          return nil
+        end
+
         request = {}
-        response = nil
 
-        lines = data.split("\r\n")
+        lines = header.split("\r\n")
 
         # extract request path
         first_line = lines.shift.match(PATH)
@@ -27,7 +32,24 @@ module EventMachine
           h = HEADER.match(line)
           request[h[1].strip] = h[2].strip if h
         end
-        request['Third-Key'] = lines.last
+
+        version = request['Sec-WebSocket-Key1'] ? 76 : 75
+        case version
+        when 75
+          if !remains.empty?
+            raise HandshakeError, "Extra bytes after header"
+          end
+        when 76
+          if remains.length < 8
+            # The whole third-key has not been received yet.
+            return nil
+          elsif remains.length > 8
+            raise HandshakeError, "Extra bytes after third key"
+          end
+          request['Third-Key'] = remains
+        else
+          raise WebSocketError, "Must not happen"
+        end
 
         unless request['Connection'] == 'Upgrade' and request['Upgrade'] == 'WebSocket'
           raise HandshakeError, "Connection and Upgrade headers required"
@@ -37,15 +59,18 @@ module EventMachine
         protocol = (secure ? "wss" : "ws")
         request['Host'] = Addressable::URI.parse("#{protocol}://"+request['Host'])
 
-        version = request['Sec-WebSocket-Key1'] ? 76 : 75
-
-        case version
-        when 75
-          Handler75.new(request, response, debug)
-        when 76
-          Handler76.new(request, response, debug)
+        if version = request['Sec-WebSocket-Draft']
+          if version == '1' || version == '2' || version == '3'
+            # We'll use handler03 - I believe they're all compatible
+            Handler03.new(connection, request, debug)
+          else
+            # According to spec should abort the connection
+            raise WebSocketError, "Unknown draft version: #{version}"
+          end
+        elsif request['Sec-WebSocket-Key1']
+          Handler76.new(connection, request, debug)
         else
-          raise WebSocketError, "Must not happen"
+          Handler75.new(connection, request, debug)
         end
       end
     end

data/lib/em-websocket/handshake75.rb

@@ -0,0 +1,21 @@
+module EventMachine
+  module WebSocket
+    module Handshake75
+      def handshake
+        location  = "#{request['Host'].scheme}://#{request['Host'].host}"
+        location << ":#{request['Host'].port}" if request['Host'].port
+        location << request['Path']
+
+        upgrade =  "HTTP/1.1 101 Web Socket Protocol Handshake\r\n"
+        upgrade << "Upgrade: WebSocket\r\n"
+        upgrade << "Connection: Upgrade\r\n"
+        upgrade << "WebSocket-Origin: #{request['Origin']}\r\n"
+        upgrade << "WebSocket-Location: #{location}\r\n\r\n"
+
+        debug [:upgrade_headers, upgrade]
+
+        return upgrade
+      end
+    end
+  end
+end

data/lib/em-websocket/handshake76.rb

@@ -0,0 +1,61 @@
+require 'digest/md5'
+
+module EventMachine
+  module WebSocket
+    module Handshake76
+      def handshake
+        challenge_response = solve_challenge(
+          request['Sec-WebSocket-Key1'],
+          request['Sec-WebSocket-Key2'],
+          request['Third-Key']
+        )
+
+        location  = "#{request['Host'].scheme}://#{request['Host'].host}"
+        location << ":#{request['Host'].port}" if request['Host'].port
+        location << request['Path']
+
+        upgrade =  "HTTP/1.1 101 WebSocket Protocol Handshake\r\n"
+        upgrade << "Upgrade: WebSocket\r\n"
+        upgrade << "Connection: Upgrade\r\n"
+        upgrade << "Sec-WebSocket-Location: #{location}\r\n"
+        upgrade << "Sec-WebSocket-Origin: #{request['Origin']}\r\n"
+        if protocol = request['Sec-WebSocket-Protocol']
+          validate_protocol!(protocol)
+          upgrade << "Sec-WebSocket-Protocol: #{protocol}\r\n"
+        end
+        upgrade << "\r\n"
+        upgrade << challenge_response
+
+        debug [:upgrade_headers, upgrade]
+
+        return upgrade
+      end
+
+      private
+
+      def solve_challenge(first, second, third)
+        # Refer to 5.2 4-9 of the draft 76
+        sum = [(extract_nums(first) / count_spaces(first))].pack("N*") +
+          [(extract_nums(second) / count_spaces(second))].pack("N*") +
+          third
+        Digest::MD5.digest(sum)
+      end
+
+      def extract_nums(string)
+        string.scan(/[0-9]/).join.to_i
+      end
+
+      def count_spaces(string)
+        spaces = string.scan(/ /).size
+        # As per 5.2.5, abort the connection if spaces are zero.
+        raise HandshakeError, "Websocket Key1 or Key2 does not contain spaces - this is a symptom of a cross-protocol attack" if spaces == 0
+        return spaces
+      end
+
+      def validate_protocol!(protocol)
+        raise HandshakeError, "Invalid WebSocket-Protocol: empty" if protocol.empty?
+        # TODO: Validate characters
+      end
+    end
+  end
+end

data/lib/em-websocket/version.rb

@@ -0,0 +1,5 @@
+module EventMachine
+  module Websocket
+    VERSION = "0.2.0"
+  end
+end

data/spec/helper.rb

@@ -1,9 +1,13 @@
 require 'rubygems'
-require 'spec'
+require 'rspec'
 require 'pp'
 require 'em-http'
 
-require 'lib/em-websocket'
+require 'em-websocket'
+
+Rspec.configure do |c|
+  c.mock_with :rspec
+end
 
 class FakeWebSocketClient < EM::Connection
   attr_writer :onopen, :onclose, :onmessage
@@ -21,7 +25,7 @@ class FakeWebSocketClient < EM::Connection
       @onopen.call if @onopen
       @state = :open
     else
-      @onmessage.call if @onmessage
+      @onmessage.call(data) if @onmessage
       @packets << data
     end
   end
@@ -55,11 +59,12 @@ def format_response(r)
 end
 
 def handler(request, secure = false)
-  EM::WebSocket::HandlerFactory.build(format_request(request), secure)
+  connection = Object.new
+  EM::WebSocket::HandlerFactory.build(connection, format_request(request), secure)
 end
 
-def send_handshake(response)
-  simple_matcher do |given|
-    given.handshake.lines.sort == format_response(response).lines.sort
+RSpec::Matchers.define :send_handshake do |response|
+  match do |actual|
+    actual.handshake.lines.sort == format_response(response).lines.sort
   end
 end