websocket-driver 0.7.1

data/ext/websocket-driver/WebsocketMaskService.java

@@ -0,0 +1,57 @@
+package com.jcoglan.websocket;
+
+import java.io.IOException;
+import org.jruby.Ruby;
+import org.jruby.RubyClass;
+import org.jruby.RubyModule;
+import org.jruby.RubyObject;
+import org.jruby.RubyString;
+import org.jruby.anno.JRubyMethod;
+import org.jruby.runtime.ObjectAllocator;
+import org.jruby.runtime.ThreadContext;
+import org.jruby.runtime.builtin.IRubyObject;
+import org.jruby.runtime.load.BasicLibraryService;
+
+public class WebsocketMaskService implements BasicLibraryService {
+    private Ruby runtime;
+
+    public boolean basicLoad(Ruby runtime) throws IOException {
+        this.runtime = runtime;
+
+        RubyModule websocket = runtime.defineModule("WebSocket");
+        RubyClass webSocketMask = websocket.defineClassUnder("Mask", runtime.getObject(), getAllocator());
+
+        webSocketMask.defineAnnotatedMethods(WebsocketMask.class);
+        return true;
+    }
+
+    ObjectAllocator getAllocator() {
+        return new ObjectAllocator() {
+            public IRubyObject allocate(Ruby runtime, RubyClass rubyClass) {
+                return new WebsocketMask(runtime, rubyClass);
+            }
+        };
+    }
+
+    public class WebsocketMask extends RubyObject {
+        public WebsocketMask(final Ruby runtime, RubyClass rubyClass) {
+            super(runtime, rubyClass);
+        }
+
+        @JRubyMethod
+        public IRubyObject mask(ThreadContext context, IRubyObject payload, IRubyObject mask) {
+            if (mask.isNil()) return payload;
+
+            byte[] payload_a = ((RubyString)payload).getBytes();
+            byte[] mask_a    = ((RubyString)mask).getBytes();
+            int i, n         = payload_a.length;
+
+            if (n == 0) return payload;
+
+            for (i = 0; i < n; i++) {
+                payload_a[i] ^= mask_a[i % 4];
+            }
+            return RubyString.newStringNoCopy(runtime, payload_a);
+        }
+    }
+}

data/ext/websocket-driver/extconf.rb

@@ -0,0 +1,4 @@
+require 'mkmf'
+extension_name = 'websocket_mask'
+dir_config(extension_name)
+create_makefile(extension_name)

data/ext/websocket-driver/websocket_mask.c

@@ -0,0 +1,32 @@
+#include <ruby.h>
+
+VALUE method_websocket_mask(VALUE self, VALUE payload, VALUE mask)
+{
+    char *payload_s, *mask_s, *unmasked_s;
+    long i, n;
+    VALUE unmasked;
+
+    if (mask == Qnil || RSTRING_LEN(mask) != 4) {
+        return payload;
+    }
+
+    payload_s = RSTRING_PTR(payload);
+    mask_s    = RSTRING_PTR(mask);
+    n         = RSTRING_LEN(payload);
+
+    unmasked   = rb_str_new(0, n);
+    unmasked_s = RSTRING_PTR(unmasked);
+
+    for (i = 0; i < n; i++) {
+        unmasked_s[i] = payload_s[i] ^ mask_s[i % 4];
+    }
+    return unmasked;
+}
+
+void Init_websocket_mask()
+{
+    VALUE WebSocket = rb_define_module("WebSocket");
+    VALUE Mask = rb_define_module_under(WebSocket, "Mask");
+
+    rb_define_singleton_method(Mask, "mask", method_websocket_mask, 2);
+}

data/lib/websocket/driver.rb

@@ -0,0 +1,233 @@
+# Protocol references:
+#
+# * http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-75
+# * http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76
+# * http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol-17
+
+require 'base64'
+require 'digest/md5'
+require 'digest/sha1'
+require 'securerandom'
+require 'set'
+require 'stringio'
+require 'uri'
+require 'websocket/extensions'
+
+module WebSocket
+  autoload :HTTP, File.expand_path('../http', __FILE__)
+
+  class Driver
+
+    root = File.expand_path('../driver', __FILE__)
+
+    begin
+      # Load C native extension
+      require 'websocket_mask'
+    rescue LoadError
+      # Fall back to pure-Ruby implementation
+      require 'websocket/mask'
+    end
+
+
+    if RUBY_PLATFORM =~ /java/
+      require 'jruby'
+      com.jcoglan.websocket.WebsocketMaskService.new.basicLoad(JRuby.runtime)
+    end
+
+    unless Mask.respond_to?(:mask)
+      def Mask.mask(payload, mask)
+        @instance ||= new
+        @instance.mask(payload, mask)
+      end
+    end
+
+    MAX_LENGTH = 0x3ffffff
+    STATES     = [:connecting, :open, :closing, :closed]
+
+    BINARY  = 'ASCII-8BIT'
+    UNICODE = 'UTF-8'
+
+    ConnectEvent = Struct.new(nil)
+    OpenEvent    = Struct.new(nil)
+    MessageEvent = Struct.new(:data)
+    PingEvent    = Struct.new(:data)
+    PongEvent    = Struct.new(:data)
+    CloseEvent   = Struct.new(:code, :reason)
+
+    ProtocolError      = Class.new(StandardError)
+    URIError           = Class.new(ArgumentError)
+    ConfigurationError = Class.new(ArgumentError)
+
+    autoload :Client,       root + '/client'
+    autoload :Draft75,      root + '/draft75'
+    autoload :Draft76,      root + '/draft76'
+    autoload :EventEmitter, root + '/event_emitter'
+    autoload :Headers,      root + '/headers'
+    autoload :Hybi,         root + '/hybi'
+    autoload :Proxy,        root + '/proxy'
+    autoload :Server,       root + '/server'
+    autoload :StreamReader, root + '/stream_reader'
+
+    include EventEmitter
+    attr_reader :protocol, :ready_state
+
+    def initialize(socket, options = {})
+      super()
+      Driver.validate_options(options, [:max_length, :masking, :require_masking, :protocols])
+
+      @socket      = socket
+      @reader      = StreamReader.new
+      @options     = options
+      @max_length  = options[:max_length] || MAX_LENGTH
+      @headers     = Headers.new
+      @queue       = []
+      @ready_state = 0
+    end
+
+    def state
+      return nil unless @ready_state >= 0
+      STATES[@ready_state]
+    end
+
+    def add_extension(extension)
+      false
+    end
+
+    def set_header(name, value)
+      return false unless @ready_state <= 0
+      @headers[name] = value
+      true
+    end
+
+    def start
+      return false unless @ready_state == 0
+
+      unless Driver.websocket?(@socket.env)
+        return fail_handshake(ProtocolError.new('Not a WebSocket request'))
+      end
+
+      begin
+        response = handshake_response
+      rescue => error
+        return fail_handshake(error)
+      end
+
+      @socket.write(response)
+      open unless @stage == -1
+      true
+    end
+
+    def text(message)
+      message = message.encode(UNICODE) unless message.encoding.name == UNICODE
+      frame(message, :text)
+    end
+
+    def binary(message)
+      false
+    end
+
+    def ping(*args)
+      false
+    end
+
+    def pong(*args)
+      false
+    end
+
+    def close(reason = nil, code = nil)
+      return false unless @ready_state == 1
+      @ready_state = 3
+      emit(:close, CloseEvent.new(nil, nil))
+      true
+    end
+
+  private
+
+    def fail_handshake(error)
+      headers = Headers.new
+      headers['Content-Type'] = 'text/plain'
+      headers['Content-Length'] = error.message.bytesize
+
+      headers = ['HTTP/1.1 400 Bad Request', headers.to_s, error.message]
+      @socket.write(headers.join("\r\n"))
+      fail(:protocol_error, error.message)
+
+      false
+    end
+
+    def fail(type, message)
+      @ready_state = 2
+      emit(:error, ProtocolError.new(message))
+      close
+    end
+
+    def open
+      @ready_state = 1
+      @queue.each { |message| frame(*message) }
+      @queue = []
+      emit(:open, OpenEvent.new)
+    end
+
+    def queue(message)
+      @queue << message
+      true
+    end
+
+    def self.client(socket, options = {})
+      Client.new(socket, options.merge(:masking => true))
+    end
+
+    def self.server(socket, options = {})
+      Server.new(socket, options.merge(:require_masking => true))
+    end
+
+    def self.rack(socket, options = {})
+      env     = socket.env
+      version = env['HTTP_SEC_WEBSOCKET_VERSION']
+      key     = env['HTTP_SEC_WEBSOCKET_KEY']
+      key1    = env['HTTP_SEC_WEBSOCKET_KEY1']
+      key2    = env['HTTP_SEC_WEBSOCKET_KEY2']
+
+      if version or key
+        Hybi.new(socket, options.merge(:require_masking => true))
+      elsif key1 or key2
+        Draft76.new(socket, options)
+      else
+        Draft75.new(socket, options)
+      end
+    end
+
+    def self.encode(string, encoding = nil)
+      case string
+        when Array then
+          string = string.pack('C*')
+          encoding ||= BINARY
+        when String then
+          encoding ||= UNICODE
+      end
+      unless string.encoding.name == encoding
+        string = string.dup if string.frozen?
+        string.force_encoding(encoding)
+      end
+      string.valid_encoding? ? string : nil
+    end
+
+    def self.validate_options(options, valid_keys)
+      options.keys.each do |key|
+        unless valid_keys.include?(key)
+          raise ConfigurationError, "Unrecognized option: #{key.inspect}"
+        end
+      end
+    end
+
+    def self.websocket?(env)
+      connection = env['HTTP_CONNECTION'] || ''
+      upgrade    = env['HTTP_UPGRADE']    || ''
+
+      env['REQUEST_METHOD'] == 'GET' and
+      connection.downcase.split(/ *, */).include?('upgrade') and
+      upgrade.downcase == 'websocket'
+    end
+
+  end
+end

data/lib/websocket/driver/client.rb

@@ -0,0 +1,140 @@
+module WebSocket
+  class Driver
+
+    class Client < Hybi
+      VALID_SCHEMES = %w[ws wss]
+
+      def self.generate_key
+        Base64.strict_encode64(SecureRandom.random_bytes(16))
+      end
+
+      attr_reader :status, :headers
+
+      def initialize(socket, options = {})
+        super
+
+        @ready_state = -1
+        @key         = Client.generate_key
+        @accept      = Hybi.generate_accept(@key)
+        @http        = HTTP::Response.new
+
+        uri = URI.parse(@socket.url)
+        unless VALID_SCHEMES.include?(uri.scheme)
+          raise URIError, "#{socket.url} is not a valid WebSocket URL"
+        end
+
+        host      = uri.host + (uri.port ? ":#{uri.port}" : '')
+        path      = (uri.path == '') ? '/' : uri.path
+        @pathname = path + (uri.query ? '?' + uri.query : '')
+
+        @headers['Host']                  = host
+        @headers['Upgrade']               = 'websocket'
+        @headers['Connection']            = 'Upgrade'
+        @headers['Sec-WebSocket-Key']     = @key
+        @headers['Sec-WebSocket-Version'] = VERSION
+
+        if @protocols.size > 0
+          @headers['Sec-WebSocket-Protocol'] = @protocols * ', '
+        end
+
+        if uri.user
+          auth = Base64.strict_encode64([uri.user, uri.password] * ':')
+          @headers['Authorization'] = 'Basic ' + auth
+        end
+      end
+
+      def version
+        "hybi-#{VERSION}"
+      end
+
+      def proxy(origin, options = {})
+        Proxy.new(self, origin, options)
+      end
+
+      def start
+        return false unless @ready_state == -1
+        @socket.write(handshake_request)
+        @ready_state = 0
+        true
+      end
+
+      def parse(chunk)
+        return if @ready_state == 3
+        return super if @ready_state > 0
+
+        @http.parse(chunk)
+        return fail_handshake('Invalid HTTP response') if @http.error?
+        return unless @http.complete?
+
+        validate_handshake
+        return if @ready_state == 3
+
+        open
+        parse(@http.body)
+      end
+
+    private 
+
+      def handshake_request
+        extensions = @extensions.generate_offer
+        @headers['Sec-WebSocket-Extensions'] = extensions if extensions
+
+        start   = "GET #{@pathname} HTTP/1.1"
+        headers = [start, @headers.to_s, '']
+        headers.join("\r\n")
+      end
+
+      def fail_handshake(message)
+        message = "Error during WebSocket handshake: #{message}"
+        @ready_state = 3
+        emit(:error, ProtocolError.new(message))
+        emit(:close, CloseEvent.new(ERRORS[:protocol_error], message))
+      end
+
+      def validate_handshake
+        @status  = @http.code
+        @headers = Headers.new(@http.headers)
+
+        unless @http.code == 101
+          return fail_handshake("Unexpected response code: #{@http.code}")
+        end
+
+        upgrade    = @http['Upgrade'] || ''
+        connection = @http['Connection'] || ''
+        accept     = @http['Sec-WebSocket-Accept'] || ''
+        protocol   = @http['Sec-WebSocket-Protocol'] || ''
+
+        if upgrade == ''
+          return fail_handshake("'Upgrade' header is missing")
+        elsif upgrade.downcase != 'websocket'
+          return fail_handshake("'Upgrade' header value is not 'WebSocket'")
+        end
+
+        if connection == ''
+          return fail_handshake("'Connection' header is missing")
+        elsif connection.downcase != 'upgrade'
+          return fail_handshake("'Connection' header value is not 'Upgrade'")
+        end
+
+        unless accept == @accept
+          return fail_handshake('Sec-WebSocket-Accept mismatch')
+        end
+
+        unless protocol == ''
+          if @protocols.include?(protocol)
+            @protocol = protocol
+          else
+            return fail_handshake('Sec-WebSocket-Protocol mismatch')
+          end
+        end
+
+        begin
+          @extensions.activate(@headers['Sec-WebSocket-Extensions'])
+        rescue ::WebSocket::Extensions::ExtensionError => error
+          return fail_handshake(error.message)
+        end
+      end
+    end
+
+  end
+end