From 5df85c39c304a26851efbcc144015537c1d04123 Mon Sep 17 00:00:00 2001
From: Lucas DRAESCHER <lucas.draescher@clever-cloud.com>
Date: Thu, 18 Jul 2024 10:21:35 +0200
Subject: [PATCH] WIP: add support for ClientAuthentication_hook

---
 pgrx-pg-sys/include/pg13.h          |  1 +
 pgrx-tests/src/tests/hooks_tests.rs | 13 +++++++++++-
 pgrx/src/hooks.rs                   | 31 +++++++++++++++++++++++++++++
 3 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/pgrx-pg-sys/include/pg13.h b/pgrx-pg-sys/include/pg13.h
index 649ba9a145..082be15af6 100644
--- a/pgrx-pg-sys/include/pg13.h
+++ b/pgrx-pg-sys/include/pg13.h
@@ -15,6 +15,7 @@
 #include "miscadmin.h"
 #include "pgstat.h"
 
+#include "libpq/auth.h"
 #include "access/amapi.h"
 #include "access/genam.h"
 #include "access/generic_xlog.h"
diff --git a/pgrx-tests/src/tests/hooks_tests.rs b/pgrx-tests/src/tests/hooks_tests.rs
index bc46d21259..fea12f65b4 100644
--- a/pgrx-tests/src/tests/hooks_tests.rs
+++ b/pgrx-tests/src/tests/hooks_tests.rs
@@ -25,6 +25,17 @@ mod tests {
             events: u32,
         }
         impl PgHooks for TestHook {
+
+            fn client_authentication(
+                &mut self,
+                port: PgBox<Port>,
+                status: i32,
+                prev_hook: fn(port: PgBox<Port>, status: i32) -> HookResult<()>,
+            ) -> HookResult<()> {
+                self.events += 1;
+                prev_hook(port, status)
+            }
+
             /// Hook before the logs are being processed by PostgreSQL itself
             fn emit_log(
                 &mut self,
@@ -128,7 +139,7 @@ mod tests {
         // To trigger the emit_log hook, we need something to log.
         // We therefore ensure the select statement will be logged.
         Spi::run("SET local log_statement to 'all'; SELECT 1").expect("SPI failed");
-        assert_eq!(8, HOOK.events);
+        assert_eq!(9, HOOK.events);
 
         // TODO:  it'd be nice to also test that .commit() and .abort() also get called
         //    but I don't see how to do that since we're running *inside* a transaction here
diff --git a/pgrx/src/hooks.rs b/pgrx/src/hooks.rs
index 965e418ce4..07bd428e3c 100644
--- a/pgrx/src/hooks.rs
+++ b/pgrx/src/hooks.rs
@@ -43,6 +43,15 @@ impl<T> Deref for HookResult<T> {
 }
 
 pub trait PgHooks {
+    fn client_authentication(
+        &mut self,
+        port: PgBox<pg_sys::Port>,
+        status: i32,
+        prev_hook: fn(port: PgBox<pg_sys::Port>, status: i32) -> HookResult<()>,
+    ) -> HookResult<()> {
+        prev_hook(port, status)
+    }
+
     /// Hook before the logs are being processed by PostgreSQL itself
     fn emit_log(
         &mut self,
@@ -186,6 +195,7 @@ pub trait PgHooks {
 
 struct Hooks {
     current_hook: Box<&'static mut (dyn PgHooks)>,
+    client_authentication: pg_sys::ClientAuthentication_hook_type,
     prev_emit_log_hook: pg_sys::emit_log_hook_type,
     prev_executor_start_hook: pg_sys::ExecutorStart_hook_type,
     prev_executor_run_hook: pg_sys::ExecutorRun_hook_type,
@@ -238,6 +248,7 @@ pub unsafe fn register_hook(hook: &'static mut (dyn PgHooks)) {
         prev_post_parse_analyze_hook: pg_sys::post_parse_analyze_hook
             .replace(pgrx_post_parse_analyze),
         prev_emit_log_hook: pg_sys::emit_log_hook.replace(pgrx_emit_log),
+        client_authentication: pg_sys::ClientAuthentication_hook.replace(pgrx_client_auth),
     });
 
     #[pg_guard]
@@ -483,6 +494,26 @@ unsafe extern "C" fn pgrx_planner(
     pgrx_planner_impl(parse, std::ptr::null(), cursor_options, bound_params)
 }
 
+#[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))]
+#[pg_guard]
+unsafe extern "C" fn pgrx_client_auth(
+    port: *mut pg_sys::Port,
+    status: i32
+) {
+    fn prev(port: PgBox<pg_sys::Port>, status: i32) -> HookResult<()> {
+        HookResult::new(unsafe {
+            match HOOKS.as_mut().unwrap().client_authentication.as_ref() {
+                None => (),
+                Some(f) => (f)(port.as_ptr(), status),
+            }
+        })
+    }
+
+    let hook = &mut HOOKS.as_mut().unwrap().current_hook;
+    hook.client_authentication(PgBox::from_pg(port), status, prev).inner
+}
+
+
 #[cfg(any(feature = "pg13", feature = "pg14", feature = "pg15", feature = "pg16"))]
 #[pg_guard]
 unsafe extern "C" fn pgrx_planner(