From 1096551b47c7902290fd1b3c0d95fb9c5412ab7a Mon Sep 17 00:00:00 2001
From: Yuli Khodorkovskiy <yuli@crunchydata.com>
Date: Wed, 27 Feb 2019 15:49:52 -0500
Subject: [PATCH] Require reset_user() when using set_user()

When using a token to escalate using set_user(), the token was not
required for resetting to the original role.

https://github.com/pgaudit/set_user/issues/22
---
 set_user.c | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/set_user.c b/set_user.c
index c9b9503..984d983 100644
--- a/set_user.c
+++ b/set_user.c
@@ -581,9 +581,27 @@ PU_hook(Node *parsetree, const char *queryString,
 				if ((strcmp(((VariableSetStmt *) parsetree)->name,
 					 "log_statement") == 0) &&
 					Block_LS)
+				{
 					ereport(ERROR,
 							(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
 							 errmsg("\"SET log_statement\" blocked by set_user config")));
+				}
+				else if ((strcmp(((VariableSetStmt *) parsetree)->name,
+					 "role") == 0))
+				{
+					ereport(ERROR,
+							(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+							 errmsg("\"RESET role\" blocked by set_user"),
+							 errhint("\"Use `SELECT reset_user();` to reset role\"")));
+				}
+				else if ((strcmp(((VariableSetStmt *) parsetree)->name,
+					 "user") == 0))
+				{
+					ereport(ERROR,
+							(errcode(ERRCODE_INSUFFICIENT_PRIVILEGE),
+							 errmsg("\"RESET user\" blocked by set_user"),
+							 errhint("\"Use `SELECT reset_user();` to reset user\"")));
+				}
 				break;
 			default:
 				break;