From 716f19d79e94541bd525f228bfe184420c7c44ad Mon Sep 17 00:00:00 2001
From: James Coglan <jcoglan@gmail.com>
Date: Sat, 10 Sep 2022 15:37:55 +0100
Subject: [PATCH] Fix handling of default ports on Ruby 3.1

---
 lib/websocket/driver.rb              |  9 ++++++
 lib/websocket/driver/client.rb       |  3 +-
 lib/websocket/driver/proxy.rb        |  4 +--
 spec/websocket/driver/client_spec.rb | 48 ++++++++++++++++++++++++++++
 4 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/lib/websocket/driver.rb b/lib/websocket/driver.rb
index d568e91..7033a82 100644
--- a/lib/websocket/driver.rb
+++ b/lib/websocket/driver.rb
@@ -42,6 +42,7 @@ module WebSocket
     end
 
     MAX_LENGTH = 0x3ffffff
+    PORTS      = { 'ws' => 80, 'wss' => 443 }
     STATES     = [:connecting, :open, :closing, :closed]
 
     BINARY  = 'ASCII-8BIT'
@@ -178,6 +179,14 @@ module WebSocket
       string.valid_encoding? ? string : nil
     end
 
+    def self.host_header(uri)
+      host = uri.host
+      if uri.port and uri.port != PORTS[uri.scheme]
+        host += ":#{uri.port}"
+      end
+      host
+    end
+
     def self.validate_options(options, valid_keys)
       options.keys.each do |key|
         unless valid_keys.include?(key)
diff --git a/lib/websocket/driver/client.rb b/lib/websocket/driver/client.rb
index 0a45534..68e10eb 100644
--- a/lib/websocket/driver/client.rb
+++ b/lib/websocket/driver/client.rb
@@ -23,11 +23,10 @@ module WebSocket
           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['Host']                  = Driver.host_header(uri)
         @headers['Upgrade']               = 'websocket'
         @headers['Connection']            = 'Upgrade'
         @headers['Sec-WebSocket-Key']     = @key
diff --git a/lib/websocket/driver/proxy.rb b/lib/websocket/driver/proxy.rb
index e4891e3..69f59a5 100644
--- a/lib/websocket/driver/proxy.rb
+++ b/lib/websocket/driver/proxy.rb
@@ -4,8 +4,6 @@ module WebSocket
     class Proxy
       include EventEmitter
 
-      PORTS = {'ws' => 80, 'wss' => 443}
-
       attr_reader :status, :headers
 
       def initialize(client, origin, options)
@@ -20,7 +18,7 @@ module WebSocket
         @state   = 0
 
         @headers = Headers.new
-        @headers['Host'] = @origin.host + (@origin.port ? ":#{@origin.port}" : '')
+        @headers['Host'] = Driver.host_header(@origin)
         @headers['Connection'] = 'keep-alive'
         @headers['Proxy-Connection'] = 'keep-alive'
 
diff --git a/spec/websocket/driver/client_spec.rb b/spec/websocket/driver/client_spec.rb
index 75b7f4f..abdde2c 100644
--- a/spec/websocket/driver/client_spec.rb
+++ b/spec/websocket/driver/client_spec.rb
@@ -121,6 +121,54 @@ describe WebSocket::Driver::Client do
         end
       end
 
+      describe "with an explicit port" do
+        let(:url) { "ws://www.example.com:3000/socket" }
+
+        it "includes the port in the Host header" do
+          expect(socket).to receive(:write).with(
+              "GET /socket HTTP/1.1\r\n" +
+              "Host: www.example.com:3000\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.start
+        end
+      end
+
+      describe "with a wss: URL" do
+        let(:url) { "wss://www.example.com/socket" }
+
+        it "does not include the port in the Host header" do
+          expect(socket).to 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" +
+              "\r\n")
+          driver.start
+        end
+      end
+
+      describe "with a wss: URL and explicit port" do
+        let(:url) { "wss://www.example.com:3000/socket" }
+
+        it "includes the port in the Host header" do
+          expect(socket).to receive(:write).with(
+              "GET /socket HTTP/1.1\r\n" +
+              "Host: www.example.com:3000\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.start
+        end
+      end
+
       describe "with custom headers" do
         before do
           driver.set_header "User-Agent", "Chrome"
-- 
2.27.0