Fix handling of default ports on Ruby 3.1

lib/websocket/driver.rb

@@ -42,6 +42,7 @@ module WebSocket
     end
 
     MAX_LENGTH = 0x3ffffff
+    PORTS      = { 'ws' => 80, 'wss' => 443 }
     STATES     = [:connecting, :open, :closing, :closed]
 
     ConnectEvent = Struct.new(nil)
@@ -209,6 +210,14 @@ module WebSocket
       data.force_encoding(encoding)
     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)

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

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'
 

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"