From 36321727909adbc5f1325a962b9d00c4c44e789a Mon Sep 17 00:00:00 2001
From: Adam Trilling <adamtrilling@gmail.com>
Date: Sat, 1 Apr 2017 15:06:52 -0400
Subject: [PATCH] Add headers parameter for HTTP::WebSocket constructors

---
 spec/std/http/web_socket_spec.cr |  1 +
 src/http/web_socket.cr           | 10 ++++++----
 src/http/web_socket/protocol.cr  |  7 +++----
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/spec/std/http/web_socket_spec.cr b/spec/std/http/web_socket_spec.cr
index 6234b444795f..39603b7ef6ad 100644
--- a/spec/std/http/web_socket_spec.cr
+++ b/spec/std/http/web_socket_spec.cr
@@ -296,4 +296,5 @@ describe HTTP::WebSocket do
   typeof(HTTP::WebSocket.new(URI.parse("ws://localhost")))
   typeof(HTTP::WebSocket.new("localhost", "/"))
   typeof(HTTP::WebSocket.new("ws://localhost"))
+  typeof(HTTP::WebSocket.new(URI.parse("ws://localhost"), headers: HTTP::Headers{"X-TEST_HEADER" => "some-text"}))
 end
diff --git a/src/http/web_socket.cr b/src/http/web_socket.cr
index 4c2be73598d0..0e59b5ca28b7 100644
--- a/src/http/web_socket.cr
+++ b/src/http/web_socket.cr
@@ -23,9 +23,11 @@ class HTTP::WebSocket
   # HTTP::WebSocket.new(URI.parse("ws://websocket.example.com/chat"))        # Creates a new WebSocket to `websocket.example.com`
   # HTTP::WebSocket.new(URI.parse("wss://websocket.example.com/chat"))       # Creates a new WebSocket with TLS to `websocket.example.com`
   # HTTP::WebSocket.new(URI.parse("http://websocket.example.com:8080/chat")) # Creates a new WebSocket to `websocket.example.com` on port `8080`
+  # HTTP::WebSocket.new(URI.parse("ws://websocket.example.com/chat"),        # Creates a new WebSocket to `websocket.example.com` with an Authorization header
+  #   HTTP::Headers{"Authorization" => "Bearer authtoken"})
   # ```
-  def self.new(uri : URI | String)
-    new(Protocol.new(uri))
+  def self.new(uri : URI | String, headers = HTTP::Headers.new)
+    new(Protocol.new(uri, headers: headers))
   end
 
   # Opens a new websocket to the target host. This will also handle the handshake
@@ -35,8 +37,8 @@ class HTTP::WebSocket
   # HTTP::WebSocket.new("websocket.example.com", "/chat")            # Creates a new WebSocket to `websocket.example.com`
   # HTTP::WebSocket.new("websocket.example.com", "/chat", tls: true) # Creates a new WebSocket with TLS to `ẁebsocket.example.com`
   # ```
-  def self.new(host : String, path : String, port = nil, tls = false)
-    new(Protocol.new(host, path, port, tls))
+  def self.new(host : String, path : String, port = nil, tls = false, headers = HTTP::Headers.new)
+    new(Protocol.new(host, path, port, tls, headers))
   end
 
   def on_ping(&@on_ping : String ->)
diff --git a/src/http/web_socket/protocol.cr b/src/http/web_socket/protocol.cr
index 581d797a6acd..90a7d54cd138 100644
--- a/src/http/web_socket/protocol.cr
+++ b/src/http/web_socket/protocol.cr
@@ -237,7 +237,7 @@ class HTTP::WebSocket::Protocol
     end
   end
 
-  def self.new(host : String, path : String, port = nil, tls = false)
+  def self.new(host : String, path : String, port = nil, tls = false, headers = HTTP::Headers.new)
     {% if flag?(:without_openssl) %}
       if tls
         raise "WebSocket TLS is disabled because `-D without_openssl` was passed at compile time"
@@ -259,7 +259,6 @@ class HTTP::WebSocket::Protocol
       end
     {% end %}
 
-    headers = HTTP::Headers.new
     headers["Host"] = "#{host}:#{port}"
     headers["Connection"] = "Upgrade"
     headers["Upgrade"] = "websocket"
@@ -277,12 +276,12 @@ class HTTP::WebSocket::Protocol
     new(socket, masked: true)
   end
 
-  def self.new(uri : URI | String)
+  def self.new(uri : URI | String, headers = HTTP::Headers.new)
     uri = URI.parse(uri) if uri.is_a?(String)
 
     if (host = uri.host) && (path = uri.path)
       tls = uri.scheme == "https" || uri.scheme == "wss"
-      return new(host, path, uri.port, tls)
+      return new(host, path, uri.port, tls, headers)
     end
 
     raise ArgumentError.new("No host or path specified which are required.")