From 30b2b2f8c519c939db47016e8588a8e6ab0e553b Mon Sep 17 00:00:00 2001
From: Samay Sharma <smilingsamay@gmail.com>
Date: Tue, 15 Feb 2022 22:28:40 -0800
Subject: [PATCH v3 2/4] Add sample extension to test custom auth provider
 hooks

This change adds a new extension to src/test/modules to
test the custom authentication provider hooks. In this
extension, we use an array to define which users to
authenticate and what passwords to use. We then get
encrypted passwords from the client and match them with
the encrypted version of the password in the array.
---
 src/include/libpq/scram.h                     |  2 +-
 src/test/modules/test_auth_provider/Makefile  | 16 ++++
 .../test_auth_provider/test_auth_provider.c   | 86 +++++++++++++++++++
 3 files changed, 103 insertions(+), 1 deletion(-)
 create mode 100644 src/test/modules/test_auth_provider/Makefile
 create mode 100644 src/test/modules/test_auth_provider/test_auth_provider.c

diff --git a/src/include/libpq/scram.h b/src/include/libpq/scram.h
index e60992a0d2..c51e848c24 100644
--- a/src/include/libpq/scram.h
+++ b/src/include/libpq/scram.h
@@ -18,7 +18,7 @@
 #include "libpq/sasl.h"
 
 /* SASL implementation callbacks */
-extern const pg_be_sasl_mech pg_be_scram_mech;
+extern PGDLLIMPORT const pg_be_sasl_mech pg_be_scram_mech;
 
 /* Routines to handle and check SCRAM-SHA-256 secret */
 extern char *pg_be_scram_build_secret(const char *password);
diff --git a/src/test/modules/test_auth_provider/Makefile b/src/test/modules/test_auth_provider/Makefile
new file mode 100644
index 0000000000..17971a5c7a
--- /dev/null
+++ b/src/test/modules/test_auth_provider/Makefile
@@ -0,0 +1,16 @@
+# src/test/modules/test_auth_provider/Makefile
+
+MODULE_big = test_auth_provider
+OBJS = test_auth_provider.o
+PGFILEDESC = "test_auth_provider - provider to test auth hooks"
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_auth_provider
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_auth_provider/test_auth_provider.c b/src/test/modules/test_auth_provider/test_auth_provider.c
new file mode 100644
index 0000000000..7c4b1f3500
--- /dev/null
+++ b/src/test/modules/test_auth_provider/test_auth_provider.c
@@ -0,0 +1,86 @@
+/* -------------------------------------------------------------------------
+ *
+ * test_auth_provider.c
+ *			example authentication provider plugin
+ *
+ * Copyright (c) 2022, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *		contrib/test_auth_provider/test_auth_provider.c
+ *
+ * -------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+#include "fmgr.h"
+#include "libpq/auth.h"
+#include "libpq/libpq.h"
+#include "libpq/scram.h"
+
+PG_MODULE_MAGIC;
+
+void _PG_init(void);
+
+static char *get_encrypted_password_for_user(char *user_name);
+
+/*
+ * List of usernames / passwords to approve. Here we are not
+ * getting passwords from Postgres but from this list. In a more real-life
+ * extension, you can fetch valid credentials and authentication tokens /
+ * passwords from an external authentication provider.
+ */
+char credentials[3][3][50] = {
+	{"bob","alice","carol"},
+	{"bob123","alice123","carol123"}
+};
+
+static int TestAuthenticationCheck(Port *port)
+{
+	int result = STATUS_ERROR;
+	char *real_pass;
+	const char *logdetail = NULL;
+
+	real_pass = get_encrypted_password_for_user(port->user_name);
+	if (real_pass)
+	{
+		result = CheckSASLAuth(&pg_be_scram_mech, port, real_pass, &logdetail);
+		pfree(real_pass);
+	}
+
+	if (result == STATUS_OK)
+		set_authn_id(port, port->user_name);
+
+	return result;
+}
+
+/*
+ * Get SCRAM encrypted version of the password for user.
+ */
+static char *
+get_encrypted_password_for_user(char *user_name)
+{
+	char *password = NULL;
+	int i;
+	for (i=0; i<3; i++)
+	{
+		if (strcmp(user_name, credentials[0][i]) == 0)
+		{
+			password = pstrdup(pg_be_scram_build_secret(credentials[1][i]));
+		}
+	}
+
+	return password;
+}
+
+static const char *TestAuthenticationError(Port *port)
+{
+	char *error_message = (char *)palloc (100);
+	sprintf(error_message, "Test authentication failed for user %s", port->user_name);
+	return error_message;
+}
+
+void
+_PG_init(void)
+{
+	RegisterAuthProvider("test", TestAuthenticationCheck, TestAuthenticationError);
+}
-- 
2.34.1