WIP: Fix support for built-in OAuth hooks

Since the buitt-in OAuth hooks in libpq can return timerfd and not jsut
a socket when you ask for the current file descriptor we are waiting on
we need to make sure to use the right Ruby class to wrap the file
descriptor, if it is not a valid socket we should use IO.

Author: jeltz

.gitignore

@@ -15,6 +15,9 @@
 /lib/2.?/
 /lib/3.?/
 /pkg/
+/spec/oauth/*.bc
+/spec/oauth/*.o
+/spec/oauth/*.so
 /tmp/
 /tmp_test_*/
 /vendor/

Gemfile

@@ -15,6 +15,7 @@ group :test do
   gem "rake-compiler", "~> 1.0"
   gem "rake-compiler-dock", "~> 1.11.0" #, git: "https://github.com/rake-compiler/rake-compiler-dock"
   gem "rspec", "~> 3.5"
+  gem "webrick", "~> 1.8"
   # "bigdecimal" is a gem on ruby-3.4+ and it's optional for ruby-pg.
   # Specs should succeed without it, but 4 examples are then excluded.
   # With bigdecimal commented out here, corresponding tests are omitted on ruby-3.4+ but are executed on ruby < 3.4.

ext/pg_connection.c

@@ -965,6 +965,19 @@ pgconn_socket(VALUE self)
 	return INT2NUM(sd);
 }
 
+#ifdef _WIN32
+#define is_socket(fd) rb_w32_is_socket(fd)
+#else
+static int
+is_socket(int fd)
+{
+    struct stat sbuf;
+
+    if (fstat(fd, &sbuf) < 0)
+        rb_sys_fail("fstat(2)");
+    return S_ISSOCK(sbuf.st_mode);
+}
+#endif
 
 VALUE
 pg_wrap_socket_io(int sd, VALUE self, VALUE *p_socket_io, int *p_ruby_sd)
@@ -983,7 +996,7 @@ pg_wrap_socket_io(int sd, VALUE self, VALUE *p_socket_io, int *p_ruby_sd)
 		*p_ruby_sd = ruby_sd = sd;
 	#endif
 
-	cSocket = rb_const_get(rb_cObject, rb_intern("BasicSocket"));
+	cSocket = rb_const_get(rb_cObject, rb_intern(is_socket(ruby_sd) ? "BasicSocket" : "IO"));
 	socket_io = rb_funcall( cSocket, rb_intern("for_fd"), 1, INT2NUM(ruby_sd));
 
 	/* Disable autoclose feature */

spec/helpers.rb

@@ -194,6 +194,7 @@ class PostgresServer
 		attr_reader :port
 		attr_reader :conninfo
 		attr_reader :unix_socket
+		attr_reader :version
 
 		### Set up a PostgreSQL database instance for testing.
 		def initialize(name, port: 23456, postgresql_conf: '')
@@ -205,6 +206,7 @@ def initialize(name, port: 23456, postgresql_conf: '')
 			@pgdata = @test_dir + 'data'
 			@logfile = @test_dir + 'setup.log'
 			@pg_bindir = pg_bindir
+			@version = pg_version
 			@unix_socket = @test_dir.to_s
 			@conninfo = "host=localhost port=#{@port} dbname=test sslrootcert=#{@pgdata + 'ruby-pg-ca-cert'} sslcert=#{@pgdata + 'ruby-pg-client-cert'} sslkey=#{@pgdata + 'ruby-pg-client-key'}"
 
@@ -267,8 +269,13 @@ def setup_cluster(postgresql_conf)
 					ssl_cert_file = 'ruby-pg-server-cert'
 					ssl_key_file = 'ruby-pg-server-key'
 					fsync = off
-					#{postgresql_conf}
 				EOT
+				if @version >= 18
+					fd.puts <<~EOT
+						oauth_validator_libraries = '#{TEST_DIRECTORY}/spec/oauth/dummy_validator'
+					EOT
+				end
+				fd.puts postgresql_conf
 			end
 
 			# Enable MD5 authentication in hba config
@@ -278,6 +285,12 @@ def setup_cluster(postgresql_conf)
 					# TYPE  DATABASE     USER              ADDRESS             METHOD
 					host    all          testusermd5       ::1/128             md5
 				EOT
+				if @version >= 18
+					fd.puts <<~EOT
+						host    all          testuseroauth     127.0.0.1/32        oauth scope=test issuer="http://localhost:#{@port + 3}"
+						host    all          testuseroauth     ::1/32              oauth scope=test issuer="http://localhost:#{@port + 3}"
+					EOT
+				end
 				fd.puts hba_content
 			end
 
@@ -340,6 +353,10 @@ def pg_bindir
 		rescue
 			nil
 		end
+
+		def pg_version
+			`#{pg_bin_path("pg_ctl")} --version`[/pg_ctl \(PostgreSQL\) (\d+)/, 1]&.to_i
+		end
 	end
 
 	class CertGenerator

spec/oauth/Makefile

@@ -0,0 +1,8 @@
+MODULES = dummy_validator
+PGFILEDESC = "dummy_validator - dummy OAuth validator"
+
+OBJS = $(WIN32RES)
+
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)

spec/oauth/dummy_validator.c

@@ -0,0 +1,29 @@
+#include "postgres.h"
+#include "fmgr.h"
+#include "libpq/oauth.h"
+
+PG_MODULE_MAGIC;
+
+static bool
+validate_token(const ValidatorModuleState *state,
+			   const char *token, const char *role,
+			   ValidatorModuleResult *res)
+{
+	if (strcmp(token, "yes") == 0)
+	{
+		res->authorized = true;
+		res->authn_id = pstrdup(role);
+	}
+	return true;
+}
+
+static const OAuthValidatorCallbacks validator_callbacks = {
+	PG_OAUTH_VALIDATOR_MAGIC,
+	.validate_cb = validate_token
+};
+
+const OAuthValidatorCallbacks *
+_PG_oauth_validator_module_init(void)
+{
+	return &validator_callbacks;
+}

spec/pg/connection_spec.rb

@@ -3010,4 +3010,63 @@ def wait_check_socket(conn)
 				.to raise_error(TypeError)
 		end
 	end
+
+	describe "OAuth support", :postgresql_18 do
+		before :all do
+			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
+
+		before :each do
+			@old_env, ENV["PGOAUTHDEBUG"] = ENV["PGOAUTHDEBUG"], "UNSAFE"
+		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
+
+		it "should work with no hook" do
+			oauth_server = start_fake_oauth(@port + 3)
+
+			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
+
+		after :each do
+			ENV["PGOAUTHDEBUG"] = @old_env
+		end
+	end
 end