From 2daa01e26bfad39e61dd5e81ba98cfa5668c4338 Mon Sep 17 00:00:00 2001
From: Sam Ruby <rubys@intertwingly.net>
Date: Fri, 9 Oct 2026 11:52:35 -0400
Subject: [PATCH] Open pooled push connections as a Net::HTTP subclass that
 includes Stages

WebPush::Connections mixed Stages into each new connection with
`http.extend Stages` just before starting it. The connection is built in
one place, WebPush::PersistentRequest, so it can be built from a class
that already includes Stages: the same method lookup, without giving
every pooled connection its own singleton class. Connections#request
now also refuses a connection that isn't one, so a caller that builds a
plain Net::HTTP gets the ArgumentError rather than a connection whose
failures can't be told apart.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 config/initializers/web_push.rb               | 2 +-
 lib/web_push/connections.rb                   | 9 +++++++--
 test/test_helpers/push_service_test_helper.rb | 2 +-
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/config/initializers/web_push.rb b/config/initializers/web_push.rb
index 006b0a0..f55f5b3 100644
--- a/config/initializers/web_push.rb
+++ b/config/initializers/web_push.rb
@@ -29,7 +29,7 @@ module WebPush::PersistentRequest
     # itself and re-resolve the endpoint host, so http.ipaddr would no longer
     # pin the destination and the DNS-rebinding guarantee would be lost.
     # Push delivery to public vendor endpoints goes direct.
-    http = Net::HTTP.new(uri.host, uri.port, nil)
+    http = (@options[:connection] ? WebPush::Connections::HTTP : Net::HTTP).new(uri.host, uri.port, nil)
     http.ipaddr = endpoint_ip
     http.use_ssl = true
     http.ssl_timeout = @options[:ssl_timeout] unless @options[:ssl_timeout].nil?
diff --git a/lib/web_push/connections.rb b/lib/web_push/connections.rb
index f151701..24a13aa 100644
--- a/lib/web_push/connections.rb
+++ b/lib/web_push/connections.rb
@@ -26,6 +26,12 @@ class WebPush::Connections
       end
   end
 
+  # The Net::HTTP a pooled delivery is opened with: Stages is part of the class, so a connection carries its stage
+  # from the start instead of having the module mixed into each one as it's opened.
+  class HTTP < Net::HTTP
+    include Stages
+  end
+
   def initialize(keep_alive_timeout: 30, max_idle: 150)
     @keep_alive_timeout = keep_alive_timeout
     @max_idle = max_idle
@@ -35,7 +41,7 @@ class WebPush::Connections
   end
 
   def request(http, request)
-    unless http.ipaddr && !http.proxy? && http.use_ssl? && !http.started?
+    unless http.is_a?(HTTP) && http.ipaddr && !http.proxy? && http.use_ssl? && !http.started?
       raise ArgumentError, "Only new, direct TLS connections pinned to an address are pooled"
     end
     address = [ http.address, http.port, http.ipaddr ]
@@ -48,7 +54,6 @@ class WebPush::Connections
       end
     end
 
-    http.extend Stages
     http.keep_alive_timeout = @keep_alive_timeout
     http.start
     send_over(http, request, reused: false).tap { checkin(address, http) }
diff --git a/test/test_helpers/push_service_test_helper.rb b/test/test_helpers/push_service_test_helper.rb
index a150561..1d4e701 100644
--- a/test/test_helpers/push_service_test_helper.rb
+++ b/test/test_helpers/push_service_test_helper.rb
@@ -109,7 +109,7 @@ module PushServiceTestHelper
 
     # Built as WebPush::PersistentRequest builds it for a delivery.
     def pinned_connection(server, ip = IP)
-      Net::HTTP.new(HOST, server.port, nil).tap do |http|
+      WebPush::Connections::HTTP.new(HOST, server.port, nil).tap do |http|
         http.ipaddr = ip
         http.use_ssl = true
       end