Support for wss/ssl in client

lib/web_socket.rb

@@ -5,23 +5,24 @@
 require "socket"
 require "uri"
 require "digest/md5"
+require 'openssl'
 
 
 class WebSocket
-    
+
     class << self
-        
+
         attr_accessor(:debug)
-        
+
     end
-    
+
     class Error < RuntimeError
-    
+
     end
-    
+
     def initialize(arg, params = {})
       if params[:server] # server
-        
+
         @server = params[:server]
         @socket = arg
         line = gets().chomp()
@@ -45,23 +46,28 @@ def initialize(arg, params = {})
             [self.origin, @server.accepted_domains, @server.origin_to_domain(self.origin)])
         end
         @handshaked = false
-        
+
       else # client
-        
+
         uri = arg.is_a?(String) ? URI.parse(arg) : arg
-        if uri.scheme == "wss"
-          raise(WebSocket::Error, "wss scheme is unimplemented")
-        elsif uri.scheme != "ws"
-          raise(WebSocket::Error, "unsupported scheme: #{uri.scheme}")
-        end
+
+        raise(WebSocket::Error, "unsupported scheme: #{uri.scheme}") unless ["ws","wss"].include?(uri.scheme)
+
         @path = (uri.path.empty? ? "/" : uri.path) + (uri.query ? "?" + uri.query : "")
         host = uri.host + (uri.port == 80 ? "" : ":#{uri.port}")
         origin = params[:origin] || "http://#{uri.host}"
         key1 = generate_key()
         key2 = generate_key()
         key3 = generate_key3()
-
-        @socket = TCPSocket.new(uri.host, uri.port || 80)
+
+        socket = TCPSocket.new(uri.host, uri.port || 80)
+
+        if uri.scheme == "ws"
+          @socket = socket
+        else
+          @socket = ssl_handshake(socket)
+        end
+
         write(
           "GET #{@path} HTTP/1.1\r\n" +
           "Upgrade: WebSocket\r\n" +
@@ -73,7 +79,7 @@ def initialize(arg, params = {})
           "\r\n" +
           "#{key3}")
         flush()
-        
+
         line = gets().chomp()
         raise(WebSocket::Error, "bad response: #{line}") if !(line =~ /\AHTTP\/1.1 101 /n)
         read_header()
@@ -88,14 +94,14 @@ def initialize(arg, params = {})
             "security digest doesn't match: %p != %p" % [reply_digest, expected_digest])
         end
         @handshaked = true
-        
+
       end
       @received = []
       @buffer = ""
     end
-    
+
     attr_reader(:server, :header, :path)
-    
+
     def handshake(status = nil, header = {})
       if @handshaked
         raise(WebSocket::Error, "handshake has already been done")
@@ -118,7 +124,7 @@ def handshake(status = nil, header = {})
       flush()
       @handshaked = true
     end
-    
+
     def send(data)
       if !@handshaked
         raise(WebSocket::Error, "call WebSocket\#handshake first")
@@ -127,7 +133,7 @@ def send(data)
       write("\x00#{data}\xff")
       flush()
     end
-    
+
     def receive()
       if !@handshaked
         raise(WebSocket::Error, "call WebSocket\#handshake first")
@@ -139,31 +145,31 @@ def receive()
       end
       return force_encoding($1, "UTF-8")
     end
-    
+
     def tcp_socket
       return @socket
     end
-    
+
     def host
       return @header["Host"]
     end
-    
+
     def origin
       return @header["Origin"]
     end
-    
+
     def location
       return "ws://#{self.host}#{@path}"
     end
-    
+
     def close()
       @socket.close()
     end
 
   private
-    
+
     NOISE_CHARS = ("\x21".."\x2f").to_a() + ("\x3a".."\x7e").to_a()
-    
+
     def read_header()
       @header = {}
       while line = gets()
@@ -181,19 +187,19 @@ def read_header()
         raise(WebSocket::Error, "invalid Connection: " + @header["Connection"])
       end
     end
-    
+
     def gets(rs = $/)
       line = @socket.gets(rs)
       $stderr.printf("recv> %p\n", line) if WebSocket.debug
       return line
     end
-    
+
     def read(num_bytes)
       str = @socket.read(num_bytes)
       $stderr.printf("recv> %p\n", str) if WebSocket.debug
       return str
     end
-    
+
     def write(data)
       if WebSocket.debug
         data.scan(/\G(.*?(\n|\z))/n) do
@@ -202,17 +208,17 @@ def write(data)
       end
       @socket.write(data)
     end
-    
+
     def flush()
       @socket.flush()
     end
-    
+
     def security_digest(key1, key2, key3)
       bytes1 = websocket_key_to_bytes(key1)
       bytes2 = websocket_key_to_bytes(key2)
       return Digest::MD5.digest(bytes1 + bytes2 + key3)
     end
-    
+
     def generate_key()
       spaces = 1 + rand(12)
       max = 0xffffffff / spaces
@@ -229,29 +235,37 @@ def generate_key()
       end
       return key
     end
-    
+
     def generate_key3()
       return [rand(0x100000000)].pack("N") + [rand(0x100000000)].pack("N")
     end
-    
+
     def websocket_key_to_bytes(key)
       num = key.gsub(/[^\d]/n, "").to_i() / key.scan(/ /).size
       return [num].pack("N")
     end
-    
+
     def force_encoding(str, encoding)
       if str.respond_to?(:force_encoding)
         return str.force_encoding(encoding)
       else
         return str
       end
     end
-
+
+    def ssl_handshake(socket)
+      ssl_context = OpenSSL::SSL::SSLContext.new()
+      ssl_socket = OpenSSL::SSL::SSLSocket.new(socket, ssl_context)
+      ssl_socket.sync_close = true
+      ssl_socket.connect
+      ssl_socket
+    end
+
 end
 
 
 class WebSocketServer
-    
+
     def initialize(params_or_uri, params = nil)
       if params
         uri = params_or_uri.is_a?(String) ? URI.parse(params_or_uri) : params_or_uri
@@ -271,9 +285,9 @@ def initialize(params_or_uri, params = nil)
         @tcp_server = TCPServer.open(@port)
       end
     end
-    
+
     attr_reader(:tcp_server, :port, :accepted_domains)
-    
+
     def run(&block)
       while true
         Thread.start(accept()) do |s|
@@ -291,20 +305,20 @@ def run(&block)
         end
       end
     end
-    
+
     def accept()
       return @tcp_server.accept()
     end
-    
+
     def accepted_origin?(origin)
       domain = origin_to_domain(origin)
       return @accepted_domains.any?(){ |d| File.fnmatch(d, domain) }
     end
-    
+
     def origin_to_domain(origin)
       return origin == "file://" ? "file://" : URI.parse(origin).host
     end
-    
+
     def create_web_socket(socket)
       ch = socket.getc()
       if ch == ?<
@@ -316,16 +330,16 @@ def create_web_socket(socket)
         return WebSocket.new(socket, :server => self)
       end
     end
-    
+
   private
-    
+
     def print_backtrace(ex)
       $stderr.printf("%s: %s (%p)\n", ex.backtrace[0], ex.message, ex.class)
       for s in ex.backtrace[1..-1]
         $stderr.printf("        %s\n", s)
       end
     end
-    
+
     # Handles Flash socket policy file request sent when web-socket-js is used:
     # http://github.com/gimite/web-socket-js/tree/master
     def send_flash_socket_policy_file(socket)
@@ -346,9 +360,9 @@ def send_flash_socket_policy_file(socket)
 
 if __FILE__ == $0
   Thread.abort_on_exception = true
-  
+
   if ARGV[0] == "server" && ARGV.size == 3
-    
+
     server = WebSocketServer.new(
       :accepted_domains => [ARGV[1]],
       :port => ARGV[2].to_i())
@@ -368,9 +382,9 @@ def send_flash_socket_policy_file(socket)
       end
       puts("Connection closed")
     end
-    
+
   elsif ARGV[0] == "client" && ARGV.size == 2
-    
+
     client = WebSocket.new(ARGV[1])
     puts("Connected")
     Thread.new() do
@@ -383,13 +397,13 @@ def send_flash_socket_policy_file(socket)
       client.send(data)
       printf("Sent: %p\n", data)
     end
-    
+
   else
-    
+
     $stderr.puts("Usage:")
     $stderr.puts("  ruby web_socket.rb server ACCEPTED_DOMAIN PORT")
     $stderr.puts("  ruby web_socket.rb client ws://HOST:PORT/")
     exit(1)
-    
+
   end
 end