From e673604a720e0748ed7a1dc9707c65ad46370f16 Mon Sep 17 00:00:00 2001
From: Mike Palmiotto <mike.palmiotto@crunchydata.com>
Date: Tue, 20 Mar 2018 22:26:37 -0400
Subject: [PATCH] Add post-set_user hooks

This commit introduces two hooks which allow another extension to do something
(and return) immediately after `set_user` and `reset_user` are called.

There is a new make target, `make install-headers`, which is responsible for
exposing the hook to other extensions. The target is prepended to the
`install` target by default.

To use the post-set_user hooks in an extension:
1) Add '-I$(includedir)' to CPPFLAGS
2) #include set_user.h in whichever file implements the hook
3) Register the 'post_reset_user_hook' and 'post_set_user_hook' with a local
implementation
---
 Makefile   | 12 ++++++++++++
 set_user.c | 23 +++++++++++++++++++++++
 set_user.h | 11 +++++++++++
 3 files changed, 46 insertions(+)
 create mode 100644 set_user.h

diff --git a/Makefile b/Makefile
index fd989ff..789f764 100644
--- a/Makefile
+++ b/Makefile
@@ -18,3 +18,15 @@ top_builddir = ../..
 include $(top_builddir)/src/Makefile.global
 include $(top_srcdir)/contrib/contrib-global.mk
 endif
+
+.PHONY: install-headers uninstall-headers
+
+install: install-headers
+
+install-headers:
+	$(INSTALL_DATA) "set_user.h" $(includedir)
+
+uninstall: uninstall-headers
+
+uninstall-headers:
+	rm "$(includedir)/set_user.h"
diff --git a/set_user.c b/set_user.c
index 8bc1f19..7070c11 100644
--- a/set_user.c
+++ b/set_user.c
@@ -31,6 +31,7 @@
 #include "postgres.h"
 
 #include "pg_config.h"
+
 #if !defined(PG_VERSION_NUM) || PG_VERSION_NUM < 90100
 /* prior to 9.1 */
 #error "This extension only builds with PostgreSQL 9.1 or later"
@@ -100,6 +101,7 @@
 #include "utils/varlena.h"
 #endif /* HAS_VARLENA_H */
 
+#include "set_user.h"
 #define WHITELIST_WILDCARD	"*"
 
 PG_MODULE_MAGIC;
@@ -109,6 +111,9 @@ static Oid save_OldUserId = InvalidOid;
 static char *reset_token = NULL;
 static ProcessUtility_hook_type prev_hook = NULL;
 
+post_reset_user_hook_type post_reset_user_hook = NULL;
+post_set_user_hook_type post_set_user_hook = NULL;
+
 #ifdef HAS_ALTER_SYSTEM
 /* 9.4 & up */
 static bool Block_AS = false;
@@ -151,6 +156,8 @@ static void PU_hook(Node *parsetree, const char *queryString,
 #endif
 #endif
 
+extern void PostSetUserHook(bool is_reset, const char *newuser);
+
 extern Datum set_user(PG_FUNCTION_ARGS);
 void _PG_init(void);
 void _PG_fini(void);
@@ -399,6 +406,7 @@ set_user(PG_FUNCTION_ARGS)
 			  newuser);
 
 	SetCurrentRoleId(NewUserId, NewUser_is_superuser);
+	PostSetUserHook(is_reset, newuser);
 
 	PG_RETURN_TEXT_P(cstring_to_text("OK"));
 }
@@ -543,3 +551,18 @@ PU_hook(Node *parsetree, const char *queryString,
 #endif
 #endif
 }
+
+/*
+ * PostSetUserHook
+ *
+ * Handler for the post_set_user_hook
+ */
+void
+PostSetUserHook(bool is_reset, const char *username)
+{
+	if (post_reset_user_hook && is_reset)
+		(*post_reset_user_hook) ();
+	else if (post_set_user_hook)
+		(*post_set_user_hook) (username);
+	return;
+}
diff --git a/set_user.h b/set_user.h
new file mode 100644
index 0000000..b1054f8
--- /dev/null
+++ b/set_user.h
@@ -0,0 +1,11 @@
+#ifndef SET_USER_H
+#define SET_USER_H
+
+/* Expose a hook for other extensions to use after reset_user finishes up. */
+typedef void (*post_reset_user_hook_type) (void);
+extern PGDLLIMPORT post_reset_user_hook_type post_reset_user_hook;
+
+/* Expose a hook for other extensions to use after set_user finishes up. */
+typedef void (*post_set_user_hook_type) (const char *username);
+extern PGDLLIMPORT post_set_user_hook_type post_set_user_hook;
+#endif