Change the OAuth hook to be per connection

The per connection hook works only in async mode, but sync API in ruby-pg is documented as for testing only and they have several flaws already.
So I don't think there is any need to support them with the OAuth hook.
Therefore on sync API the only supported option is non-hooked OAuth.

This moves common code to helpers and hooked OAuth to async specs.

Connecting with a OAuth hook within a Ractor is currently also not possible, but might be changed in future.

Author: larskanis

lib/pg.rb

@@ -84,8 +84,8 @@ def self.version_string( include_buildnum=nil )
 
 
   ### Convenience alias for PG::Connection.new.
-  def self.connect( *args, &block )
-    Connection.new( *args, &block )
+  def self.connect( *args, **kwargs, &block )
+    Connection.new( *args, **kwargs, &block )
   end
 
   if defined?(Ractor.make_shareable)

lib/pg/connection.rb

@@ -867,8 +867,8 @@ class << self
 		#   It's still possible to do load balancing with +load_balance_hosts+ set to +random+ and to increase the number of connections a node gets, when the hostname is provided multiple times in the host string.
 		#   This is because in non-timeout cases the host is tried multiple times.
 		#
-		def new(*args)
-			conn = connect_to_hosts(*args)
+		def new(*args, **kwargs)
+			conn = connect_to_hosts(*args, **kwargs)
 
 			if block_given?
 				begin
@@ -919,8 +919,8 @@ def new(*args)
 				port: dests.map{|d| d[2] }.join(","))
 		end
 
-		private def connect_to_hosts(*args)
-			option_string = parse_connect_args(*args)
+		private def connect_to_hosts(*args, set_auth_data_hook: nil, **kwargs)
+			option_string = parse_connect_args(*args, **kwargs)
 			iopts = PG::Connection.conninfo_parse(option_string).each_with_object({}){|h, o| o[h[:keyword].to_sym] = h[:val] if h[:val] }
 			iopts = PG::Connection.conndefaults.each_with_object({}){|h, o| o[h[:keyword].to_sym] = h[:val] if h[:val] }.merge(iopts)
 
@@ -943,6 +943,12 @@ def new(*args)
 			end
 			conn = self.connect_start(iopts) or
 										raise(PG::Error, "Unable to create a new connection")
+			if set_auth_data_hook
+				@@auth_mutex.synchronize do
+					@@pgconn_map[conn.send(:pgconn_address)] = conn
+				end
+				conn.instance_variable_set(:@auth_data_hook, set_auth_data_hook)
+			end
 
 			raise PG::ConnectionBad, conn.error_message if conn.status == PG::CONNECTION_BAD
 
@@ -1090,5 +1096,23 @@ def async_api=(enable)
 		end
 	end
 
+	if PG.respond_to?(:set_auth_data_hook)
+		def call_auth_data_hook(data)
+			@auth_data_hook.call(self, data)
+		end
+
+		@@auth_mutex = Mutex.new
+		@@pgconn_map = ObjectSpace::WeakMap.new
+
+		PG.set_auth_data_hook do |conn_num, data|
+			@@auth_mutex.synchronize do
+				conn = @@pgconn_map[conn_num]
+				if conn
+					conn.call_auth_data_hook(data)
+				end
+			end
+		end
+	end
+
 	self.async_api = true
 end # class PG::Connection

spec/helpers.rb

@@ -673,6 +673,38 @@ def with_env_vars(**kwargs)
 	def set_etc_hosts(hostaddr, hostname)
 		system "sudo --non-interactive sed -i '/.* #{hostname}$/{h;s/.*/#{hostaddr} #{hostname}/};${x;/^$/{s//#{hostaddr} #{hostname}/;H};x}' /etc/hosts" or skip("unable to change /etc/hosts file")
 	end
+
+	def build_oauth_validator
+		skip "requires a PostgreSQL 18 cluster" unless $pg_server.version >= 18
+
+		system "make", "-s", "-C", (TEST_DIRECTORY + "spec/oauth").to_s
+		raise "Building OAuth validator library failed!" unless $?.success?
+
+		require 'webrick'
+
+		PG.connect(@conninfo) do |conn|
+			conn.exec("DROP USER IF EXISTS testuseroauth")
+			conn.exec("CREATE USER testuseroauth")
+		end
+	end
+
+	def start_fake_oauth(port)
+		server = WEBrick::HTTPServer.new(Port: port, Logger: WEBrick::Log.new(nil, WEBrick::BasicLog::WARN))
+		server.mount_proc("/.well-known/openid-configuration") do |req, res|
+			res["Content-Type"] = "application/json"
+			res.body = %!{"issuer":"http://localhost:#{port}","token_endpoint":"http://localhost:#{port}/token","device_authorization_endpoint":"http://localhost:#{@port + 3}/devauth"}!
+		end
+		server.mount_proc("/devauth") do |req, res|
+			res["Content-Type"] = "application/json"
+			res.body = %!{"device_code":"42","user_code":"666","verification_uri":"http://localhost:#{port}/verify","expires_in":60}!
+		end
+		server.mount_proc("/token") do |req, res|
+			res["Content-Type"] = "application/json"
+			res.body = %!{"access_token":"yes","token_type":""}!
+		end
+		Thread.new { server.start }
+		server
+	end
 end
 
 RSpec.configure do |config|

spec/pg/connection_async_spec.rb

@@ -156,4 +156,163 @@ def interrupt_thread(exc=nil)
 		expect( conn.hostaddr ).to eq( "::1" )
 		expect( conn.port ).to eq( @port )
 	end
+
+	describe "option set_auth_data_hook", :postgresql_18  do
+		before :all do
+			build_oauth_validator
+		end
+
+		before :each do
+			@old_env, ENV["PGOAUTHDEBUG"] = ENV["PGOAUTHDEBUG"], "UNSAFE"
+		end
+
+		it "should call prompt oauth device hook" do
+			oauth_server = start_fake_oauth(@port + 3)
+
+			verification_uri, user_code, verification_uri_complete, expires_in = nil, nil, nil, nil
+			conn1, conn2 = nil, nil
+
+			hook = proc do |conn, data|
+				case data
+				when PG::PromptOAuthDevice
+					conn1 = conn
+					verification_uri = data.verification_uri
+					user_code = data.user_code
+					verification_uri_complete = data.verification_uri_complete
+					expires_in = data.expires_in
+					true
+				end
+			end
+
+			begin
+				PG.connect("host=localhost port=#{@port} dbname=test user=testuseroauth oauth_issuer=http://localhost:#{@port + 3} oauth_client_id=foo", set_auth_data_hook: hook) do |conn|
+					conn.exec("SELECT 1")
+					conn2 = conn
+				end
+			rescue PG::ConnectionBad => e
+				if e.message =~ /no OAuth flows are available/
+					skip "requires libpq-oauth to be installed"
+				end
+				raise
+			ensure
+				oauth_server.shutdown
+			end
+
+			expect(conn1).to eq(conn2)
+			expect(verification_uri).to eq("http://localhost:#{@port + 3}/verify")
+			expect(user_code).to eq("666")
+			expect(verification_uri_complete).to eq(nil)
+			expect(expires_in).to eq(60)
+		end
+
+		it "should call oauth bearer request hook" do
+			openid_configuration, scope = nil, nil
+			conn1, conn2 = nil, nil
+
+			hook = proc do |conn, data|
+				case data
+				when PG::OAuthBearerRequest
+					conn1 = conn
+					openid_configuration = data.openid_configuration
+					scope = data.scope
+					data.token = "yes"
+					true
+				end
+			end
+
+			PG.connect(host: "localhost", port: @port, dbname: "test", user: "testuseroauth", oauth_issuer: "http://localhost:#{@port + 3}", oauth_client_id: "foo", set_auth_data_hook: hook) do |conn|
+				conn.exec("SELECT 1")
+				conn2 = conn
+			end
+
+			expect(conn1).to eq(conn2)
+			expect(openid_configuration).to eq("http://localhost:#{@port + 3}/.well-known/openid-configuration")
+			expect(scope).to eq("test")
+		end
+
+		it "shouldn't garbage collect PG::Connection in use" do
+			conn1 = nil
+			hook = proc do |conn, data|
+				case data
+				when PG::OAuthBearerRequest
+					data.token = "yes"
+					conn1 = conn
+					true
+				end
+			end
+
+			GC.stress = true
+			begin
+				conn = PG.connect(host: "localhost", port: @port, dbname: "test", user: "testuseroauth", oauth_issuer: "http://localhost:#{@port + 3}", oauth_client_id: "foo", set_auth_data_hook: hook)
+			ensure
+				GC.stress = false
+			end
+			conn.exec("SELECT 1")
+
+			expect(conn1).to eq(conn)
+		end
+
+		it "should garbage collect PG::Connection after use" do
+			hook = proc do |conn, data|
+				case data
+				when PG::OAuthBearerRequest
+					conn1 = conn
+					openid_configuration = data.openid_configuration
+					scope = data.scope
+					data.token = "yes"
+					true
+				end
+			end
+
+			10.times do
+				conn = PG.connect(host: "localhost", port: @port, dbname: "test", user: "testuseroauth", oauth_issuer: "http://localhost:#{@port + 3}", oauth_client_id: "foo", set_auth_data_hook: hook)
+				conn.exec("SELECT 1")
+			end
+
+			before = PG::Connection.class_variable_get(:@@pgconn_map).keys.size
+			GC.start
+			after = PG::Connection.class_variable_get(:@@pgconn_map).keys.size
+
+			expect(before - after).to be_between(2, 20)
+		end
+
+		# TODO: Is resetting the global hook still useful, when the hook is per connection?
+		# it "should reset the hook when called without block" do
+		# 	oauth_server = start_fake_oauth(@port + 3)
+		#
+		# 	PG.set_auth_data_hook do |conn_num, data|
+		# 		raise "broken hook"
+		# 	end
+		#
+		# 	expect do
+		# 		PG.connect("host=localhost port=#{@port} dbname=test user=testuseroauth oauth_issuer=http://localhost:#{@port + 3} oauth_client_id=foo") {}
+		# 	end.to raise_error("broken hook")
+		#
+		# 	PG.set_auth_data_hook
+		#
+		# 	begin
+		# 		PG.connect("host=localhost port=#{@port} dbname=test user=testuseroauth oauth_issuer=http://localhost:#{@port + 3} oauth_client_id=foo") do |conn|
+		# 			conn.exec("SELECT 1")
+		# 		end
+		# 	rescue PG::ConnectionBad => e
+		# 		if e.message =~ /no OAuth flows are available/
+		# 			skip "requires libpq-oauth to be installed"
+		# 		end
+		# 		raise
+		# 	ensure
+		# 		oauth_server.shutdown
+		# 	end
+		# end
+
+		# around :example do |ex|
+		# 	GC.stress = true
+		# 	ex.run
+		# 	GC.stress = false
+		# end
+
+		after :each do
+			# PG.set_auth_data_hook
+			ENV["PGOAUTHDEBUG"] = @old_env
+		end
+	end
 end