websocket-rack 0.1.0

data/lib/rack/websocket/framing76.rb

@@ -0,0 +1,115 @@
+# encoding: BINARY
+
+module Rack
+  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(newdata)
+        debug [:message, @data]
+
+        # This algorithm comes straight from the spec
+        # http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76#section-5.3
+
+        error = false
+
+        while !error
+          return if @data.size == 0
+
+          pointer = 0
+          frame_type = @data.getbyte(pointer)
+          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.getbyte(pointer)
+              b = @data.getbyte(pointer)
+              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.getbyte(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
+
+            if @data.getbyte(0) != 0x00
+              # Close the connection since this buffer can never match
+              @connection.close_with_error(DataError.new("Invalid frame received"))
+            end
+
+            # Addition to the spec to protect against malicious requests
+            if @data.size > MAXIMUM_FRAME_LENGTH
+              @connection.close_with_error(DataError.new("Frame length too long (#{@data.size} bytes)"))
+              return false
+            end
+
+            # Optimization to avoid calling slice! unnecessarily
+            error = true and next unless newdata =~ /\xff/
+
+            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/rack/websocket/handler.rb

@@ -0,0 +1,43 @@
+module Rack
+  module WebSocket
+    class Handler
+      include Debugger
+
+      attr_reader :request, :state
+
+      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(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/rack/websocket/handler03.rb

@@ -0,0 +1,14 @@
+module Rack
+  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/rack/websocket/handler75.rb

@@ -0,0 +1,8 @@
+module Rack
+  module WebSocket
+    class Handler75 < Handler
+      include Handshake75
+      include Framing76
+    end
+  end
+end

data/lib/rack/websocket/handler76.rb

@@ -0,0 +1,11 @@
+module Rack
+  module WebSocket
+    class Handler76 < Handler
+      include Handshake76
+      include Framing76
+
+      # "\377\000" is octet version and "\xff\x00" is hex version
+      TERMINATE_STRING = "\xff\x00"
+    end
+  end
+end

data/lib/rack/websocket/handler_factory.rb

@@ -0,0 +1,61 @@
+module Rack
+  module WebSocket
+    class HandlerFactory
+
+      def self.build(connection, data, debug = false)
+        request = Rack::Request.new(data)
+
+        unless request.env['rack.input'].nil?
+          request.env['rack.input'].rewind
+          remains = request.env['rack.input'].read
+        else
+          # The whole header has not been received yet.
+          return nil
+        end
+
+        unless request.get?
+          raise HandshakeError, "Must be GET request"
+        end
+
+        version = request.env['HTTP_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.env['HTTP_THIRD_KEY'] = remains
+        else
+          raise WebSocketError, "Must not happen"
+        end
+
+        unless request.env['HTTP_CONNECTION'] == 'Upgrade' and request.env['HTTP_UPGRADE'] == 'WebSocket'
+          raise HandshakeError, "Connection and Upgrade headers required"
+        end
+
+        # transform headers
+        request.env['rack.url_scheme'] = (request.scheme == 'https' ? "wss" : "ws")
+
+        if version = request.env['HTTP_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.env['HTTP_SEC_WEBSOCKET_KEY1']
+          Handler76.new(connection, request, debug)
+        else
+          Handler75.new(connection, request, debug)
+        end
+      end
+    end
+  end
+end

data/lib/rack/websocket/handshake75.rb

@@ -0,0 +1,21 @@
+module Rack
+  module WebSocket
+    module Handshake75
+      def handshake
+        location  = "#{request.env['rack.url_scheme']}://#{request.host}"
+        location << ":#{request.port}" if request.port > 0
+        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.env['HTTP_ORIGIN']}\r\n"
+        upgrade << "WebSocket-Location: #{location}\r\n\r\n"
+
+        debug [:upgrade_headers, upgrade]
+
+        return upgrade
+      end
+    end
+  end
+end

data/lib/rack/websocket/handshake76.rb

@@ -0,0 +1,71 @@
+require 'digest/md5'
+
+module Rack
+  module WebSocket
+    module Handshake76
+      def handshake
+        challenge_response = solve_challenge(
+          request.env['HTTP_SEC_WEBSOCKET_KEY1'],
+          request.env['HTTP_SEC_WEBSOCKET_KEY2'],
+          request.env['HTTP_THIRD_KEY']
+        )
+
+        location  = "#{request.env['rack.url_scheme']}://#{request.host}"
+        location << ":#{request.port}" if request.port > 0
+        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.env['HTTP_ORIGIN']}\r\n"
+        if protocol = request.env['HTTP_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 = [numbers_over_spaces(first)].pack("N*") +
+          [numbers_over_spaces(second)].pack("N*") +
+          third
+        Digest::MD5.digest(sum)
+      end
+
+      def numbers_over_spaces(string)
+        numbers = string.scan(/[0-9]/).join.to_i
+
+        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
+
+        # As per 5.2.6, abort if numbers is not an integral multiple of spaces
+        if numbers % spaces != 0
+          raise HandshakeError, "Invalid Key #{string.inspect}"
+        end
+
+        quotient = numbers / spaces
+
+        if quotient > 2**32-1
+          raise HandshakeError, "Challenge computation out of range for key #{string.inspect}"
+        end
+
+        return quotient
+      end
+
+      def validate_protocol!(protocol)
+        raise HandshakeError, "Invalid WebSocket-Protocol: empty" if protocol.empty?
+        # TODO: Validate characters
+      end
+    end
+  end
+end

data/lib/rack/websocket.rb

@@ -0,0 +1,40 @@
+require 'rack'
+
+module Rack
+  module WebSocket
+    VERSION = "0.1.0"
+    ROOT_PATH = ::File.expand_path(::File.dirname(__FILE__))
+
+    class WebSocketError < RuntimeError; end
+    class HandshakeError < WebSocketError; end
+    class DataError < WebSocketError; end
+
+    autoload :Application,    "#{ROOT_PATH}/websocket/application"
+    autoload :Connection,     "#{ROOT_PATH}/websocket/connection"
+    autoload :Debugger,       "#{ROOT_PATH}/websocket/debugger"
+    autoload :Framing03,      "#{ROOT_PATH}/websocket/framing03"
+    autoload :Framing76,      "#{ROOT_PATH}/websocket/framing76"
+    autoload :Handler,        "#{ROOT_PATH}/websocket/handler"
+    autoload :Handler03,      "#{ROOT_PATH}/websocket/handler03"
+    autoload :Handler75,      "#{ROOT_PATH}/websocket/handler75"
+    autoload :Handler76,      "#{ROOT_PATH}/websocket/handler76"
+    autoload :HandlerFactory, "#{ROOT_PATH}/websocket/handler_factory"
+    autoload :Handshake75,    "#{ROOT_PATH}/websocket/handshake75"
+    autoload :Handshake76,    "#{ROOT_PATH}/websocket/handshake76"
+
+    module Extensions
+      autoload :Thin,         "#{ROOT_PATH}/websocket/extensions/thin"
+    end
+
+  end
+end
+
+::Thin.send(:include, ::Rack::WebSocket::Extensions::Thin) if defined?(Thin)
+
+unless ''.respond_to?(:getbyte)
+  class String
+    def getbyte(i)
+      self[i]
+    end
+  end
+end

data/spec/helper.rb

@@ -0,0 +1,44 @@
+require 'rubygems'
+require 'rspec'
+require 'pp'
+require 'stringio'
+
+require 'rack/websocket'
+
+Rspec.configure do |c|
+  c.mock_with :rspec
+end
+
+def format_request(r)
+  data = {}
+  data['REQUEST_METHOD'] = r[:method] if r[:method]
+  data['PATH_INFO'] = r[:path] if r[:path]
+  data['SERVER_PORT'] = r[:port] if r[:port] && r[:port] != 80
+  r[:headers].each do |key, value|
+    data['HTTP_' + key.upcase.gsub('-','_')] = value
+  end
+  data['rack.input'] = StringIO.new(r[:body]) if r[:body]
+  # data = "#{r[:method]} #{r[:path]} HTTP/1.1\r\n"
+  # header_lines = r[:headers].map { |k,v| "#{k}: #{v}" }
+  # data << [header_lines, '', r[:body]].join("\r\n")
+  data
+end
+
+def format_response(r)
+  data = "HTTP/1.1 101 WebSocket Protocol Handshake\r\n"
+  header_lines = r[:headers].map { |k,v| "#{k}: #{v}" }
+  data << [header_lines, '', r[:body]].join("\r\n")
+  data
+end
+
+def handler(request, secure = false)
+  connection = Object.new
+  secure_hash = secure ? {'rack.url_scheme' => 'https'} : {}
+  Rack::WebSocket::HandlerFactory.build(connection, format_request(request).merge(secure_hash))
+end
+
+RSpec::Matchers.define :send_handshake do |response|
+  match do |actual|
+    actual.handshake.lines.sort == format_response(response).lines.sort
+  end
+end

data/spec/integration/draft03_spec.rb

@@ -0,0 +1,252 @@
+# require 'helper'
+# 
+# describe "draft03" do
+#   before :each do
+#     @request = {
+#       :port => 80,
+#       :method => "GET",
+#       :path => "/demo",
+#       :headers => {
+#         'Host' => 'example.com',
+#         'Connection' => 'Upgrade',
+#         'Sec-WebSocket-Key2' => '12998 5 Y3 1  .P00',
+#         'Sec-WebSocket-Protocol' => 'sample',
+#         'Upgrade' => 'WebSocket',
+#         'Sec-WebSocket-Key1' => '4 @1  46546xW%0l 1 5',
+#         'Origin' => 'http://example.com',
+#         'Sec-WebSocket-Draft' => '3'
+#       },
+#       :body => '^n:ds[4U'
+#     }
+# 
+#     @response = {
+#       :headers => {
+#         "Upgrade" => "WebSocket",
+#         "Connection" => "Upgrade",
+#         "Sec-WebSocket-Location" => "ws://example.com/demo",
+#         "Sec-WebSocket-Origin" => "http://example.com",
+#         "Sec-WebSocket-Protocol" => "sample"
+#       },
+#       :body => "8jKS\'y:G*Co,Wxa-"
+#     }
+#   end
+# 
+#   # These examples are straight from the spec
+#   # http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol-03#section-4.6
+#   describe "examples from the spec" do
+#     it "should accept a single-frame text message" do
+#       EM.run do
+#         EM::WebSocket.start(:host => "0.0.0.0", :port => 12345) { |ws|
+#           ws.onmessage { |msg|
+#             msg.should == 'Hello'
+#             EM.stop
+#           }
+#         }
+# 
+#         # Create a fake client which sends draft 76 handshake
+#         connection = EM.connect('0.0.0.0', 12345, FakeWebSocketClient)
+#         connection.send_data(format_request(@request))
+# 
+#         # Send frame
+#         connection.onopen = lambda {
+#           connection.send_data("\x04\x05Hello")
+#         }
+#       end
+#     end
+#     
+#     it "should accept a fragmented text message" do
+#       EM.run do
+#         EM::WebSocket.start(:host => "0.0.0.0", :port => 12345) { |ws|
+#           ws.onmessage { |msg|
+#             msg.should == 'Hello'
+#             EM.stop
+#           }
+#         }
+# 
+#         # Create a fake client which sends draft 76 handshake
+#         connection = EM.connect('0.0.0.0', 12345, FakeWebSocketClient)
+#         connection.send_data(format_request(@request))
+# 
+#         # Send frame
+#         connection.onopen = lambda {
+#           connection.send_data("\x84\x03Hel")
+#           connection.send_data("\x00\x02lo")
+#         }
+#       end
+#     end
+#     
+#     it "should accept a ping request and respond with the same body" do
+#       EM.run do
+#         EM::WebSocket.start(:host => "0.0.0.0", :port => 12345) { |ws| }
+# 
+#         # Create a fake client which sends draft 76 handshake
+#         connection = EM.connect('0.0.0.0', 12345, FakeWebSocketClient)
+#         connection.send_data(format_request(@request))
+# 
+#         # Send frame
+#         connection.onopen = lambda {
+#           connection.send_data("\x02\x05Hello")
+#         }
+#         
+#         connection.onmessage = lambda { |frame|
+#           next if frame.nil?
+#           frame.should == "\x03\x05Hello"
+#           EM.stop
+#         }
+#       end
+#     end
+#     
+#     it "should accept a 256 bytes binary message in a single frame" do
+#       EM.run do
+#         data = "a" * 256
+#         
+#         EM::WebSocket.start(:host => "0.0.0.0", :port => 12345) { |ws|
+#           ws.onmessage { |msg|
+#             msg.should == data
+#             EM.stop
+#           }
+#         }
+# 
+#         # Create a fake client which sends draft 76 handshake
+#         connection = EM.connect('0.0.0.0', 12345, FakeWebSocketClient)
+#         connection.send_data(format_request(@request))
+# 
+#         # Send frame
+#         connection.onopen = lambda {
+#           connection.send_data("\x05\x7E\x01\x00" + data)
+#         }
+#       end
+#     end
+#     
+#     it "should accept a 64KiB binary message in a single frame" do
+#       EM.run do
+#         data = "a" * 65536
+#         
+#         EM::WebSocket.start(:host => "0.0.0.0", :port => 12345) { |ws|
+#           ws.onmessage { |msg|
+#             msg.should == data
+#             EM.stop
+#           }
+#         }
+# 
+#         # Create a fake client which sends draft 76 handshake
+#         connection = EM.connect('0.0.0.0', 12345, FakeWebSocketClient)
+#         connection.send_data(format_request(@request))
+# 
+#         # Send frame
+#         connection.onopen = lambda {
+#           connection.send_data("\x05\x7F\x00\x00\x00\x00\x00\x01\x00\x00" + data)
+#         }
+#       end
+#     end
+#   end
+# 
+#   describe "close handling" do
+#     it "should respond to a new close frame with a close frame" do
+#       EM.run do
+#         EM::WebSocket.start(:host => "0.0.0.0", :port => 12345) { |ws| }
+# 
+#         # Create a fake client which sends draft 76 handshake
+#         connection = EM.connect('0.0.0.0', 12345, FakeWebSocketClient)
+#         connection.send_data(format_request(@request))
+# 
+#         # Send close frame
+#         connection.onopen = lambda {
+#           connection.send_data("\x01\x00")
+#         }
+# 
+#         # Check that close ack received
+#         connection.onmessage = lambda { |frame|
+#           frame.should == "\x01\x00"
+#           EM.stop
+#         }
+#       end
+#     end
+# 
+#     it "should close the connection on receiving a close acknowlegement" do
+#       EM.run do
+#         ack_received = false
+# 
+#         EM::WebSocket.start(:host => "0.0.0.0", :port => 12345) { |ws|
+#           ws.onopen {
+#             # 2. Send a close frame
+#             EM.next_tick {
+#               ws.close_websocket
+#             }
+#           }
+#         }
+# 
+#         # 1. Create a fake client which sends draft 76 handshake
+#         connection = EM.connect('0.0.0.0', 12345, FakeWebSocketClient)
+#         connection.send_data(format_request(@request))
+# 
+#         # 3. Check that close frame recieved and acknowlege it
+#         connection.onmessage = lambda { |frame|
+#           frame.should == "\x01\x00"
+#           ack_received = true
+#           connection.send_data("\x01\x00")
+#         }
+# 
+#         # 4. Check that connection is closed _after_ the ack
+#         connection.onclose = lambda {
+#           ack_received.should == true
+#           EM.stop
+#         }
+#       end
+#     end
+# 
+#     it "should not allow data frame to be sent after close frame sent" do
+#       EM.run {
+#         EM::WebSocket.start(:host => "0.0.0.0", :port => 12345) { |ws|
+#           ws.onopen {
+#             # 2. Send a close frame
+#             EM.next_tick {
+#               ws.close_websocket
+#             }
+# 
+#             # 3. Check that exception raised if I attempt to send more data
+#             EM.add_timer(0.1) {
+#               lambda {
+#                 ws.send('hello world')
+#               }.should raise_error(EM::WebSocket::WebSocketError, 'Cannot send data frame since connection is closing')
+#               EM.stop
+#             }
+#           }
+#         }
+# 
+#         # 1. Create a fake client which sends draft 76 handshake
+#         connection = EM.connect('0.0.0.0', 12345, FakeWebSocketClient)
+#         connection.send_data(format_request(@request))
+#       }
+#     end
+# 
+#     it "should still respond to control frames after close frame sent" do
+#       EM.run {
+#         EM::WebSocket.start(:host => "0.0.0.0", :port => 12345) { |ws|
+#           ws.onopen {
+#             # 2. Send a close frame
+#             EM.next_tick {
+#               ws.close_websocket
+#             }
+#           }
+#         }
+# 
+#         # 1. Create a fake client which sends draft 76 handshake
+#         connection = EM.connect('0.0.0.0', 12345, FakeWebSocketClient)
+#         connection.send_data(format_request(@request))
+# 
+#         connection.onmessage = lambda { |frame|
+#           if frame == "\x01\x00"
+#             # 3. After the close frame is received send a ping frame, but
+#             # don't respond with a close ack
+#             connection.send_data("\x02\x05Hello")
+#           else
+#             # 4. Check that the pong is received
+#             frame.should == "\x03\x05Hello"
+#             EM.stop
+#           end
+#         }
+#       }
+#     end
+#   end
+# end