Add API for setting custom handshake headers.

README.md

@@ -238,6 +238,12 @@ describing the error.
 Sets the callback block to execute when the socket becomes closed. The `event`
 object has `code` and `reason` attributes.
 
+#### `driver.set_header(name, value)`
+
+Sets a custom header to be sent as part of the handshake response, either from
+the server or from the client. Must be called before `start`, since this is
+when the headers are serialized and sent.
+
 #### `driver.start`
 
 Initiates the protocol by sending the handshake - either the response for a

lib/websocket/driver.rb

@@ -7,6 +7,7 @@
 require 'base64'
 require 'digest/md5'
 require 'digest/sha1'
+require 'set'
 require 'stringio'
 require 'uri'
 
@@ -47,6 +48,7 @@ module WebSocket
     autoload :Draft75,      root + '/draft75'
     autoload :Draft76,      root + '/draft76'
     autoload :EventEmitter, root + '/event_emitter'
+    autoload :Headers,      root + '/headers'
     autoload :Hybi,         root + '/hybi'
     autoload :Server,       root + '/server'
 
@@ -58,6 +60,7 @@ module WebSocket
 
       @socket      = socket
       @options     = options
+      @headers     = Headers.new
       @queue       = []
       @ready_state = 0
     end
@@ -67,6 +70,12 @@ module WebSocket
       STATES[@ready_state]
     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
       @socket.write(handshake_response)

lib/websocket/driver/client.rb

@@ -54,7 +54,7 @@ module WebSocket
           headers << "Sec-WebSocket-Protocol: #{@protocols * ', '}"
         end
 
-        (headers + ['', '']).join("\r\n")
+        (headers + [@headers.to_s, '']).join("\r\n")
       end
 
       def fail_handshake(message)

lib/websocket/driver/draft75.rb

@@ -71,6 +71,7 @@ module WebSocket
         upgrade << "Connection: Upgrade\r\n"
         upgrade << "WebSocket-Origin: #{@socket.env['HTTP_ORIGIN']}\r\n"
         upgrade << "WebSocket-Location: #{@socket.url}\r\n"
+        upgrade << @headers.to_s
         upgrade << "\r\n"
         upgrade
       end

lib/websocket/driver/draft76.rb

@@ -37,6 +37,7 @@ module WebSocket
         upgrade << "Connection: Upgrade\r\n"
         upgrade << "Sec-WebSocket-Origin: #{@socket.env['HTTP_ORIGIN']}\r\n"
         upgrade << "Sec-WebSocket-Location: #{@socket.url}\r\n"
+        upgrade << @headers.to_s
         upgrade << "\r\n"
         upgrade
       end

lib/websocket/driver/headers.rb

@@ -0,0 +1,26 @@
+module WebSocket
+  class Driver
+
+    class Headers
+      ALLOWED_DUPLICATES = %w(set-cookie set-cookie2 warning www-authenticate)
+
+      def initialize
+        @sent  = Set.new
+        @lines = []
+      end
+
+      def []=(name, value)
+        return if value.nil?
+        key = HTTP.normalize_header(name)
+        return unless @sent.add?(key) or ALLOWED_DUPLICATES.include?(key)
+        @lines << "#{name.strip}: #{value.to_s.strip}\r\n"
+      end
+
+      def to_s
+        @lines.join('')
+      end
+    end
+
+  end
+end
+

lib/websocket/driver/hybi.rb

@@ -213,7 +213,7 @@ module WebSocket
           end
         end
 
-        (headers + ['','']).join("\r\n")
+        (headers + [@headers.to_s, '']).join("\r\n")
       end
 
       def shutdown(code, reason)

lib/websocket/driver/server.rb

@@ -24,13 +24,13 @@ module WebSocket
         url
       end
 
-      %w[start state frame text binary ping close].each do |method|
+      %w[set_header start state frame text binary ping close].each do |method|
         define_method(method) do |*args|
           if @delegate
             @delegate.__send__(method, *args)
           else
             @queue << [method, args]
-            false
+            true
           end
         end
       end

lib/websocket/http.rb

@@ -7,6 +7,10 @@ module WebSocket
     autoload :Request,  root + '/request'
     autoload :Response, root + '/response'
 
+    def self.normalize_header(name)
+      name.to_s.strip.downcase.gsub(/^http_/, '').gsub(/_/, '-')
+    end
+
   end
 end
 

lib/websocket/http/headers.rb

@@ -61,14 +61,10 @@ module WebSocket
 
       def header_line(line)
         return false unless parsed = line.scan(HEADER_LINE).first
-        @headers[normalize_header(parsed[0])] = parsed[1].strip
+        @headers[HTTP.normalize_header(parsed[0])] = parsed[1].strip
         true
       end
 
-      def normalize_header(name)
-        name.downcase.gsub(/^http_/, '').gsub(/_/, '-')
-      end
-
       def string_buffer
         @buffer.pack('C*')
       end

lib/websocket/http/request.rb

@@ -28,9 +28,9 @@ module WebSocket
 
       def complete
         super
-        @headers.each do |key, value|
-          rack_name = key.upcase.gsub(/-/, '_')
-          rack_name = "HTTP_#{rack_name}" unless RESERVED_HEADERS.include?(key)
+        @headers.each do |name, value|
+          rack_name = name.upcase.gsub(/-/, '_')
+          rack_name = "HTTP_#{rack_name}" unless RESERVED_HEADERS.include?(name)
           @env[rack_name] = value
         end
         if host = @env['HTTP_HOST']

lib/websocket/http/response.rb

@@ -9,7 +9,7 @@ module WebSocket
       attr_reader :code
 
       def [](name)
-        @headers[normalize_header(name)]
+        @headers[HTTP.normalize_header(name)]
       end
 
       def body

spec/websocket/driver/client_spec.rb

@@ -84,6 +84,25 @@ describe WebSocket::Driver::Client do
         end
       end
 
+      describe "with custom headers" do
+        before do
+          driver.set_header "User-Agent", "Chrome"
+        end
+
+        it "writes the handshake with custom headers" do
+          socket.should_receive(:write).with(
+              "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" +
+              "User-Agent: Chrome\r\n" +
+              "\r\n")
+          driver.start
+        end
+      end
+
       it "changes the state to :connecting" do
         driver.start
         driver.state.should == :connecting

spec/websocket/driver/hybi_spec.rb

@@ -85,6 +85,23 @@ describe WebSocket::Driver::Hybi do
         end
       end
 
+      describe "with custom headers" do
+        before do
+          driver.set_header "Authorization", "Bearer WAT"
+        end
+
+        it "writes the handshake with custom headers" do
+          socket.should_receive(:write).with(
+              "HTTP/1.1 101 Switching Protocols\r\n" +
+              "Upgrade: websocket\r\n" +
+              "Connection: Upgrade\r\n" +
+              "Sec-WebSocket-Accept: JdiiuafpBKRqD7eol0y4vJDTsTs=\r\n" +
+              "Authorization: Bearer WAT\r\n" +
+              "\r\n")
+          driver.start
+        end
+      end
+
       it "triggers the onopen event" do
         driver.start
         @open.should == true