v18-0009-WIP-Python-OAuth-provider-implementation.patch

From 563046557804b1d397d74dfc6bd83c188fc30907 Mon Sep 17 00:00:00 2001
From: Jacob Champion <jacob.champion@enterprisedb.com>
Date: Mon, 26 Feb 2024 16:24:32 -0800
Subject: [PATCH v18 9/9] WIP: Python OAuth provider implementation

---
 src/test/modules/oauth_validator/Makefile     |   2 +
 src/test/modules/oauth_validator/meson.build  |   3 +
 .../modules/oauth_validator/t/001_server.pl   |  12 +-
 .../modules/oauth_validator/t/oauth_server.py |  91 +++++++++++
 src/test/perl/PostgreSQL/Test/OAuthServer.pm  | 141 +++---------------
 5 files changed, 124 insertions(+), 125 deletions(-)
 create mode 100755 src/test/modules/oauth_validator/t/oauth_server.py

diff --git a/src/test/modules/oauth_validator/Makefile b/src/test/modules/oauth_validator/Makefile
index 1f874cd7f2..e93e01455a 100644
--- a/src/test/modules/oauth_validator/Makefile
+++ b/src/test/modules/oauth_validator/Makefile
@@ -1,3 +1,5 @@
+export PYTHON
+
 MODULES = validator
 PGFILEDESC = "validator - test OAuth validator module"
 
diff --git a/src/test/modules/oauth_validator/meson.build b/src/test/modules/oauth_validator/meson.build
index d9c1d1d577..3feba6f826 100644
--- a/src/test/modules/oauth_validator/meson.build
+++ b/src/test/modules/oauth_validator/meson.build
@@ -29,5 +29,8 @@ tests += {
     'tests': [
       't/001_server.pl',
     ],
+    'env': {
+      'PYTHON': python.path(),
+    },
   },
 }
diff --git a/src/test/modules/oauth_validator/t/001_server.pl b/src/test/modules/oauth_validator/t/001_server.pl
index 49e04b0afe..bbfa69e442 100644
--- a/src/test/modules/oauth_validator/t/001_server.pl
+++ b/src/test/modules/oauth_validator/t/001_server.pl
@@ -34,20 +34,16 @@ $node->append_conf('postgresql.conf', "shared_preload_libraries = 'validator'\n"
 $node->append_conf('postgresql.conf', "oauth_validator_library = 'validator'\n");
 $node->start;
 
-reset_pg_hba($node, 'all', 'all', 'oauth issuer="127.0.0.1:18080" scope="openid postgres"');
-
-my $webserver = PostgreSQL::Test::OAuthServer->new(18080);
+my $webserver = PostgreSQL::Test::OAuthServer->new();
+$webserver->run();
 
 my $port = $webserver->port();
-
-is($port, 18080, "Port is 18080");
-
-$webserver->setup();
-$webserver->run();
+reset_pg_hba($node, 'all', 'all', 'oauth issuer="127.0.0.1:' . $port . '" scope="openid postgres"');
 
 $node->connect_ok("dbname=postgres oauth_client_id=f02c6361-0635", "connect",
 				  expected_stderr => qr@Visit https://example\.com/ and enter the code: postgresuser@);
 
+$webserver->stop();
 $node->stop;
 
 done_testing();
diff --git a/src/test/modules/oauth_validator/t/oauth_server.py b/src/test/modules/oauth_validator/t/oauth_server.py
new file mode 100755
index 0000000000..7fa0b05a18
--- /dev/null
+++ b/src/test/modules/oauth_validator/t/oauth_server.py
@@ -0,0 +1,91 @@
+#! /usr/bin/env python3
+
+import http.server
+import json
+import os
+import sys
+
+
+class OAuthHandler(http.server.BaseHTTPRequestHandler):
+    JsonObject = dict[str, object]  # TypeAlias is not available until 3.10
+
+    def do_GET(self):
+        if self.path == "/.well-known/openid-configuration":
+            resp = self.config()
+        else:
+            self.send_error(404, "Not Found")
+            return
+
+        self._send_json(resp)
+
+    def do_POST(self):
+        if self.path == "/authorize":
+            resp = self.authorization()
+        elif self.path == "/token":
+            resp = self.token()
+        else:
+            self.send_error(404, "Not Found")
+            return
+
+        self._send_json(resp)
+
+    def _send_json(self, js: JsonObject) -> None:
+        """
+        Sends the provided JSON dict as an application/json response.
+        """
+
+        resp = json.dumps(js).encode("ascii")
+
+        self.send_response(200, "OK")
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(resp)))
+        self.end_headers()
+
+        self.wfile.write(resp)
+
+    def config(self) -> JsonObject:
+        port = self.server.socket.getsockname()[1]
+
+        return {
+            "issuer": f"http://localhost:{port}",
+            "token_endpoint": f"http://localhost:{port}/token",
+            "device_authorization_endpoint": f"http://localhost:{port}/authorize",
+            "response_types_supported": ["token"],
+            "subject_types_supported": ["public"],
+            "id_token_signing_alg_values_supported": ["RS256"],
+            "grant_types_supported": ["urn:ietf:params:oauth:grant-type:device_code"],
+        }
+
+    def authorization(self) -> JsonObject:
+        return {
+            "device_code": "postgres",
+            "user_code": "postgresuser",
+            "interval": 0,
+            "verification_uri": "https://example.com/",
+            "expires-in": 5,
+        }
+
+    def token(self) -> JsonObject:
+        return {
+            "access_token": "9243959234",
+            "token_type": "bearer",
+        }
+
+
+def main():
+    s = http.server.HTTPServer(("127.0.0.1", 0), OAuthHandler)
+
+    # Give the parent the port number to contact (this is also the signal that
+    # we're ready to receive requests).
+    port = s.socket.getsockname()[1]
+    print(port)
+
+    stdout = sys.stdout.fileno()
+    sys.stdout.close()
+    os.close(stdout)
+
+    s.serve_forever()  # we expect our parent to send a termination signal
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/test/perl/PostgreSQL/Test/OAuthServer.pm b/src/test/perl/PostgreSQL/Test/OAuthServer.pm
index 3ac90c3d0f..d96733f531 100644
--- a/src/test/perl/PostgreSQL/Test/OAuthServer.pm
+++ b/src/test/perl/PostgreSQL/Test/OAuthServer.pm
@@ -5,6 +5,7 @@ package PostgreSQL::Test::OAuthServer;
 use warnings;
 use strict;
 use threads;
+use Scalar::Util;
 use Socket;
 use IO::Select;
 
@@ -13,27 +14,13 @@ local *server_socket;
 sub new
 {
 	my $class = shift;
-	my $port = shift;
 
 	my $self = {};
 	bless($self, $class);
 
-	$self->{'port'} = $port;
-
 	return $self;
 }
 
-sub setup
-{
-	my $self = shift;
-	my $tcp = getprotobyname('tcp');
-
-	socket($self->{'socket'}, PF_INET, SOCK_STREAM, $tcp)
-		or die "no socket";
-	setsockopt($self->{'socket'}, SOL_SOCKET, SO_REUSEADDR, pack("l", 1));
-	bind($self->{'socket'}, sockaddr_in($self->{'port'}, INADDR_ANY));
-}
-
 sub port
 {
 	my $self = shift;
@@ -44,115 +31,35 @@ sub port
 sub run
 {
 	my $self = shift;
+	my $port;
 
-	my $server_thread = threads->create(\&_listen, $self);
-	$server_thread->detach();
+	my $pid = open(my $read_fh, "-|", $ENV{PYTHON}, "t/oauth_server.py")
+		// die "failed to start OAuth server: $!";
+
+	read($read_fh, $port, 7) // die "failed to read port number: $!";
+	chomp $port;
+	die "server did not advertise a valid port"
+		unless Scalar::Util::looks_like_number($port);
+
+	$self->{'pid'} = $pid;
+	$self->{'port'} = $port;
+	$self->{'child'} = $read_fh;
+
+	print("# OAuth provider (PID $pid) is listening on port $port\n");
 }
 
-sub _listen
+sub stop
 {
 	my $self = shift;
 
-	listen($self->{'socket'}, SOMAXCONN) or die "fail to listen: $!";
-
-	while (1)
-	{
-		my $fh;
-		my %request;
-		my $remote = accept($fh, $self->{'socket'});
-		binmode $fh;
-
-		my ($method, $object, $prot) = split(/ /, <$fh>);
-		$request{'method'} = $method;
-		$request{'object'} = $object;
-		chomp($request{'object'});
-
-		local $/ = Socket::CRLF;
-		my $c = 0;
-		while(<$fh>)
-		{
-			chomp;
-			# Headers
-			if (/:/)
-			{
-				my ($field, $value) = split(/:/, $_, 2);
-				$value =~ s/^\s+//;
-				$request{'headers'}{lc $field} = $value;
-			}
-			# POST data
-			elsif (/^$/)
-			{
-				read($fh, $request{'content'}, $request{'headers'}{'content-length'})
-					if defined $request{'headers'}{'content-length'};
-				last;
-			}
-		}
-
-		# Debug printing
-		# print ": read ".$request{'method'} . ";" . $request{'object'}.";\n";
-		# foreach my $h (keys(%{$request{'headers'}}))
-		#{
-		#	printf ": headers: " . $request{'headers'}{$h} . "\n";
-		#}
-		#printf ": POST: " . $request{'content'} . "\n" if defined $request{'content'};
-
-		if ($request{'object'} eq '/.well-known/openid-configuration')
-		{
-			print $fh "HTTP/1.0 200 OK\r\nServer: Postgres Regress\r\n";
-			print $fh "Content-Type: application/json\r\n";
-			print $fh "\r\n";
-			print $fh <<EOR;
-			{
-				"issuer": "http://localhost:$self->{'port'}",
-				"token_endpoint": "http://localhost:$self->{'port'}/token",
-				"device_authorization_endpoint": "http://localhost:$self->{'port'}/authorize",
-				"response_types_supported": ["token"],
-				"subject_types_supported": ["public"],
-				"id_token_signing_alg_values_supported": ["RS256"],
-				"grant_types_supported": ["urn:ietf:params:oauth:grant-type:device_code"]
-			}
-EOR
-		}
-		elsif ($request{'object'} eq '/authorize')
-		{
-			print ": returning device_code\n";
-			print $fh "HTTP/1.0 200 OK\r\nServer: Postgres Regress\r\n";
-			print $fh "Content-Type: application/json\r\n";
-			print $fh "\r\n";
-			print $fh <<EOR;
-			{
-				"device_code": "postgres",
-				"user_code" : "postgresuser",
-				"interval" : 0,
-				"verification_uri" : "https://example.com/",
-				"expires-in": 5
-			}
-EOR
-		}
-		elsif ($request{'object'} eq '/token')
-		{
-			print ": returning token\n";
-			print $fh "HTTP/1.0 200 OK\r\nServer: Postgres Regress\r\n";
-			print $fh "Content-Type: application/json\r\n";
-			print $fh "\r\n";
-			print $fh <<EOR;
-			{
-				"access_token": "9243959234",
-				"token_type": "bearer"
-			}
-EOR
-		}
-		else
-		{
-			print ": returning default\n";
-			print $fh "HTTP/1.0 200 OK\r\nServer: Postgres Regress\r\n";
-			print $fh "Content-Type: text/html\r\n";
-			print $fh "\r\n";
-			print $fh "Ok\n";
-		}
-
-		close($fh);
-	}
+	print("# Sending SIGTERM to OAuth provider PID: $self->{'pid'}\n");
+
+	kill(15, $self->{'pid'});
+	$self->{'pid'} = undef;
+
+	# Closing the popen() handle waits for the process to exit.
+	close($self->{'child'});
+	$self->{'child'} = undef;
 }
 
 1;
-- 
2.34.1