From a2b4d75bf8311ff94d8693f6bd284fcdb29598d7 Mon Sep 17 00:00:00 2001
From: Jacob Champion <jacob.champion@enterprisedb.com>
Date: Mon, 7 Oct 2024 16:56:26 -0700
Subject: [PATCH v30 1/5] Make SASL max message length configurable

The proposed OAUTHBEARER SASL mechanism will need to allow larger
messages in the exchange, since tokens are sent directly by the client.
Move this limit into the pg_be_sasl_mech struct so that it can be
changed per-mechanism.
---
 src/backend/libpq/auth-sasl.c  | 10 +---------
 src/backend/libpq/auth-scram.c |  4 +++-
 src/include/libpq/sasl.h       | 11 +++++++++++
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/src/backend/libpq/auth-sasl.c b/src/backend/libpq/auth-sasl.c
index 08b24d90b4..4039e7fa3e 100644
--- a/src/backend/libpq/auth-sasl.c
+++ b/src/backend/libpq/auth-sasl.c
@@ -20,14 +20,6 @@
 #include "libpq/pqformat.h"
 #include "libpq/sasl.h"
 
-/*
- * Maximum accepted size of SASL messages.
- *
- * The messages that the server or libpq generate are much smaller than this,
- * but have some headroom.
- */
-#define PG_MAX_SASL_MESSAGE_LENGTH	1024
-
 /*
  * Perform a SASL exchange with a libpq client, using a specific mechanism
  * implementation.
@@ -103,7 +95,7 @@ CheckSASLAuth(const pg_be_sasl_mech *mech, Port *port, char *shadow_pass,
 
 		/* Get the actual SASL message */
 		initStringInfo(&buf);
-		if (pq_getmessage(&buf, PG_MAX_SASL_MESSAGE_LENGTH))
+		if (pq_getmessage(&buf, mech->max_message_length))
 		{
 			/* EOF - pq_getmessage already logged error */
 			pfree(buf.data);
diff --git a/src/backend/libpq/auth-scram.c b/src/backend/libpq/auth-scram.c
index 03ddddc3c2..e4be4d499e 100644
--- a/src/backend/libpq/auth-scram.c
+++ b/src/backend/libpq/auth-scram.c
@@ -113,7 +113,9 @@ static int	scram_exchange(void *opaq, const char *input, int inputlen,
 const pg_be_sasl_mech pg_be_scram_mech = {
 	scram_get_mechanisms,
 	scram_init,
-	scram_exchange
+	scram_exchange,
+
+	PG_MAX_SASL_MESSAGE_LENGTH
 };
 
 /*
diff --git a/src/include/libpq/sasl.h b/src/include/libpq/sasl.h
index 7a1f970cca..3f2c02b8f2 100644
--- a/src/include/libpq/sasl.h
+++ b/src/include/libpq/sasl.h
@@ -26,6 +26,14 @@
 #define PG_SASL_EXCHANGE_SUCCESS		1
 #define PG_SASL_EXCHANGE_FAILURE		2
 
+/*
+ * Maximum accepted size of SASL messages.
+ *
+ * The messages that the server or libpq generate are much smaller than this,
+ * but have some headroom.
+ */
+#define PG_MAX_SASL_MESSAGE_LENGTH	1024
+
 /*
  * Backend SASL mechanism callbacks.
  *
@@ -127,6 +135,9 @@ typedef struct pg_be_sasl_mech
 							 const char *input, int inputlen,
 							 char **output, int *outputlen,
 							 const char **logdetail);
+
+	/* The maximum size allowed for client SASLResponses. */
+	int			max_message_length;
 } pg_be_sasl_mech;
 
 /* Common implementation for auth.c */
-- 
2.34.1