Refactor proxy support into a separate class that represents the proxy connection as a real object.

README.md

@@ -197,17 +197,33 @@ sent back by the server:
 
 The client driver supports connections via HTTP proxies using the `CONNECT`
 method. Instead of sending the WebSocket handshake immediately, it will send a
-`CONNECT` request, wait for a `200` response, and then proceed as normal. To use
-this feature, set the `:proxy` option to the HTTP origin of the proxy, including
-any authorization information.
+`CONNECT` request, wait for a `200` response, and then proceed as normal.
+
+To use this feature, call `proxy = driver.proxy(url)` where `url` is the origin
+of the proxy, including a username and password if required. This produces an
+object that manages the process of connecting via the proxy. You should call
+`proxy.start` to begin the connection process, and pass data you receive via the
+socket to `proxy.parse(data)`. You should use these methods _instead_ of
+`driver.start` and `driver.parse(data)`.
+
+In the event that proxy connection fails, `proxy` will emit an `:error`. You can
+inspect the proxy's response via `proxy.status` and `proxy.headers`.
 
 ```rb
-driver = WebSocket::Driver.client(socket, :proxy => 'http://username:password@proxy.example.com')
+proxy.on :error do |error|
+  puts error.message
+  puts proxy.status
+  puts proxy.headers.inspect
+end
 ```
 
-It is up to you to set up a TCP connection to the proxy yourself. If you prefer,
-you can perform the `CONNECT` logic yourself and then hand off the TCP
-connection to the client driver to continue the handshake process.
+You can pass additional options to the proxy to control it. Before calling
+`proxy.start` you can set custom headers using `proxy.set_header`:
+
+```rb
+proxy.set_header('User-Agent', 'ruby')
+proxy.start
+```
 
 
 ### Driver API

lib/websocket/driver.rb

@@ -51,6 +51,7 @@ module WebSocket
     autoload :EventEmitter, root + '/event_emitter'
     autoload :Headers,      root + '/headers'
     autoload :Hybi,         root + '/hybi'
+    autoload :Proxy,        root + '/proxy'
     autoload :Server,       root + '/server'
 
     include EventEmitter

lib/websocket/driver/client.rb

@@ -2,8 +2,6 @@ module WebSocket
   class Driver
 
     class Client < Hybi
-      PORTS = {'ws' => 80, 'wss' => 443}
-
       def self.generate_key
         Base64.encode64((1..16).map { rand(255).chr } * '').strip
       end
@@ -13,8 +11,7 @@ module WebSocket
       def initialize(socket, options = {})
         super
 
-        @proxy       = options[:proxy]
-        @ready_state = @proxy ? -2 : -1
+        @ready_state = -1
         @key         = Client.generate_key
         @accept      = Hybi.generate_accept(@key)
         @http        = HTTP::Response.new
@@ -24,54 +21,29 @@ module WebSocket
         'hybi-13'
       end
 
+      def proxy(origin, options = {})
+        Proxy.new(self, origin, options)
+      end
+
       def start
-        return false unless @ready_state < 0
-        if @ready_state == -2
-          @socket.write(Driver.encode(proxy_connect_request, :binary))
-          @ready_state = -1
-        elsif @ready_state == -1
-          @socket.write(Driver.encode(handshake_request, :binary))
-          @ready_state = 0
-        end
+        return false unless @ready_state == -1
+        @socket.write(Driver.encode(handshake_request, :binary))
+        @ready_state = 0
         true
       end
 
       def parse(buffer)
         return super if @ready_state > 0
+
         @http.parse(buffer)
         return fail_handshake('Invalid HTTP response') if @http.error?
-        return unless @http.complete?
-
-        if @ready_state == -1
-          validate_proxy_response
-        else
-          validate_handshake
-        end
 
+        validate_handshake if @http.complete?
         parse(@http.body) if @ready_state == 1
       end
 
     private 
 
-      def proxy_connect_request
-        proxy = URI.parse(@proxy)
-        uri   = URI.parse(@socket.url)
-        port  = uri.port || PORTS[uri.scheme]
-
-        headers = [ "CONNECT #{uri.host}:#{port} HTTP/1.1",
-                    "Host: #{uri.host}",
-                    "Connection: keep-alive",
-                    "Proxy-Connection: keep-alive"
-                  ]
-
-        if proxy.user
-          auth = Base64.encode64([proxy.user, proxy.password] * ':').gsub(/\n/, '')
-          headers << "Proxy-Authorization: Basic #{auth}"
-        end
-
-        (headers + ['', '']).join("\r\n")
-      end
-
       def handshake_request
         uri   = URI.parse(@socket.url)
         host  = uri.host + (uri.port ? ":#{uri.port}" : '')
@@ -98,18 +70,6 @@ module WebSocket
         (headers + [@headers.to_s, '']).join("\r\n")
       end
 
-      def validate_proxy_response
-        if @http.code == 200
-          @http = HTTP::Response.new
-          start
-        else
-          message = "Can't establish a connection to the server at #{@socket.url}"
-          emit(:error, ProtocolError.new(message))
-          @ready_state = 3
-          emit(:close, CloseEvent.new(1006, ''))
-        end
-      end
-
       def fail_handshake(message)
         message = "Error during WebSocket handshake: #{message}"
         emit(:error, ProtocolError.new(message))

lib/websocket/driver/proxy.rb

@@ -0,0 +1,81 @@
+module WebSocket
+  class Driver
+
+    class Proxy
+      include EventEmitter
+
+      PORTS = {'ws' => 80, 'wss' => 443}
+
+      attr_reader :status, :headers
+
+      def initialize(client, origin, options)
+        super()
+
+        @client = client
+        @headers = Headers.new
+        @http    = HTTP::Response.new
+        @socket  = client.instance_variable_get(:@socket)
+        @origin  = URI.parse(@socket.url)
+        @url     = URI.parse(origin)
+        @options = options
+        @state   = 0
+      end
+
+      def set_header(name, value)
+        return false unless @state == 0
+        @headers[name] = value
+        true
+      end
+
+      def start
+        return false unless @state == 0
+        @state = 1
+
+        host = @origin.host + (@origin.port ? ":#{@origin.port}" : '')
+        port = @origin.port || PORTS[@origin.scheme]
+
+        headers = [ "CONNECT #{@origin.host}:#{port} HTTP/1.1",
+                    "Host: #{host}",
+                    "Connection: keep-alive",
+                    "Proxy-Connection: keep-alive"
+                  ]
+
+        if @url.user
+          auth = Base64.encode64([@url.user, @url.password] * ':').gsub(/\n/, '')
+          headers << "Proxy-Authorization: Basic #{auth}"
+        end
+
+        @socket.write((headers + [@headers.to_s, '']).join("\r\n"))
+        true
+      end
+
+      def parse(buffer)
+        return @delegate.parse(buffer) if @delegate
+
+        @http.parse(buffer)
+        return unless @http.complete?
+
+        @status = @http.code
+        @headers = Headers.new(@http.headers)
+
+        if @status != 200
+          message = "Can't establish a connection to the server at #{@socket.url}"
+          emit(:error, ProtocolError.new(message))
+          @ready_state = 3
+          return emit(:close, CloseEvent.new(1006, ''))
+        end
+
+        @socket.start_tls if @origin.scheme == 'wss'
+        configure_delegate(@client)
+        @client.start
+      end
+
+    private
+
+      def configure_delegate(delegate)
+        @delegate = delegate
+      end
+    end
+
+  end
+end

spec/websocket/driver/client_spec.rb

@@ -132,35 +132,6 @@ describe WebSocket::Driver::Client do
         end
       end
 
-      describe "using a proxy" do
-        let(:options) { {:proxy => "http://proxy.example.com"} }
-
-        it "sends a CONNECT request to the proxy" do
-          expect(socket).to receive(:write).with(
-              "CONNECT www.example.com:80 HTTP/1.1\r\n" +
-              "Host: www.example.com\r\n" +
-              "Connection: keep-alive\r\n" +
-              "Proxy-Connection: keep-alive\r\n" +
-              "\r\n")
-          driver.start
-        end
-      end
-
-      describe "using an authenticated proxy" do
-        let(:options) { {:proxy => "http://user:pass@proxy.example.com"} }
-
-        it "sends a CONNECT request to the proxy" do
-          expect(socket).to receive(:write).with(
-              "CONNECT www.example.com:80 HTTP/1.1\r\n" +
-              "Host: www.example.com\r\n" +
-              "Connection: keep-alive\r\n" +
-              "Proxy-Connection: keep-alive\r\n" +
-              "Proxy-Authorization: Basic dXNlcjpwYXNz\r\n" +
-              "\r\n")
-          driver.start
-        end
-      end
-
       it "changes the state to :connecting" do
         driver.start
         expect(driver.state).to eq :connecting
@@ -168,32 +139,74 @@ describe WebSocket::Driver::Client do
     end
   end
 
-  describe "in the :connecting state" do
-    before { driver.start }
+  describe "using a proxy" do
+    it "sends a CONNECT request" do
+      proxy = driver.proxy("http://proxy.example.com")
+      expect(socket).to receive(:write).with(
+          "CONNECT www.example.com:80 HTTP/1.1\r\n" +
+          "Host: www.example.com\r\n" +
+          "Connection: keep-alive\r\n" +
+          "Proxy-Connection: keep-alive\r\n" +
+          "\r\n")
+      proxy.start
+    end
 
-    describe "using a proxy" do
-      let(:options) { {:proxy => "http://proxy.example.com"} }
+    it "sends an authenticated CONNECT request" do
+      proxy = driver.proxy("http://user:pass@proxy.example.com")
+      expect(socket).to receive(:write).with(
+          "CONNECT www.example.com:80 HTTP/1.1\r\n" +
+          "Host: www.example.com\r\n" +
+          "Connection: keep-alive\r\n" +
+          "Proxy-Connection: keep-alive\r\n" +
+          "Proxy-Authorization: Basic dXNlcjpwYXNz\r\n" +
+          "\r\n")
+      proxy.start
+    end
 
-      it "writes the handshake request when the proxy connects" do
+    it "sends a CONNECT request with custom headers" do
+      proxy = driver.proxy("http://proxy.example.com")
+      proxy.set_header("User-Agent", "Chrome")
+      expect(socket).to receive(:write).with(
+          "CONNECT www.example.com:80 HTTP/1.1\r\n" +
+          "Host: www.example.com\r\n" +
+          "Connection: keep-alive\r\n" +
+          "Proxy-Connection: keep-alive\r\n" +
+          "User-Agent: Chrome\r\n" +
+          "\r\n")
+      proxy.start
+    end
+
+    describe "receiving a response" do
+      let(:proxy) { driver.proxy("http://proxy.example.com") }
+
+      before do
+        proxy.on(:error)   { |e| @error = e }
+        proxy.on(:close)   { |e| @close = [e.code, e.reason] }
+      end
+
+      it "emits a WebSocket handshake when the proxy connects" do
         expect(socket).to receive(:write).with(
-            "GET /socket HTTP/1.1\r\n" + 
+            "GET /socket HTTP/1.1\r\n" +
             "Host: www.example.com\r\n" +
             "Upgrade: websocket\r\n" +
             "Connection: Upgrade\r\n" +
             "Sec-WebSocket-Key: 2vBVWg4Qyk3ZoM/5d3QD9Q==\r\n" +
             "Sec-WebSocket-Version: 13\r\n" +
             "\r\n")
-        driver.parse("HTTP/1.1 200 OK\r\n\r\n")
+        proxy.parse("HTTP/1.1 200 OK\r\n\r\n")
       end
 
-      it "emits an error if the proxy does not connect" do
-        expect(socket).to_not receive(:write)
-        driver.parse("HTTP/1.1 403 Forbidden\r\n\r\n")
+      it "emits an 'error' event if the proxy does not connect" do
+        proxy.parse("HTTP/1.1 403 Forbidden\r\n\r\n")
+        expect(@open).to eq false
         expect(@error.message).to eq "Can't establish a connection to the server at ws://www.example.com/socket"
         expect(@close).to eq [1006, ""]
-        expect(driver.state).to eq :closed
       end
     end
+  end
+
+  describe "in the :connecting state" do
+    before { driver.start }
 
     describe "with a valid response" do
       before { driver.parse(response) }