public inbox for [email protected]  
help / color / mirror / Atom feed
From: Jacob Champion <[email protected]>
To: [email protected] <[email protected]>
To: [email protected] <[email protected]>
Cc: [email protected] <[email protected]>
Subject: Re: [PoC] Federated Authn/z with OAUTHBEARER
Date: Wed, 25 Aug 2021 18:41:39 +0000
Message-ID: <[email protected]> (raw)
In-Reply-To: <[email protected]>
References: <[email protected]>
	<[email protected]>
	<[email protected]>

On Tue, 2021-06-22 at 23:22 +0000, Jacob Champion wrote:
> On Fri, 2021-06-18 at 11:31 +0300, Heikki Linnakangas wrote:
> > 
> > A few small things caught my eye in the backend oauth_exchange function:
> > 
> > > +       /* Handle the client's initial message. */
> > > +       p = strdup(input);
> > 
> > this strdup() should be pstrdup().
> 
> Thanks, I'll fix that in the next re-roll.
> 
> > In the same function, there are a bunch of reports like this:
> > 
> > >                    ereport(ERROR,
> > > +                          (errcode(ERRCODE_PROTOCOL_VIOLATION),
> > > +                           errmsg("malformed OAUTHBEARER message"),
> > > +                           errdetail("Comma expected, but found character \"%s\".",
> > > +                                     sanitize_char(*p))));
> > 
> > I don't think the double quotes are needed here, because sanitize_char 
> > will return quotes if it's a single character. So it would end up 
> > looking like this: ... found character "'x'".
> 
> I'll fix this too. Thanks!

v2, attached, incorporates Heikki's suggested fixes and also rebases on
top of latest HEAD, which had the SASL refactoring changes committed
last month.

The biggest change from the last patchset is 0001, an attempt at
enabling jsonapi in the frontend without the use of palloc(), based on
suggestions by Michael and Tom from last commitfest. I've also made
some improvements to the pytest suite. No major changes to the OAuth
implementation yet.

--Jacob


Attachments:

  [text/x-patch] v2-0001-common-jsonapi-support-FRONTEND-clients.patch (20.4K, ../[email protected]/2-v2-0001-common-jsonapi-support-FRONTEND-clients.patch)
  download | inline diff:
From 8c4b82940efb7e0f0f33ac915d5f7969a36e3644 Mon Sep 17 00:00:00 2001
From: Jacob Champion <[email protected]>
Date: Mon, 3 May 2021 15:38:26 -0700
Subject: [PATCH v2 1/5] common/jsonapi: support FRONTEND clients

Based on a patch by Michael Paquier.

For frontend code, use PQExpBuffer instead of StringInfo. This requires
us to track allocation failures so that we can return JSON_OUT_OF_MEMORY
as needed. json_errdetail() now allocates its error message inside
memory owned by the JsonLexContext, so clients don't need to worry about
freeing it.

For convenience, the backend now has destroyJsonLexContext() to mirror
other create/destroy APIs. The frontend has init/term versions of the
API to handle stack-allocated JsonLexContexts.

We can now partially revert b44669b2ca, now that json_errdetail() works
correctly.
---
 src/backend/utils/adt/jsonfuncs.c             |   4 +-
 src/bin/pg_verifybackup/parse_manifest.c      |  13 +-
 src/bin/pg_verifybackup/t/005_bad_manifest.pl |   2 +-
 src/common/Makefile                           |   2 +-
 src/common/jsonapi.c                          | 290 +++++++++++++-----
 src/include/common/jsonapi.h                  |  47 ++-
 6 files changed, 270 insertions(+), 88 deletions(-)

diff --git a/src/backend/utils/adt/jsonfuncs.c b/src/backend/utils/adt/jsonfuncs.c
index 5fd54b64b5..fa39751188 100644
--- a/src/backend/utils/adt/jsonfuncs.c
+++ b/src/backend/utils/adt/jsonfuncs.c
@@ -723,9 +723,7 @@ json_object_keys(PG_FUNCTION_ARGS)
 		pg_parse_json_or_ereport(lex, sem);
 		/* keys are now in state->result */
 
-		pfree(lex->strval->data);
-		pfree(lex->strval);
-		pfree(lex);
+		destroyJsonLexContext(lex);
 		pfree(sem);
 
 		MemoryContextSwitchTo(oldcontext);
diff --git a/src/bin/pg_verifybackup/parse_manifest.c b/src/bin/pg_verifybackup/parse_manifest.c
index c7ccc78c70..6cedb7435f 100644
--- a/src/bin/pg_verifybackup/parse_manifest.c
+++ b/src/bin/pg_verifybackup/parse_manifest.c
@@ -119,7 +119,7 @@ void
 json_parse_manifest(JsonManifestParseContext *context, char *buffer,
 					size_t size)
 {
-	JsonLexContext *lex;
+	JsonLexContext lex = {0};
 	JsonParseErrorType json_error;
 	JsonSemAction sem;
 	JsonManifestParseState parse;
@@ -129,8 +129,8 @@ json_parse_manifest(JsonManifestParseContext *context, char *buffer,
 	parse.state = JM_EXPECT_TOPLEVEL_START;
 	parse.saw_version_field = false;
 
-	/* Create a JSON lexing context. */
-	lex = makeJsonLexContextCstringLen(buffer, size, PG_UTF8, true);
+	/* Initialize a JSON lexing context. */
+	initJsonLexContextCstringLen(&lex, buffer, size, PG_UTF8, true);
 
 	/* Set up semantic actions. */
 	sem.semstate = &parse;
@@ -145,14 +145,17 @@ json_parse_manifest(JsonManifestParseContext *context, char *buffer,
 	sem.scalar = json_manifest_scalar;
 
 	/* Run the actual JSON parser. */
-	json_error = pg_parse_json(lex, &sem);
+	json_error = pg_parse_json(&lex, &sem);
 	if (json_error != JSON_SUCCESS)
-		json_manifest_parse_failure(context, "parsing failed");
+		json_manifest_parse_failure(context, json_errdetail(json_error, &lex));
 	if (parse.state != JM_EXPECT_EOF)
 		json_manifest_parse_failure(context, "manifest ended unexpectedly");
 
 	/* Verify the manifest checksum. */
 	verify_manifest_checksum(&parse, buffer, size);
+
+	/* Clean up. */
+	termJsonLexContext(&lex);
 }
 
 /*
diff --git a/src/bin/pg_verifybackup/t/005_bad_manifest.pl b/src/bin/pg_verifybackup/t/005_bad_manifest.pl
index 4f5b8f5a49..9f8a100a71 100644
--- a/src/bin/pg_verifybackup/t/005_bad_manifest.pl
+++ b/src/bin/pg_verifybackup/t/005_bad_manifest.pl
@@ -16,7 +16,7 @@ my $tempdir = TestLib::tempdir;
 
 test_bad_manifest(
 	'input string ended unexpectedly',
-	qr/could not parse backup manifest: parsing failed/,
+	qr/could not parse backup manifest: The input string ended unexpectedly/,
 	<<EOM);
 {
 EOM
diff --git a/src/common/Makefile b/src/common/Makefile
index 880722fcf5..5ecb09a8c4 100644
--- a/src/common/Makefile
+++ b/src/common/Makefile
@@ -40,7 +40,7 @@ override CPPFLAGS += -DVAL_LDFLAGS_EX="\"$(LDFLAGS_EX)\""
 override CPPFLAGS += -DVAL_LDFLAGS_SL="\"$(LDFLAGS_SL)\""
 override CPPFLAGS += -DVAL_LIBS="\"$(LIBS)\""
 
-override CPPFLAGS := -DFRONTEND -I. -I$(top_srcdir)/src/common $(CPPFLAGS)
+override CPPFLAGS := -DFRONTEND -I. -I$(top_srcdir)/src/common -I$(libpq_srcdir) $(CPPFLAGS)
 LIBS += $(PTHREAD_LIBS)
 
 # If you add objects here, see also src/tools/msvc/Mkvcbuild.pm
diff --git a/src/common/jsonapi.c b/src/common/jsonapi.c
index 5504072b4f..3a9620f739 100644
--- a/src/common/jsonapi.c
+++ b/src/common/jsonapi.c
@@ -20,10 +20,39 @@
 #include "common/jsonapi.h"
 #include "mb/pg_wchar.h"
 
-#ifndef FRONTEND
+#ifdef FRONTEND
+#include "pqexpbuffer.h"
+#else
+#include "lib/stringinfo.h"
 #include "miscadmin.h"
 #endif
 
+/*
+ * In backend, we will use palloc/pfree along with StringInfo.  In frontend, use
+ * malloc and PQExpBuffer, and return JSON_OUT_OF_MEMORY on out-of-memory.
+ */
+#ifdef FRONTEND
+
+#define STRDUP(s) strdup(s)
+#define ALLOC(size) malloc(size)
+
+#define appendStrVal		appendPQExpBuffer
+#define appendStrValChar	appendPQExpBufferChar
+#define createStrVal		createPQExpBuffer
+#define resetStrVal			resetPQExpBuffer
+
+#else /* !FRONTEND */
+
+#define STRDUP(s) pstrdup(s)
+#define ALLOC(size) palloc(size)
+
+#define appendStrVal		appendStringInfo
+#define appendStrValChar	appendStringInfoChar
+#define createStrVal		makeStringInfo
+#define resetStrVal			resetStringInfo
+
+#endif
+
 /*
  * The context of the parser is maintained by the recursive descent
  * mechanism, but is passed explicitly to the error reporting routine
@@ -132,10 +161,12 @@ IsValidJsonNumber(const char *str, int len)
 	return (!numeric_error) && (total_len == dummy_lex.input_length);
 }
 
+#ifndef FRONTEND
+
 /*
  * makeJsonLexContextCstringLen
  *
- * lex constructor, with or without StringInfo object for de-escaped lexemes.
+ * lex constructor, with or without a string object for de-escaped lexemes.
  *
  * Without is better as it makes the processing faster, so only make one
  * if really required.
@@ -145,13 +176,66 @@ makeJsonLexContextCstringLen(char *json, int len, int encoding, bool need_escape
 {
 	JsonLexContext *lex = palloc0(sizeof(JsonLexContext));
 
+	initJsonLexContextCstringLen(lex, json, len, encoding, need_escapes);
+
+	return lex;
+}
+
+void
+destroyJsonLexContext(JsonLexContext *lex)
+{
+	termJsonLexContext(lex);
+	pfree(lex);
+}
+
+#endif /* !FRONTEND */
+
+void
+initJsonLexContextCstringLen(JsonLexContext *lex, char *json, int len, int encoding, bool need_escapes)
+{
 	lex->input = lex->token_terminator = lex->line_start = json;
 	lex->line_number = 1;
 	lex->input_length = len;
 	lex->input_encoding = encoding;
-	if (need_escapes)
-		lex->strval = makeStringInfo();
-	return lex;
+	lex->parse_strval = need_escapes;
+	if (lex->parse_strval)
+	{
+		/*
+		 * This call can fail in FRONTEND code. We defer error handling to time
+		 * of use (json_lex_string()) since there's no way to signal failure
+		 * here, and we might not need to parse any strings anyway.
+		 */
+		lex->strval = createStrVal();
+	}
+	lex->errormsg = NULL;
+}
+
+void
+termJsonLexContext(JsonLexContext *lex)
+{
+	static const JsonLexContext empty = {0};
+
+	if (lex->strval)
+	{
+#ifdef FRONTEND
+		destroyPQExpBuffer(lex->strval);
+#else
+		pfree(lex->strval->data);
+		pfree(lex->strval);
+#endif
+	}
+
+	if (lex->errormsg)
+	{
+#ifdef FRONTEND
+		destroyPQExpBuffer(lex->errormsg);
+#else
+		pfree(lex->errormsg->data);
+		pfree(lex->errormsg);
+#endif
+	}
+
+	*lex = empty;
 }
 
 /*
@@ -217,7 +301,7 @@ json_count_array_elements(JsonLexContext *lex, int *elements)
 	 * etc, so doing this with a copy makes that safe.
 	 */
 	memcpy(&copylex, lex, sizeof(JsonLexContext));
-	copylex.strval = NULL;		/* not interested in values here */
+	copylex.parse_strval = false;		/* not interested in values here */
 	copylex.lex_level++;
 
 	count = 0;
@@ -279,14 +363,21 @@ parse_scalar(JsonLexContext *lex, JsonSemAction *sem)
 	/* extract the de-escaped string value, or the raw lexeme */
 	if (lex_peek(lex) == JSON_TOKEN_STRING)
 	{
-		if (lex->strval != NULL)
-			val = pstrdup(lex->strval->data);
+		if (lex->parse_strval)
+		{
+			val = STRDUP(lex->strval->data);
+			if (val == NULL)
+				return JSON_OUT_OF_MEMORY;
+		}
 	}
 	else
 	{
 		int			len = (lex->token_terminator - lex->token_start);
 
-		val = palloc(len + 1);
+		val = ALLOC(len + 1);
+		if (val == NULL)
+			return JSON_OUT_OF_MEMORY;
+
 		memcpy(val, lex->token_start, len);
 		val[len] = '\0';
 	}
@@ -320,8 +411,12 @@ parse_object_field(JsonLexContext *lex, JsonSemAction *sem)
 
 	if (lex_peek(lex) != JSON_TOKEN_STRING)
 		return report_parse_error(JSON_PARSE_STRING, lex);
-	if ((ostart != NULL || oend != NULL) && lex->strval != NULL)
-		fname = pstrdup(lex->strval->data);
+	if ((ostart != NULL || oend != NULL) && lex->parse_strval)
+	{
+		fname = STRDUP(lex->strval->data);
+		if (fname == NULL)
+			return JSON_OUT_OF_MEMORY;
+	}
 	result = json_lex(lex);
 	if (result != JSON_SUCCESS)
 		return result;
@@ -368,6 +463,10 @@ parse_object(JsonLexContext *lex, JsonSemAction *sem)
 	JsonParseErrorType result;
 
 #ifndef FRONTEND
+	/*
+	 * TODO: clients need some way to put a bound on stack growth. Parse level
+	 * limits maybe?
+	 */
 	check_stack_depth();
 #endif
 
@@ -676,8 +775,15 @@ json_lex_string(JsonLexContext *lex)
 	int			len;
 	int			hi_surrogate = -1;
 
-	if (lex->strval != NULL)
-		resetStringInfo(lex->strval);
+	if (lex->parse_strval)
+	{
+#ifdef FRONTEND
+		/* make sure initialization succeeded */
+		if (lex->strval == NULL)
+			return JSON_OUT_OF_MEMORY;
+#endif
+		resetStrVal(lex->strval);
+	}
 
 	Assert(lex->input_length > 0);
 	s = lex->token_start;
@@ -737,7 +843,7 @@ json_lex_string(JsonLexContext *lex)
 						return JSON_UNICODE_ESCAPE_FORMAT;
 					}
 				}
-				if (lex->strval != NULL)
+				if (lex->parse_strval)
 				{
 					/*
 					 * Combine surrogate pairs.
@@ -797,19 +903,19 @@ json_lex_string(JsonLexContext *lex)
 
 						unicode_to_utf8(ch, (unsigned char *) utf8str);
 						utf8len = pg_utf_mblen((unsigned char *) utf8str);
-						appendBinaryStringInfo(lex->strval, utf8str, utf8len);
+						appendBinaryPQExpBuffer(lex->strval, utf8str, utf8len);
 					}
 					else if (ch <= 0x007f)
 					{
 						/* The ASCII range is the same in all encodings */
-						appendStringInfoChar(lex->strval, (char) ch);
+						appendPQExpBufferChar(lex->strval, (char) ch);
 					}
 					else
 						return JSON_UNICODE_HIGH_ESCAPE;
 #endif							/* FRONTEND */
 				}
 			}
-			else if (lex->strval != NULL)
+			else if (lex->parse_strval)
 			{
 				if (hi_surrogate != -1)
 					return JSON_UNICODE_LOW_SURROGATE;
@@ -819,22 +925,22 @@ json_lex_string(JsonLexContext *lex)
 					case '"':
 					case '\\':
 					case '/':
-						appendStringInfoChar(lex->strval, *s);
+						appendStrValChar(lex->strval, *s);
 						break;
 					case 'b':
-						appendStringInfoChar(lex->strval, '\b');
+						appendStrValChar(lex->strval, '\b');
 						break;
 					case 'f':
-						appendStringInfoChar(lex->strval, '\f');
+						appendStrValChar(lex->strval, '\f');
 						break;
 					case 'n':
-						appendStringInfoChar(lex->strval, '\n');
+						appendStrValChar(lex->strval, '\n');
 						break;
 					case 'r':
-						appendStringInfoChar(lex->strval, '\r');
+						appendStrValChar(lex->strval, '\r');
 						break;
 					case 't':
-						appendStringInfoChar(lex->strval, '\t');
+						appendStrValChar(lex->strval, '\t');
 						break;
 					default:
 						/* Not a valid string escape, so signal error. */
@@ -858,12 +964,12 @@ json_lex_string(JsonLexContext *lex)
 			}
 
 		}
-		else if (lex->strval != NULL)
+		else if (lex->parse_strval)
 		{
 			if (hi_surrogate != -1)
 				return JSON_UNICODE_LOW_SURROGATE;
 
-			appendStringInfoChar(lex->strval, *s);
+			appendStrValChar(lex->strval, *s);
 		}
 
 	}
@@ -871,6 +977,11 @@ json_lex_string(JsonLexContext *lex)
 	if (hi_surrogate != -1)
 		return JSON_UNICODE_LOW_SURROGATE;
 
+#ifdef FRONTEND
+	if (lex->parse_strval && PQExpBufferBroken(lex->strval))
+		return JSON_OUT_OF_MEMORY;
+#endif
+
 	/* Hooray, we found the end of the string! */
 	lex->prev_token_terminator = lex->token_terminator;
 	lex->token_terminator = s + 1;
@@ -1043,72 +1154,93 @@ report_parse_error(JsonParseContext ctx, JsonLexContext *lex)
 	return JSON_SUCCESS;		/* silence stupider compilers */
 }
 
-
-#ifndef FRONTEND
-/*
- * Extract the current token from a lexing context, for error reporting.
- */
-static char *
-extract_token(JsonLexContext *lex)
-{
-	int			toklen = lex->token_terminator - lex->token_start;
-	char	   *token = palloc(toklen + 1);
-
-	memcpy(token, lex->token_start, toklen);
-	token[toklen] = '\0';
-	return token;
-}
-
 /*
  * Construct a detail message for a JSON error.
  *
- * Note that the error message generated by this routine may not be
- * palloc'd, making it unsafe for frontend code as there is no way to
- * know if this can be safery pfree'd or not.
+ * The returned allocation is either static or owned by the JsonLexContext and
+ * should not be freed.
  */
 char *
 json_errdetail(JsonParseErrorType error, JsonLexContext *lex)
 {
+	int		toklen = lex->token_terminator - lex->token_start;
+
+	if (error == JSON_OUT_OF_MEMORY)
+	{
+		/* Short circuit. Allocating anything for this case is unhelpful. */
+		return _("out of memory");
+	}
+
+	if (lex->errormsg)
+		resetStrVal(lex->errormsg);
+	else
+		lex->errormsg = createStrVal();
+
 	switch (error)
 	{
 		case JSON_SUCCESS:
 			/* fall through to the error code after switch */
 			break;
 		case JSON_ESCAPING_INVALID:
-			return psprintf(_("Escape sequence \"\\%s\" is invalid."),
-							extract_token(lex));
+			appendStrVal(lex->errormsg,
+						 _("Escape sequence \"\\%.*s\" is invalid."),
+						 toklen, lex->token_start);
+			break;
 		case JSON_ESCAPING_REQUIRED:
-			return psprintf(_("Character with value 0x%02x must be escaped."),
-							(unsigned char) *(lex->token_terminator));
+			appendStrVal(lex->errormsg,
+						 _("Character with value 0x%02x must be escaped."),
+						 (unsigned char) *(lex->token_terminator));
+			break;
 		case JSON_EXPECTED_END:
-			return psprintf(_("Expected end of input, but found \"%s\"."),
-							extract_token(lex));
+			appendStrVal(lex->errormsg,
+						 _("Expected end of input, but found \"%.*s\"."),
+						 toklen, lex->token_start);
+			break;
 		case JSON_EXPECTED_ARRAY_FIRST:
-			return psprintf(_("Expected array element or \"]\", but found \"%s\"."),
-							extract_token(lex));
+			appendStrVal(lex->errormsg,
+						 _("Expected array element or \"]\", but found \"%.*s\"."),
+						 toklen, lex->token_start);
+			break;
 		case JSON_EXPECTED_ARRAY_NEXT:
-			return psprintf(_("Expected \",\" or \"]\", but found \"%s\"."),
-							extract_token(lex));
+			appendStrVal(lex->errormsg,
+						 _("Expected \",\" or \"]\", but found \"%.*s\"."),
+						 toklen, lex->token_start);
+			break;
 		case JSON_EXPECTED_COLON:
-			return psprintf(_("Expected \":\", but found \"%s\"."),
-							extract_token(lex));
+			appendStrVal(lex->errormsg,
+						 _("Expected \":\", but found \"%.*s\"."),
+						 toklen, lex->token_start);
+			break;
 		case JSON_EXPECTED_JSON:
-			return psprintf(_("Expected JSON value, but found \"%s\"."),
-							extract_token(lex));
+			appendStrVal(lex->errormsg,
+						 _("Expected JSON value, but found \"%.*s\"."),
+						 toklen, lex->token_start);
+			break;
 		case JSON_EXPECTED_MORE:
 			return _("The input string ended unexpectedly.");
 		case JSON_EXPECTED_OBJECT_FIRST:
-			return psprintf(_("Expected string or \"}\", but found \"%s\"."),
-							extract_token(lex));
+			appendStrVal(lex->errormsg,
+						 _("Expected string or \"}\", but found \"%.*s\"."),
+						 toklen, lex->token_start);
+			break;
 		case JSON_EXPECTED_OBJECT_NEXT:
-			return psprintf(_("Expected \",\" or \"}\", but found \"%s\"."),
-							extract_token(lex));
+			appendStrVal(lex->errormsg,
+						 _("Expected \",\" or \"}\", but found \"%.*s\"."),
+						 toklen, lex->token_start);
+			break;
 		case JSON_EXPECTED_STRING:
-			return psprintf(_("Expected string, but found \"%s\"."),
-							extract_token(lex));
+			appendStrVal(lex->errormsg,
+						 _("Expected string, but found \"%.*s\"."),
+						 toklen, lex->token_start);
+			break;
 		case JSON_INVALID_TOKEN:
-			return psprintf(_("Token \"%s\" is invalid."),
-							extract_token(lex));
+			appendStrVal(lex->errormsg,
+						 _("Token \"%.*s\" is invalid."),
+						 toklen, lex->token_start);
+			break;
+		case JSON_OUT_OF_MEMORY:
+			/* should have been handled above; use the error path */
+			break;
 		case JSON_UNICODE_CODE_POINT_ZERO:
 			return _("\\u0000 cannot be converted to text.");
 		case JSON_UNICODE_ESCAPE_FORMAT:
@@ -1122,12 +1254,22 @@ json_errdetail(JsonParseErrorType error, JsonLexContext *lex)
 			return _("Unicode low surrogate must follow a high surrogate.");
 	}
 
-	/*
-	 * We don't use a default: case, so that the compiler will warn about
-	 * unhandled enum values.  But this needs to be here anyway to cover the
-	 * possibility of an incorrect input.
-	 */
-	elog(ERROR, "unexpected json parse error type: %d", (int) error);
-	return NULL;
-}
+	/* Note that lex->errormsg can be NULL in FRONTEND code. */
+	if (lex->errormsg && !lex->errormsg->data[0])
+	{
+		/*
+		 * We don't use a default: case, so that the compiler will warn about
+		 * unhandled enum values.  But this needs to be here anyway to cover the
+		 * possibility of an incorrect input.
+		 */
+		appendStrVal(lex->errormsg,
+					 "unexpected json parse error type: %d", (int) error);
+	}
+
+#ifdef FRONTEND
+	if (PQExpBufferBroken(lex->errormsg))
+		return _("out of memory while constructing error description");
 #endif
+
+	return lex->errormsg->data;
+}
diff --git a/src/include/common/jsonapi.h b/src/include/common/jsonapi.h
index ec3dfce9c3..dc71ab2cd3 100644
--- a/src/include/common/jsonapi.h
+++ b/src/include/common/jsonapi.h
@@ -14,8 +14,6 @@
 #ifndef JSONAPI_H
 #define JSONAPI_H
 
-#include "lib/stringinfo.h"
-
 typedef enum
 {
 	JSON_TOKEN_INVALID,
@@ -48,6 +46,7 @@ typedef enum
 	JSON_EXPECTED_OBJECT_NEXT,
 	JSON_EXPECTED_STRING,
 	JSON_INVALID_TOKEN,
+	JSON_OUT_OF_MEMORY,
 	JSON_UNICODE_CODE_POINT_ZERO,
 	JSON_UNICODE_ESCAPE_FORMAT,
 	JSON_UNICODE_HIGH_ESCAPE,
@@ -55,6 +54,17 @@ typedef enum
 	JSON_UNICODE_LOW_SURROGATE
 } JsonParseErrorType;
 
+/*
+ * Don't depend on the internal type header for strval; if callers need access
+ * then they can include the appropriate header themselves.
+ */
+#ifdef FRONTEND
+#define StrValType PQExpBufferData
+#else
+#define StrValType StringInfoData
+#endif
+
+typedef struct StrValType StrValType;
 
 /*
  * All the fields in this structure should be treated as read-only.
@@ -81,7 +91,9 @@ typedef struct JsonLexContext
 	int			lex_level;
 	int			line_number;	/* line number, starting from 1 */
 	char	   *line_start;		/* where that line starts within input */
-	StringInfo	strval;
+	bool		parse_strval;
+	StrValType *strval;			/* only used if parse_strval == true */
+	StrValType *errormsg;
 } JsonLexContext;
 
 typedef void (*json_struct_action) (void *state);
@@ -141,9 +153,10 @@ extern JsonSemAction nullSemAction;
  */
 extern JsonParseErrorType json_count_array_elements(JsonLexContext *lex,
 													int *elements);
+#ifndef FRONTEND
 
 /*
- * constructor for JsonLexContext, with or without strval element.
+ * allocating constructor for JsonLexContext, with or without strval element.
  * If supplied, the strval element will contain a de-escaped version of
  * the lexeme. However, doing this imposes a performance penalty, so
  * it should be avoided if the de-escaped lexeme is not required.
@@ -153,6 +166,32 @@ extern JsonLexContext *makeJsonLexContextCstringLen(char *json,
 													int encoding,
 													bool need_escapes);
 
+/*
+ * Counterpart to makeJsonLexContextCstringLen(): clears and deallocates lex.
+ * The context pointer should not be used after this call.
+ */
+extern void destroyJsonLexContext(JsonLexContext *lex);
+
+#endif /* !FRONTEND */
+
+/*
+ * stack constructor for JsonLexContext, with or without strval element.
+ * If supplied, the strval element will contain a de-escaped version of
+ * the lexeme. However, doing this imposes a performance penalty, so
+ * it should be avoided if the de-escaped lexeme is not required.
+ */
+extern void initJsonLexContextCstringLen(JsonLexContext *lex,
+										 char *json,
+										 int len,
+										 int encoding,
+										 bool need_escapes);
+
+/*
+ * Counterpart to initJsonLexContextCstringLen(): clears the contents of lex,
+ * but does not deallocate lex itself.
+ */
+extern void termJsonLexContext(JsonLexContext *lex);
+
 /* lex one token */
 extern JsonParseErrorType json_lex(JsonLexContext *lex);
 
-- 
2.25.1



  [text/x-patch] v2-0002-libpq-add-OAUTHBEARER-SASL-mechanism.patch (35.7K, ../[email protected]/3-v2-0002-libpq-add-OAUTHBEARER-SASL-mechanism.patch)
  download | inline diff:
From 52ac4bd25ca19735eb2bd863e8b1549ccbe6560a Mon Sep 17 00:00:00 2001
From: Jacob Champion <[email protected]>
Date: Tue, 13 Apr 2021 10:27:27 -0700
Subject: [PATCH v2 2/5] libpq: add OAUTHBEARER SASL mechanism

DO NOT USE THIS PROOF OF CONCEPT IN PRODUCTION.

Implement OAUTHBEARER (RFC 7628) and OAuth 2.0 Device Authorization
Grants (RFC 8628) on the client side. When speaking to a OAuth-enabled
server, it looks a bit like this:

    $ psql 'host=example.org oauth_client_id=f02c6361-0635-...'
    Visit https://oauth.example.org/login and enter the code: FPQ2-M4BG

The OAuth issuer must support device authorization. No other OAuth flows
are currently implemented.

The client implementation requires libiddawc and its development
headers. Configure --with-oauth (and --with-includes/--with-libraries to
point at the iddawc installation, if it's in a custom location).

Several TODOs:
- don't retry forever if the server won't accept our token
- perform several sanity checks on the OAuth issuer's responses
- handle cases where the client has been set up with an issuer and
  scope, but the Postgres server wants to use something different
- improve error debuggability during the OAuth handshake
- ...and more.
---
 configure                            | 100 ++++
 configure.ac                         |  19 +
 src/Makefile.global.in               |   1 +
 src/include/common/oauth-common.h    |  19 +
 src/include/pg_config.h.in           |   6 +
 src/interfaces/libpq/Makefile        |   7 +-
 src/interfaces/libpq/fe-auth-oauth.c | 745 +++++++++++++++++++++++++++
 src/interfaces/libpq/fe-auth-sasl.h  |   5 +-
 src/interfaces/libpq/fe-auth-scram.c |   6 +-
 src/interfaces/libpq/fe-auth.c       |  42 +-
 src/interfaces/libpq/fe-auth.h       |   3 +
 src/interfaces/libpq/fe-connect.c    |  38 ++
 src/interfaces/libpq/libpq-int.h     |   8 +
 13 files changed, 980 insertions(+), 19 deletions(-)
 create mode 100644 src/include/common/oauth-common.h
 create mode 100644 src/interfaces/libpq/fe-auth-oauth.c

diff --git a/configure b/configure
index 7542fe30a1..2ddbe9a1d9 100755
--- a/configure
+++ b/configure
@@ -713,6 +713,7 @@ with_uuid
 with_readline
 with_systemd
 with_selinux
+with_oauth
 with_ldap
 with_krb_srvnam
 krb_srvtab
@@ -856,6 +857,7 @@ with_krb_srvnam
 with_pam
 with_bsd_auth
 with_ldap
+with_oauth
 with_bonjour
 with_selinux
 with_systemd
@@ -1562,6 +1564,7 @@ Optional Packages:
   --with-pam              build with PAM support
   --with-bsd-auth         build with BSD Authentication support
   --with-ldap             build with LDAP support
+  --with-oauth            build with OAuth 2.0 support
   --with-bonjour          build with Bonjour support
   --with-selinux          build with SELinux support
   --with-systemd          build with systemd support
@@ -8144,6 +8147,42 @@ $as_echo "$with_ldap" >&6; }
 
 
 
+#
+# OAuth 2.0
+#
+{ $as_echo "$as_me:${as_lineno-$LINENO}: checking whether to build with OAuth support" >&5
+$as_echo_n "checking whether to build with OAuth support... " >&6; }
+
+
+
+# Check whether --with-oauth was given.
+if test "${with_oauth+set}" = set; then :
+  withval=$with_oauth;
+  case $withval in
+    yes)
+
+$as_echo "#define USE_OAUTH 1" >>confdefs.h
+
+      ;;
+    no)
+      :
+      ;;
+    *)
+      as_fn_error $? "no argument expected for --with-oauth option" "$LINENO" 5
+      ;;
+  esac
+
+else
+  with_oauth=no
+
+fi
+
+
+{ $as_echo "$as_me:${as_lineno-$LINENO}: result: $with_oauth" >&5
+$as_echo "$with_oauth" >&6; }
+
+
+
 #
 # Bonjour
 #
@@ -13084,6 +13123,56 @@ fi
 
 
 
+if test "$with_oauth" = yes ; then
+  { $as_echo "$as_me:${as_lineno-$LINENO}: checking for i_init_session in -liddawc" >&5
+$as_echo_n "checking for i_init_session in -liddawc... " >&6; }
+if ${ac_cv_lib_iddawc_i_init_session+:} false; then :
+  $as_echo_n "(cached) " >&6
+else
+  ac_check_lib_save_LIBS=$LIBS
+LIBS="-liddawc  $LIBS"
+cat confdefs.h - <<_ACEOF >conftest.$ac_ext
+/* end confdefs.h.  */
+
+/* Override any GCC internal prototype to avoid an error.
+   Use char because int might match the return type of a GCC
+   builtin and then its argument prototype would still apply.  */
+#ifdef __cplusplus
+extern "C"
+#endif
+char i_init_session ();
+int
+main ()
+{
+return i_init_session ();
+  ;
+  return 0;
+}
+_ACEOF
+if ac_fn_c_try_link "$LINENO"; then :
+  ac_cv_lib_iddawc_i_init_session=yes
+else
+  ac_cv_lib_iddawc_i_init_session=no
+fi
+rm -f core conftest.err conftest.$ac_objext \
+    conftest$ac_exeext conftest.$ac_ext
+LIBS=$ac_check_lib_save_LIBS
+fi
+{ $as_echo "$as_me:${as_lineno-$LINENO}: result: $ac_cv_lib_iddawc_i_init_session" >&5
+$as_echo "$ac_cv_lib_iddawc_i_init_session" >&6; }
+if test "x$ac_cv_lib_iddawc_i_init_session" = xyes; then :
+  cat >>confdefs.h <<_ACEOF
+#define HAVE_LIBIDDAWC 1
+_ACEOF
+
+  LIBS="-liddawc $LIBS"
+
+else
+  as_fn_error $? "library 'iddawc' is required for OAuth support" "$LINENO" 5
+fi
+
+fi
+
 # for contrib/sepgsql
 if test "$with_selinux" = yes; then
   { $as_echo "$as_me:${as_lineno-$LINENO}: checking for security_compute_create_name in -lselinux" >&5
@@ -13978,6 +14067,17 @@ fi
 
 done
 
+fi
+
+if test "$with_oauth" != no; then
+  ac_fn_c_check_header_mongrel "$LINENO" "iddawc.h" "ac_cv_header_iddawc_h" "$ac_includes_default"
+if test "x$ac_cv_header_iddawc_h" = xyes; then :
+
+else
+  as_fn_error $? "header file <iddawc.h> is required for OAuth" "$LINENO" 5
+fi
+
+
 fi
 
 if test "$PORTNAME" = "win32" ; then
diff --git a/configure.ac b/configure.ac
index ed3cdb9a8e..22026476d9 100644
--- a/configure.ac
+++ b/configure.ac
@@ -851,6 +851,17 @@ AC_MSG_RESULT([$with_ldap])
 AC_SUBST(with_ldap)
 
 
+#
+# OAuth 2.0
+#
+AC_MSG_CHECKING([whether to build with OAuth support])
+PGAC_ARG_BOOL(with, oauth, no,
+              [build with OAuth 2.0 support],
+              [AC_DEFINE([USE_OAUTH], 1, [Define to 1 to build with OAuth 2.0 support. (--with-oauth)])])
+AC_MSG_RESULT([$with_oauth])
+AC_SUBST(with_oauth)
+
+
 #
 # Bonjour
 #
@@ -1321,6 +1332,10 @@ fi
 AC_SUBST(LDAP_LIBS_FE)
 AC_SUBST(LDAP_LIBS_BE)
 
+if test "$with_oauth" = yes ; then
+  AC_CHECK_LIB(iddawc, i_init_session, [], [AC_MSG_ERROR([library 'iddawc' is required for OAuth support])])
+fi
+
 # for contrib/sepgsql
 if test "$with_selinux" = yes; then
   AC_CHECK_LIB(selinux, security_compute_create_name, [],
@@ -1531,6 +1546,10 @@ elif test "$with_uuid" = ossp ; then
       [AC_MSG_ERROR([header file <ossp/uuid.h> or <uuid.h> is required for OSSP UUID])])])
 fi
 
+if test "$with_oauth" != no; then
+  AC_CHECK_HEADER(iddawc.h, [], [AC_MSG_ERROR([header file <iddawc.h> is required for OAuth])])
+fi
+
 if test "$PORTNAME" = "win32" ; then
    AC_CHECK_HEADERS(crtdefs.h)
 fi
diff --git a/src/Makefile.global.in b/src/Makefile.global.in
index 6e2f224cc4..d67912711e 100644
--- a/src/Makefile.global.in
+++ b/src/Makefile.global.in
@@ -193,6 +193,7 @@ with_ldap	= @with_ldap@
 with_libxml	= @with_libxml@
 with_libxslt	= @with_libxslt@
 with_llvm	= @with_llvm@
+with_oauth	= @with_oauth@
 with_system_tzdata = @with_system_tzdata@
 with_uuid	= @with_uuid@
 with_zlib	= @with_zlib@
diff --git a/src/include/common/oauth-common.h b/src/include/common/oauth-common.h
new file mode 100644
index 0000000000..3fa95ac7e8
--- /dev/null
+++ b/src/include/common/oauth-common.h
@@ -0,0 +1,19 @@
+/*-------------------------------------------------------------------------
+ *
+ * oauth-common.h
+ *		Declarations for helper functions used for OAuth/OIDC authentication
+ *
+ * Portions Copyright (c) 1996-2021, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * src/include/common/oauth-common.h
+ *
+ *-------------------------------------------------------------------------
+ */
+#ifndef OAUTH_COMMON_H
+#define OAUTH_COMMON_H
+
+/* Name of SASL mechanism per IANA */
+#define OAUTHBEARER_NAME "OAUTHBEARER"
+
+#endif /* OAUTH_COMMON_H */
diff --git a/src/include/pg_config.h.in b/src/include/pg_config.h.in
index 15ffdd895a..f82ab38536 100644
--- a/src/include/pg_config.h.in
+++ b/src/include/pg_config.h.in
@@ -331,6 +331,9 @@
 /* Define to 1 if you have the `crypto' library (-lcrypto). */
 #undef HAVE_LIBCRYPTO
 
+/* Define to 1 if you have the `iddawc' library (-liddawc). */
+#undef HAVE_LIBIDDAWC
+
 /* Define to 1 if you have the `ldap' library (-lldap). */
 #undef HAVE_LIBLDAP
 
@@ -926,6 +929,9 @@
 /* Define to select named POSIX semaphores. */
 #undef USE_NAMED_POSIX_SEMAPHORES
 
+/* Define to 1 to build with OAuth 2.0 support. (--with-oauth) */
+#undef USE_OAUTH
+
 /* Define to 1 to build with OpenSSL support. (--with-ssl=openssl) */
 #undef USE_OPENSSL
 
diff --git a/src/interfaces/libpq/Makefile b/src/interfaces/libpq/Makefile
index 7cbdeb589b..3cdf19294b 100644
--- a/src/interfaces/libpq/Makefile
+++ b/src/interfaces/libpq/Makefile
@@ -62,6 +62,11 @@ OBJS += \
 	fe-secure-gssapi.o
 endif
 
+ifeq ($(with_oauth),yes)
+OBJS += \
+	fe-auth-oauth.o
+endif
+
 ifeq ($(PORTNAME), cygwin)
 override shlib = cyg$(NAME)$(DLSUFFIX)
 endif
@@ -83,7 +88,7 @@ endif
 # that are built correctly for use in a shlib.
 SHLIB_LINK_INTERNAL = -lpgcommon_shlib -lpgport_shlib
 ifneq ($(PORTNAME), win32)
-SHLIB_LINK += $(filter -lcrypt -ldes -lcom_err -lcrypto -lk5crypto -lkrb5 -lgssapi_krb5 -lgss -lgssapi -lssl -lsocket -lnsl -lresolv -lintl -lm, $(LIBS)) $(LDAP_LIBS_FE) $(PTHREAD_LIBS)
+SHLIB_LINK += $(filter -lcrypt -ldes -lcom_err -lcrypto -lk5crypto -lkrb5 -lgssapi_krb5 -lgss -lgssapi -lssl -liddawc -lsocket -lnsl -lresolv -lintl -lm, $(LIBS)) $(LDAP_LIBS_FE) $(PTHREAD_LIBS)
 else
 SHLIB_LINK += $(filter -lcrypt -ldes -lcom_err -lcrypto -lk5crypto -lkrb5 -lgssapi32 -lssl -lsocket -lnsl -lresolv -lintl -lm $(PTHREAD_LIBS), $(LIBS)) $(LDAP_LIBS_FE)
 endif
diff --git a/src/interfaces/libpq/fe-auth-oauth.c b/src/interfaces/libpq/fe-auth-oauth.c
new file mode 100644
index 0000000000..91d2c69f16
--- /dev/null
+++ b/src/interfaces/libpq/fe-auth-oauth.c
@@ -0,0 +1,745 @@
+/*-------------------------------------------------------------------------
+ *
+ * fe-auth-oauth.c
+ *	   The front-end (client) implementation of OAuth/OIDC authentication.
+ *
+ * Portions Copyright (c) 1996-2021, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * IDENTIFICATION
+ *	  src/interfaces/libpq/fe-auth-oauth.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include <iddawc.h>
+
+#include "postgres_fe.h"
+
+#include "common/base64.h"
+#include "common/hmac.h"
+#include "common/jsonapi.h"
+#include "common/oauth-common.h"
+#include "fe-auth.h"
+#include "mb/pg_wchar.h"
+
+/* The exported OAuth callback mechanism. */
+static void *oauth_init(PGconn *conn, const char *password,
+						const char *sasl_mechanism);
+static void oauth_exchange(void *opaq, bool final,
+						   char *input, int inputlen,
+						   char **output, int *outputlen,
+						   bool *done, bool *success);
+static bool oauth_channel_bound(void *opaq);
+static void oauth_free(void *opaq);
+
+const pg_fe_sasl_mech pg_oauth_mech = {
+	oauth_init,
+	oauth_exchange,
+	oauth_channel_bound,
+	oauth_free,
+};
+
+typedef enum
+{
+	FE_OAUTH_INIT,
+	FE_OAUTH_BEARER_SENT,
+	FE_OAUTH_SERVER_ERROR,
+} fe_oauth_state_enum;
+
+typedef struct
+{
+	fe_oauth_state_enum state;
+
+	PGconn	   *conn;
+} fe_oauth_state;
+
+static void *
+oauth_init(PGconn *conn, const char *password,
+		   const char *sasl_mechanism)
+{
+	fe_oauth_state *state;
+
+	/*
+	 * We only support one SASL mechanism here; anything else is programmer
+	 * error.
+	 */
+	Assert(sasl_mechanism != NULL);
+	Assert(!strcmp(sasl_mechanism, OAUTHBEARER_NAME));
+
+	state = malloc(sizeof(*state));
+	if (!state)
+		return NULL;
+
+	state->state = FE_OAUTH_INIT;
+	state->conn = conn;
+
+	return state;
+}
+
+static const char *
+iddawc_error_string(int errcode)
+{
+	switch (errcode)
+	{
+		case I_OK:
+			return "I_OK";
+
+		case I_ERROR:
+			return "I_ERROR";
+
+		case I_ERROR_PARAM:
+			return "I_ERROR_PARAM";
+
+		case I_ERROR_MEMORY:
+			return "I_ERROR_MEMORY";
+
+		case I_ERROR_UNAUTHORIZED:
+			return "I_ERROR_UNAUTHORIZED";
+
+		case I_ERROR_SERVER:
+			return "I_ERROR_SERVER";
+	}
+
+	return "<unknown>";
+}
+
+static void
+iddawc_error(PGconn *conn, int errcode, const char *msg)
+{
+	appendPQExpBufferStr(&conn->errorMessage, libpq_gettext(msg));
+	appendPQExpBuffer(&conn->errorMessage,
+					  libpq_gettext(" (iddawc error %s)\n"),
+					  iddawc_error_string(errcode));
+}
+
+static void
+iddawc_request_error(PGconn *conn, struct _i_session *i, int err, const char *msg)
+{
+	const char *error_code;
+	const char *desc;
+
+	appendPQExpBuffer(&conn->errorMessage, "%s: ", libpq_gettext(msg));
+
+	error_code = i_get_str_parameter(i, I_OPT_ERROR);
+	if (!error_code)
+	{
+		/*
+		 * The server didn't give us any useful information, so just print the
+		 * error code.
+		 */
+		appendPQExpBuffer(&conn->errorMessage,
+						  libpq_gettext("(iddawc error %s)\n"),
+						  iddawc_error_string(err));
+		return;
+	}
+
+	/* If the server gave a string description, print that too. */
+	desc = i_get_str_parameter(i, I_OPT_ERROR_DESCRIPTION);
+	if (desc)
+		appendPQExpBuffer(&conn->errorMessage, "%s ", desc);
+
+	appendPQExpBuffer(&conn->errorMessage, "(%s)\n", error_code);
+}
+
+static char *
+get_auth_token(PGconn *conn)
+{
+	PQExpBuffer	token_buf = NULL;
+	struct _i_session session;
+	int			err;
+	int			auth_method;
+	bool		user_prompted = false;
+	const char *verification_uri;
+	const char *user_code;
+	const char *access_token;
+	const char *token_type;
+	char	   *token = NULL;
+
+	if (!conn->oauth_discovery_uri)
+		return strdup(""); /* ask the server for one */
+
+	i_init_session(&session);
+
+	if (!conn->oauth_client_id)
+	{
+		/* We can't talk to a server without a client identifier. */
+		appendPQExpBufferStr(&conn->errorMessage,
+							 libpq_gettext("no oauth_client_id is set for the connection"));
+		goto cleanup;
+	}
+
+	token_buf = createPQExpBuffer();
+
+	if (!token_buf)
+		goto cleanup;
+
+	err = i_set_str_parameter(&session, I_OPT_OPENID_CONFIG_ENDPOINT, conn->oauth_discovery_uri);
+	if (err)
+	{
+		iddawc_error(conn, err, "failed to set OpenID config endpoint");
+		goto cleanup;
+	}
+
+	err = i_get_openid_config(&session);
+	if (err)
+	{
+		iddawc_error(conn, err, "failed to fetch OpenID discovery document");
+		goto cleanup;
+	}
+
+	if (!i_get_str_parameter(&session, I_OPT_TOKEN_ENDPOINT))
+	{
+		appendPQExpBufferStr(&conn->errorMessage,
+							 libpq_gettext("issuer has no token endpoint"));
+		goto cleanup;
+	}
+
+	if (!i_get_str_parameter(&session, I_OPT_DEVICE_AUTHORIZATION_ENDPOINT))
+	{
+		appendPQExpBufferStr(&conn->errorMessage,
+							 libpq_gettext("issuer does not support device authorization"));
+		goto cleanup;
+	}
+
+	err = i_set_response_type(&session, I_RESPONSE_TYPE_DEVICE_CODE);
+	if (err)
+	{
+		iddawc_error(conn, err, "failed to set device code response type");
+		goto cleanup;
+	}
+
+	auth_method = I_TOKEN_AUTH_METHOD_NONE;
+	if (conn->oauth_client_secret && *conn->oauth_client_secret)
+		auth_method = I_TOKEN_AUTH_METHOD_SECRET_BASIC;
+
+	err = i_set_parameter_list(&session,
+		I_OPT_CLIENT_ID, conn->oauth_client_id,
+		I_OPT_CLIENT_SECRET, conn->oauth_client_secret,
+		I_OPT_TOKEN_METHOD, auth_method,
+		I_OPT_SCOPE, conn->oauth_scope,
+		I_OPT_NONE
+	);
+	if (err)
+	{
+		iddawc_error(conn, err, "failed to set client identifier");
+		goto cleanup;
+	}
+
+	err = i_run_device_auth_request(&session);
+	if (err)
+	{
+		iddawc_request_error(conn, &session, err,
+							"failed to obtain device authorization");
+		goto cleanup;
+	}
+
+	verification_uri = i_get_str_parameter(&session, I_OPT_DEVICE_AUTH_VERIFICATION_URI);
+	if (!verification_uri)
+	{
+		appendPQExpBufferStr(&conn->errorMessage,
+							 libpq_gettext("issuer did not provide a verification URI"));
+		goto cleanup;
+	}
+
+	user_code = i_get_str_parameter(&session, I_OPT_DEVICE_AUTH_USER_CODE);
+	if (!user_code)
+	{
+		appendPQExpBufferStr(&conn->errorMessage,
+							 libpq_gettext("issuer did not provide a user code"));
+		goto cleanup;
+	}
+
+	/*
+	 * Poll the token endpoint until either the user logs in and authorizes the
+	 * use of a token, or a hard failure occurs. We perform one ping _before_
+	 * prompting the user, so that we don't make them do the work of logging in
+	 * only to find that the token endpoint is completely unreachable.
+	 */
+	err = i_run_token_request(&session);
+	while (err)
+	{
+		const char *error_code;
+		uint		interval;
+
+		error_code = i_get_str_parameter(&session, I_OPT_ERROR);
+
+		/*
+		 * authorization_pending and slow_down are the only acceptable errors;
+		 * anything else and we bail.
+		 */
+		if (!error_code || (strcmp(error_code, "authorization_pending")
+							&& strcmp(error_code, "slow_down")))
+		{
+			iddawc_request_error(conn, &session, err,
+								"OAuth token retrieval failed");
+			goto cleanup;
+		}
+
+		if (!user_prompted)
+		{
+			/*
+			 * Now that we know the token endpoint isn't broken, give the user
+			 * the login instructions.
+			 */
+			pqInternalNotice(&conn->noticeHooks,
+							 "Visit %s and enter the code: %s",
+							 verification_uri, user_code);
+
+			user_prompted = true;
+		}
+
+		/*
+		 * We are required to wait between polls; the server tells us how long.
+		 * TODO: if interval's not set, we need to default to five seconds
+		 * TODO: sanity check the interval
+		 */
+		interval = i_get_int_parameter(&session, I_OPT_DEVICE_AUTH_INTERVAL);
+
+		/*
+		 * A slow_down error requires us to permanently increase our retry
+		 * interval by five seconds. RFC 8628, Sec. 3.5.
+		 */
+		if (!strcmp(error_code, "slow_down"))
+		{
+			interval += 5;
+			i_set_int_parameter(&session, I_OPT_DEVICE_AUTH_INTERVAL, interval);
+		}
+
+		sleep(interval);
+
+		/*
+		 * XXX Reset the error code before every call, because iddawc won't do
+		 * that for us. This matters if the server first sends a "pending" error
+		 * code, then later hard-fails without sending an error code to
+		 * overwrite the first one.
+		 *
+		 * That we have to do this at all seems like a bug in iddawc.
+		 */
+		i_set_str_parameter(&session, I_OPT_ERROR, NULL);
+
+		err = i_run_token_request(&session);
+	}
+
+	access_token = i_get_str_parameter(&session, I_OPT_ACCESS_TOKEN);
+	token_type = i_get_str_parameter(&session, I_OPT_TOKEN_TYPE);
+
+	if (!access_token || !token_type || strcasecmp(token_type, "Bearer"))
+	{
+		appendPQExpBufferStr(&conn->errorMessage,
+							 libpq_gettext("issuer did not provide a bearer token"));
+		goto cleanup;
+	}
+
+	appendPQExpBufferStr(token_buf, "Bearer ");
+	appendPQExpBufferStr(token_buf, access_token);
+
+	if (PQExpBufferBroken(token_buf))
+		goto cleanup;
+
+	token = strdup(token_buf->data);
+
+cleanup:
+	if (token_buf)
+		destroyPQExpBuffer(token_buf);
+	i_clean_session(&session);
+
+	return token;
+}
+
+#define kvsep "\x01"
+
+static char *
+client_initial_response(PGconn *conn)
+{
+	static const char * const resp_format = "n,," kvsep "auth=%s" kvsep kvsep;
+
+	PQExpBuffer	token_buf;
+	PQExpBuffer	discovery_buf = NULL;
+	char	   *token = NULL;
+	char	   *response = NULL;
+
+	token_buf = createPQExpBuffer();
+	if (!token_buf)
+		goto cleanup;
+
+	/*
+	 * If we don't yet have a discovery URI, but the user gave us an explicit
+	 * issuer, use the .well-known discovery URI for that issuer.
+	 */
+	if (!conn->oauth_discovery_uri && conn->oauth_issuer)
+	{
+		discovery_buf = createPQExpBuffer();
+		if (!discovery_buf)
+			goto cleanup;
+
+		appendPQExpBufferStr(discovery_buf, conn->oauth_issuer);
+		appendPQExpBufferStr(discovery_buf, "/.well-known/openid-configuration");
+
+		if (PQExpBufferBroken(discovery_buf))
+			goto cleanup;
+
+		conn->oauth_discovery_uri = strdup(discovery_buf->data);
+	}
+
+	token = get_auth_token(conn);
+	if (!token)
+		goto cleanup;
+
+	appendPQExpBuffer(token_buf, resp_format, token);
+	if (PQExpBufferBroken(token_buf))
+		goto cleanup;
+
+	response = strdup(token_buf->data);
+
+cleanup:
+	if (token)
+		free(token);
+	if (discovery_buf)
+		destroyPQExpBuffer(discovery_buf);
+	if (token_buf)
+		destroyPQExpBuffer(token_buf);
+
+	return response;
+}
+
+#define ERROR_STATUS_FIELD "status"
+#define ERROR_SCOPE_FIELD "scope"
+#define ERROR_OPENID_CONFIGURATION_FIELD "openid-configuration"
+
+struct json_ctx
+{
+	char		   *errmsg; /* any non-NULL value stops all processing */
+	PQExpBufferData errbuf; /* backing memory for errmsg */
+	int				nested; /* nesting level (zero is the top) */
+
+	const char	   *target_field_name; /* points to a static allocation */
+	char		  **target_field;      /* see below */
+
+	/* target_field, if set, points to one of the following: */
+	char		   *status;
+	char		   *scope;
+	char		   *discovery_uri;
+};
+
+#define oauth_json_has_error(ctx) \
+	(PQExpBufferDataBroken((ctx)->errbuf) || (ctx)->errmsg)
+
+#define oauth_json_set_error(ctx, ...) \
+	do { \
+		appendPQExpBuffer(&(ctx)->errbuf, __VA_ARGS__); \
+		(ctx)->errmsg = (ctx)->errbuf.data; \
+	} while (0)
+
+static void
+oauth_json_object_start(void *state)
+{
+	struct json_ctx	   *ctx = state;
+
+	if (oauth_json_has_error(ctx))
+		return; /* short-circuit */
+
+	if (ctx->target_field)
+	{
+		Assert(ctx->nested == 1);
+
+		oauth_json_set_error(ctx,
+							 libpq_gettext("field \"%s\" must be a string"),
+							 ctx->target_field_name);
+	}
+
+	++ctx->nested;
+}
+
+static void
+oauth_json_object_end(void *state)
+{
+	struct json_ctx	   *ctx = state;
+
+	if (oauth_json_has_error(ctx))
+		return; /* short-circuit */
+
+	--ctx->nested;
+}
+
+static void
+oauth_json_object_field_start(void *state, char *name, bool isnull)
+{
+	struct json_ctx	   *ctx = state;
+
+	if (oauth_json_has_error(ctx))
+	{
+		/* short-circuit */
+		free(name);
+		return;
+	}
+
+	if (ctx->nested == 1)
+	{
+		if (!strcmp(name, ERROR_STATUS_FIELD))
+		{
+			ctx->target_field_name = ERROR_STATUS_FIELD;
+			ctx->target_field = &ctx->status;
+		}
+		else if (!strcmp(name, ERROR_SCOPE_FIELD))
+		{
+			ctx->target_field_name = ERROR_SCOPE_FIELD;
+			ctx->target_field = &ctx->scope;
+		}
+		else if (!strcmp(name, ERROR_OPENID_CONFIGURATION_FIELD))
+		{
+			ctx->target_field_name = ERROR_OPENID_CONFIGURATION_FIELD;
+			ctx->target_field = &ctx->discovery_uri;
+		}
+	}
+
+	free(name);
+}
+
+static void
+oauth_json_array_start(void *state)
+{
+	struct json_ctx	   *ctx = state;
+
+	if (oauth_json_has_error(ctx))
+		return; /* short-circuit */
+
+	if (!ctx->nested)
+	{
+		ctx->errmsg = libpq_gettext("top-level element must be an object");
+	}
+	else if (ctx->target_field)
+	{
+		Assert(ctx->nested == 1);
+
+		oauth_json_set_error(ctx,
+							 libpq_gettext("field \"%s\" must be a string"),
+							 ctx->target_field_name);
+	}
+}
+
+static void
+oauth_json_scalar(void *state, char *token, JsonTokenType type)
+{
+	struct json_ctx	   *ctx = state;
+
+	if (oauth_json_has_error(ctx))
+	{
+		/* short-circuit */
+		free(token);
+		return;
+	}
+
+	if (!ctx->nested)
+	{
+		ctx->errmsg = libpq_gettext("top-level element must be an object");
+	}
+	else if (ctx->target_field)
+	{
+		Assert(ctx->nested == 1);
+
+		if (type == JSON_TOKEN_STRING)
+		{
+			*ctx->target_field = token;
+
+			ctx->target_field = NULL;
+			ctx->target_field_name = NULL;
+
+			return; /* don't free the token we're using */
+		}
+
+		oauth_json_set_error(ctx,
+							 libpq_gettext("field \"%s\" must be a string"),
+							 ctx->target_field_name);
+	}
+
+	free(token);
+}
+
+static bool
+handle_oauth_sasl_error(PGconn *conn, char *msg, int msglen)
+{
+	JsonLexContext		lex = {0};
+	JsonSemAction		sem = {0};
+	JsonParseErrorType	err;
+	struct json_ctx		ctx = {0};
+	char			   *errmsg = NULL;
+
+	/* Sanity check. */
+	if (strlen(msg) != msglen)
+	{
+		appendPQExpBufferStr(&conn->errorMessage,
+							 libpq_gettext("server's error message contained an embedded NULL"));
+		return false;
+	}
+
+	initJsonLexContextCstringLen(&lex, msg, msglen, PG_UTF8, true);
+
+	initPQExpBuffer(&ctx.errbuf);
+	sem.semstate = &ctx;
+
+	sem.object_start = oauth_json_object_start;
+	sem.object_end = oauth_json_object_end;
+	sem.object_field_start = oauth_json_object_field_start;
+	sem.array_start = oauth_json_array_start;
+	sem.scalar = oauth_json_scalar;
+
+	err = pg_parse_json(&lex, &sem);
+
+	if (err != JSON_SUCCESS)
+	{
+		errmsg = json_errdetail(err, &lex);
+	}
+	else if (PQExpBufferDataBroken(ctx.errbuf))
+	{
+		errmsg = libpq_gettext("out of memory");
+	}
+	else if (ctx.errmsg)
+	{
+		errmsg = ctx.errmsg;
+	}
+
+	if (errmsg)
+		appendPQExpBuffer(&conn->errorMessage,
+						  libpq_gettext("failed to parse server's error response: %s"),
+						  errmsg);
+
+	/* Don't need the error buffer or the JSON lexer anymore. */
+	termPQExpBuffer(&ctx.errbuf);
+	termJsonLexContext(&lex);
+
+	if (errmsg)
+		return false;
+
+	/* TODO: what if these override what the user already specified? */
+	if (ctx.discovery_uri)
+	{
+		if (conn->oauth_discovery_uri)
+			free(conn->oauth_discovery_uri);
+
+		conn->oauth_discovery_uri = ctx.discovery_uri;
+	}
+
+	if (ctx.scope)
+	{
+		if (conn->oauth_scope)
+			free(conn->oauth_scope);
+
+		conn->oauth_scope = ctx.scope;
+	}
+	/* TODO: missing error scope should clear any existing connection scope */
+
+	if (!ctx.status)
+	{
+		appendPQExpBuffer(&conn->errorMessage,
+						  libpq_gettext("server sent error response without a status"));
+		return false;
+	}
+
+	if (!strcmp(ctx.status, "invalid_token"))
+	{
+		/*
+		 * invalid_token is the only error code we'll automatically retry for,
+		 * but only if we have enough information to do so.
+		 */
+		if (conn->oauth_discovery_uri)
+			conn->oauth_want_retry = true;
+	}
+	/* TODO: include status in hard failure message */
+
+	return true;
+}
+
+static void
+oauth_exchange(void *opaq, bool final,
+			   char *input, int inputlen,
+			   char **output, int *outputlen,
+			   bool *done, bool *success)
+{
+	fe_oauth_state *state = opaq;
+	PGconn	   *conn = state->conn;
+
+	*done = false;
+	*success = false;
+	*output = NULL;
+	*outputlen = 0;
+
+	switch (state->state)
+	{
+		case FE_OAUTH_INIT:
+			Assert(inputlen == -1);
+
+			*output = client_initial_response(conn);
+			if (!*output)
+				goto error;
+
+			*outputlen = strlen(*output);
+			state->state = FE_OAUTH_BEARER_SENT;
+
+			break;
+
+		case FE_OAUTH_BEARER_SENT:
+			if (final)
+			{
+				/* TODO: ensure there is no message content here. */
+				*done = true;
+				*success = true;
+
+				break;
+			}
+
+			/*
+			 * Error message sent by the server.
+			 */
+			if (!handle_oauth_sasl_error(conn, input, inputlen))
+				goto error;
+
+			/*
+			 * Respond with the required dummy message (RFC 7628, sec. 3.2.3).
+			 */
+			*output = strdup(kvsep);
+			*outputlen = strlen(*output); /* == 1 */
+
+			state->state = FE_OAUTH_SERVER_ERROR;
+			break;
+
+		case FE_OAUTH_SERVER_ERROR:
+			/*
+			 * After an error, the server should send an error response to fail
+			 * the SASL handshake, which is handled in higher layers.
+			 *
+			 * If we get here, the server either sent *another* challenge which
+			 * isn't defined in the RFC, or completed the handshake successfully
+			 * after telling us it was going to fail. Neither is acceptable.
+			 */
+			appendPQExpBufferStr(&conn->errorMessage,
+								 libpq_gettext("server sent additional OAuth data after error\n"));
+			goto error;
+
+		default:
+			appendPQExpBufferStr(&conn->errorMessage,
+								 libpq_gettext("invalid OAuth exchange state\n"));
+			goto error;
+	}
+
+	return;
+
+error:
+	*done = true;
+	*success = false;
+}
+
+static bool
+oauth_channel_bound(void *opaq)
+{
+	/* This mechanism does not support channel binding. */
+	return false;
+}
+
+static void
+oauth_free(void *opaq)
+{
+	fe_oauth_state *state = opaq;
+
+	free(state);
+}
diff --git a/src/interfaces/libpq/fe-auth-sasl.h b/src/interfaces/libpq/fe-auth-sasl.h
index 3d7ee576f2..0920102908 100644
--- a/src/interfaces/libpq/fe-auth-sasl.h
+++ b/src/interfaces/libpq/fe-auth-sasl.h
@@ -65,6 +65,8 @@ typedef struct pg_fe_sasl_mech
 	 *
 	 *	state:	   The opaque mechanism state returned by init()
 	 *
+	 *	final:	   true if the server has sent a final exchange outcome
+	 *
 	 *	input:	   The challenge data sent by the server, or NULL when
 	 *			   generating a client-first initial response (that is, when
 	 *			   the server expects the client to send a message to start
@@ -92,7 +94,8 @@ typedef struct pg_fe_sasl_mech
 	 *			   Ignored if *done is false.
 	 *--------
 	 */
-	void		(*exchange) (void *state, char *input, int inputlen,
+	void		(*exchange) (void *state, bool final,
+							 char *input, int inputlen,
 							 char **output, int *outputlen,
 							 bool *done, bool *success);
 
diff --git a/src/interfaces/libpq/fe-auth-scram.c b/src/interfaces/libpq/fe-auth-scram.c
index 4337e89ce9..489cbeda50 100644
--- a/src/interfaces/libpq/fe-auth-scram.c
+++ b/src/interfaces/libpq/fe-auth-scram.c
@@ -24,7 +24,8 @@
 /* The exported SCRAM callback mechanism. */
 static void *scram_init(PGconn *conn, const char *password,
 						const char *sasl_mechanism);
-static void scram_exchange(void *opaq, char *input, int inputlen,
+static void scram_exchange(void *opaq, bool final,
+						   char *input, int inputlen,
 						   char **output, int *outputlen,
 						   bool *done, bool *success);
 static bool scram_channel_bound(void *opaq);
@@ -205,7 +206,8 @@ scram_free(void *opaq)
  * Exchange a SCRAM message with backend.
  */
 static void
-scram_exchange(void *opaq, char *input, int inputlen,
+scram_exchange(void *opaq, bool final,
+			   char *input, int inputlen,
 			   char **output, int *outputlen,
 			   bool *done, bool *success)
 {
diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c
index 3421ed4685..0b5b91962a 100644
--- a/src/interfaces/libpq/fe-auth.c
+++ b/src/interfaces/libpq/fe-auth.c
@@ -39,6 +39,7 @@
 #endif
 
 #include "common/md5.h"
+#include "common/oauth-common.h"
 #include "common/scram-common.h"
 #include "fe-auth.h"
 #include "fe-auth-sasl.h"
@@ -423,7 +424,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 	bool		success;
 	const char *selected_mechanism;
 	PQExpBufferData mechanism_buf;
-	char	   *password;
+	char	   *password = NULL;
 
 	initPQExpBuffer(&mechanism_buf);
 
@@ -445,8 +446,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 	/*
 	 * Parse the list of SASL authentication mechanisms in the
 	 * AuthenticationSASL message, and select the best mechanism that we
-	 * support.  SCRAM-SHA-256-PLUS and SCRAM-SHA-256 are the only ones
-	 * supported at the moment, listed by order of decreasing importance.
+	 * support.  Mechanisms are listed by order of decreasing importance.
 	 */
 	selected_mechanism = NULL;
 	for (;;)
@@ -486,6 +486,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 				{
 					selected_mechanism = SCRAM_SHA_256_PLUS_NAME;
 					conn->sasl = &pg_scram_mech;
+					conn->password_needed = true;
 				}
 #else
 				/*
@@ -523,7 +524,17 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 		{
 			selected_mechanism = SCRAM_SHA_256_NAME;
 			conn->sasl = &pg_scram_mech;
+			conn->password_needed = true;
 		}
+#ifdef USE_OAUTH
+		else if (strcmp(mechanism_buf.data, OAUTHBEARER_NAME) == 0 &&
+				!selected_mechanism)
+		{
+			selected_mechanism = OAUTHBEARER_NAME;
+			conn->sasl = &pg_oauth_mech;
+			conn->password_needed = false;
+		}
+#endif
 	}
 
 	if (!selected_mechanism)
@@ -548,18 +559,19 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 
 	/*
 	 * First, select the password to use for the exchange, complaining if
-	 * there isn't one.  Currently, all supported SASL mechanisms require a
-	 * password, so we can just go ahead here without further distinction.
+	 * there isn't one and the SASL mechanism needs it.
 	 */
-	conn->password_needed = true;
-	password = conn->connhost[conn->whichhost].password;
-	if (password == NULL)
-		password = conn->pgpass;
-	if (password == NULL || password[0] == '\0')
+	if (conn->password_needed)
 	{
-		appendPQExpBufferStr(&conn->errorMessage,
-							 PQnoPasswordSupplied);
-		goto error;
+		password = conn->connhost[conn->whichhost].password;
+		if (password == NULL)
+			password = conn->pgpass;
+		if (password == NULL || password[0] == '\0')
+		{
+			appendPQExpBufferStr(&conn->errorMessage,
+								 PQnoPasswordSupplied);
+			goto error;
+		}
 	}
 
 	Assert(conn->sasl);
@@ -577,7 +589,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 		goto oom_error;
 
 	/* Get the mechanism-specific Initial Client Response, if any */
-	conn->sasl->exchange(conn->sasl_state,
+	conn->sasl->exchange(conn->sasl_state, false,
 						 NULL, -1,
 						 &initialresponse, &initialresponselen,
 						 &done, &success);
@@ -658,7 +670,7 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
 	/* For safety and convenience, ensure the buffer is NULL-terminated. */
 	challenge[payloadlen] = '\0';
 
-	conn->sasl->exchange(conn->sasl_state,
+	conn->sasl->exchange(conn->sasl_state, final,
 						 challenge, payloadlen,
 						 &output, &outputlen,
 						 &done, &success);
diff --git a/src/interfaces/libpq/fe-auth.h b/src/interfaces/libpq/fe-auth.h
index 63927480ee..03bea124a6 100644
--- a/src/interfaces/libpq/fe-auth.h
+++ b/src/interfaces/libpq/fe-auth.h
@@ -26,4 +26,7 @@ extern char *pg_fe_getauthname(PQExpBuffer errorMessage);
 extern const pg_fe_sasl_mech pg_scram_mech;
 extern char *pg_fe_scram_build_secret(const char *password);
 
+/* Mechanisms in fe-auth-oauth.c */
+extern const pg_fe_sasl_mech pg_oauth_mech;
+
 #endif							/* FE_AUTH_H */
diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index 49eec3e835..ba9c097060 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -344,6 +344,23 @@ static const internalPQconninfoOption PQconninfoOptions[] = {
 		"Target-Session-Attrs", "", 15, /* sizeof("prefer-standby") = 15 */
 	offsetof(struct pg_conn, target_session_attrs)},
 
+	/* OAuth v2 */
+	{"oauth_issuer", NULL, NULL, NULL,
+		"OAuth-Issuer", "", 40,
+	offsetof(struct pg_conn, oauth_issuer)},
+
+	{"oauth_client_id", NULL, NULL, NULL,
+		"OAuth-Client-ID", "", 40,
+	offsetof(struct pg_conn, oauth_client_id)},
+
+	{"oauth_client_secret", NULL, NULL, NULL,
+		"OAuth-Client-Secret", "", 40,
+	offsetof(struct pg_conn, oauth_client_secret)},
+
+	{"oauth_scope", NULL, NULL, NULL,
+		"OAuth-Scope", "", 15,
+	offsetof(struct pg_conn, oauth_scope)},
+
 	/* Terminating entry --- MUST BE LAST */
 	{NULL, NULL, NULL, NULL,
 	NULL, NULL, 0}
@@ -606,6 +623,7 @@ pqDropServerData(PGconn *conn)
 	conn->write_err_msg = NULL;
 	conn->be_pid = 0;
 	conn->be_key = 0;
+	/* conn->oauth_want_retry = false; TODO */
 }
 
 
@@ -3355,6 +3373,16 @@ keep_going:						/* We will come back to here until there is
 					/* Check to see if we should mention pgpassfile */
 					pgpassfileWarning(conn);
 
+#ifdef USE_OAUTH
+					if (conn->sasl == &pg_oauth_mech
+						&& conn->oauth_want_retry)
+					{
+						/* TODO: only allow retry once */
+						need_new_connection = true;
+						goto keep_going;
+					}
+#endif
+
 #ifdef ENABLE_GSS
 
 					/*
@@ -4129,6 +4157,16 @@ freePGconn(PGconn *conn)
 		free(conn->rowBuf);
 	if (conn->target_session_attrs)
 		free(conn->target_session_attrs);
+	if (conn->oauth_issuer)
+		free(conn->oauth_issuer);
+	if (conn->oauth_discovery_uri)
+		free(conn->oauth_discovery_uri);
+	if (conn->oauth_client_id)
+		free(conn->oauth_client_id);
+	if (conn->oauth_client_secret)
+		free(conn->oauth_client_secret);
+	if (conn->oauth_scope)
+		free(conn->oauth_scope);
 	termPQExpBuffer(&conn->errorMessage);
 	termPQExpBuffer(&conn->workBuffer);
 
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 490458adef..3d20482550 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -394,6 +394,14 @@ struct pg_conn
 	char	   *ssl_max_protocol_version;	/* maximum TLS protocol version */
 	char	   *target_session_attrs;	/* desired session properties */
 
+	/* OAuth v2 */
+	char	   *oauth_issuer;			/* token issuer URL */
+	char	   *oauth_discovery_uri;	/* URI of the issuer's discovery document */
+	char	   *oauth_client_id;		/* client identifier */
+	char	   *oauth_client_secret;	/* client secret */
+	char	   *oauth_scope;			/* access token scope */
+	bool		oauth_want_retry;		/* should we retry on failure? */
+
 	/* Optional file to write trace info to */
 	FILE	   *Pfdebug;
 	int			traceFlags;
-- 
2.25.1



  [text/x-patch] v2-0003-backend-add-OAUTHBEARER-SASL-mechanism.patch (38.6K, ../[email protected]/4-v2-0003-backend-add-OAUTHBEARER-SASL-mechanism.patch)
  download | inline diff:
From d09a00b4d52a5ed578ee5cd7623108ebdd12f202 Mon Sep 17 00:00:00 2001
From: Jacob Champion <[email protected]>
Date: Tue, 4 May 2021 16:21:11 -0700
Subject: [PATCH v2 3/5] backend: add OAUTHBEARER SASL mechanism

DO NOT USE THIS PROOF OF CONCEPT IN PRODUCTION.

Implement OAUTHBEARER (RFC 7628) on the server side. This adds a new
auth method, oauth, to pg_hba.

Because OAuth implementations vary so wildly, and bearer token
validation is heavily dependent on the issuing party, authn/z is done by
communicating with an external program: the oauth_validator_command.
This command must do the following:

1. Receive the bearer token by reading its contents from a file
   descriptor passed from the server. (The numeric value of this
   descriptor may be inserted into the oauth_validator_command using the
   %f specifier.)

   This MUST be the first action the command performs. The server will
   not begin reading stdout from the command until the token has been
   read in full, so if the command tries to print anything and hits a
   buffer limit, the backend will deadlock and time out.

2. Validate the bearer token. The correct way to do this depends on the
   issuer, but it generally involves either cryptographic operations to
   prove that the token was issued by a trusted party, or the
   presentation of the bearer token to some other party so that _it_ can
   perform validation.

   The command MUST maintain confidentiality of the bearer token, since
   in most cases it can be used just like a password. (There are ways to
   cryptographically bind tokens to client certificates, but they are
   way beyond the scope of this commit message.)

   If the token cannot be validated, the command must exit with a
   non-zero status. Further authentication/authorization is pointless if
   the bearer token wasn't issued by someone you trust.

3. Authenticate the user, authorize the user, or both:

   a. To authenticate the user, use the bearer token to retrieve some
      trusted identifier string for the end user. The exact process for
      this is, again, issuer-dependent. The command should print the
      authenticated identity string to stdout, followed by a newline.

      If the user cannot be authenticated, the validator should not
      print anything to stdout. It should also exit with a non-zero
      status, unless the token may be used to authorize the connection
      through some other means (see below).

      On a success, the command may then exit with a zero success code.
      By default, the server will then check to make sure the identity
      string matches the role that is being used (or matches a usermap
      entry, if one is in use).

   b. To optionally authorize the user, in combination with the HBA
      option trust_validator_authz=1 (see below), the validator simply
      returns a zero exit code if the client should be allowed to
      connect with its presented role (which can be passed to the
      command using the %r specifier), or a non-zero code otherwise.

      The hard part is in determining whether the given token truly
      authorizes the client to use the given role, which must
      unfortunately be left as an exercise to the reader.

      This obviously requires some care, as a poorly implemented token
      validator may silently open the entire database to anyone with a
      bearer token. But it may be a more portable approach, since OAuth
      is designed as an authorization framework, not an authentication
      framework. For example, the user's bearer token could carry an
      "allow_superuser_access" claim, which would authorize pseudonymous
      database access as any role. It's then up to the OAuth system
      administrators to ensure that allow_superuser_access is doled out
      only to the proper users.

   c. It's possible that the user can be successfully authenticated but
      isn't authorized to connect. In this case, the command may print
      the authenticated ID and then fail with a non-zero exit code.
      (This makes it easier to see what's going on in the Postgres
      logs.)

4. Token validators may optionally log to stderr. This will be printed
   verbatim into the Postgres server logs.

The oauth method supports the following HBA options (but note that two
of them are not optional, since we have no way of choosing sensible
defaults):

  issuer: Required. The URL of the OAuth issuing party, which the client
          must contact to receive a bearer token.

          Some real-world examples as of time of writing:
          - https://accounts.google.com
          - https://login.microsoft.com/[tenant-id]/v2.0

  scope:  Required. The OAuth scope(s) required for the server to
          authenticate and/or authorize the user. This is heavily
          deployment-specific, but a simple example is "openid email".

  map:    Optional. Specify a standard PostgreSQL user map; this works
          the same as with other auth methods such as peer. If a map is
          not specified, the user ID returned by the token validator
          must exactly match the role that's being requested (but see
          trust_validator_authz, below).

  trust_validator_authz:
          Optional. When set to 1, this allows the token validator to
          take full control of the authorization process. Standard user
          mapping is skipped: if the validator command succeeds, the
          client is allowed to connect under its desired role and no
          further checks are done.

Unlike the client, servers support OAuth without needing to be built
against libiddawc (since the responsibility for "speaking" OAuth/OIDC
correctly is delegated entirely to the oauth_validator_command).

Several TODOs:
- port to platforms other than "modern Linux"
- overhaul the communication with oauth_validator_command, which is
  currently a bad hack on OpenPipeStream()
- implement more sanity checks on the OAUTHBEARER message format and
  tokens sent by the client
- implement more helpful handling of HBA misconfigurations
- properly interpolate JSON when generating error responses
- use logdetail during auth failures
- deal with role names that can't be safely passed to system() without
  shell-escaping
- allow passing the configured issuer to the oauth_validator_command, to
  deal with multi-issuer setups
- ...and more.
---
 src/backend/libpq/Makefile     |   1 +
 src/backend/libpq/auth-oauth.c | 797 +++++++++++++++++++++++++++++++++
 src/backend/libpq/auth-sasl.c  |  10 +-
 src/backend/libpq/auth-scram.c |   4 +-
 src/backend/libpq/auth.c       |  26 +-
 src/backend/libpq/hba.c        |  29 +-
 src/backend/utils/misc/guc.c   |  12 +
 src/include/libpq/auth.h       |  17 +
 src/include/libpq/hba.h        |   8 +-
 src/include/libpq/oauth.h      |  24 +
 src/include/libpq/sasl.h       |  11 +
 11 files changed, 907 insertions(+), 32 deletions(-)
 create mode 100644 src/backend/libpq/auth-oauth.c
 create mode 100644 src/include/libpq/oauth.h

diff --git a/src/backend/libpq/Makefile b/src/backend/libpq/Makefile
index 6d385fd6a4..98eb2a8242 100644
--- a/src/backend/libpq/Makefile
+++ b/src/backend/libpq/Makefile
@@ -15,6 +15,7 @@ include $(top_builddir)/src/Makefile.global
 # be-fsstubs is here for historical reasons, probably belongs elsewhere
 
 OBJS = \
+	auth-oauth.o \
 	auth-sasl.o \
 	auth-scram.o \
 	auth.o \
diff --git a/src/backend/libpq/auth-oauth.c b/src/backend/libpq/auth-oauth.c
new file mode 100644
index 0000000000..c47211132c
--- /dev/null
+++ b/src/backend/libpq/auth-oauth.c
@@ -0,0 +1,797 @@
+/*-------------------------------------------------------------------------
+ *
+ * auth-oauth.c
+ *	  Server-side implementation of the SASL OAUTHBEARER mechanism.
+ *
+ * See the following RFC for more details:
+ * - RFC 7628: https://tools.ietf.org/html/rfc7628
+ *
+ * Portions Copyright (c) 1996-2021, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * src/backend/libpq/auth-oauth.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include <unistd.h>
+#include <fcntl.h>
+
+#include "common/oauth-common.h"
+#include "lib/stringinfo.h"
+#include "libpq/auth.h"
+#include "libpq/hba.h"
+#include "libpq/oauth.h"
+#include "libpq/sasl.h"
+#include "storage/fd.h"
+
+/* GUC */
+char *oauth_validator_command;
+
+static void  oauth_get_mechanisms(Port *port, StringInfo buf);
+static void *oauth_init(Port *port, const char *selected_mech, const char *shadow_pass);
+static int   oauth_exchange(void *opaq, const char *input, int inputlen,
+							char **output, int *outputlen, char **logdetail);
+
+/* Mechanism declaration */
+const pg_be_sasl_mech pg_be_oauth_mech = {
+	oauth_get_mechanisms,
+	oauth_init,
+	oauth_exchange,
+
+	PG_MAX_AUTH_TOKEN_LENGTH,
+};
+
+
+typedef enum
+{
+	OAUTH_STATE_INIT = 0,
+	OAUTH_STATE_ERROR,
+	OAUTH_STATE_FINISHED,
+} oauth_state;
+
+struct oauth_ctx
+{
+	oauth_state	state;
+	Port	   *port;
+	const char *issuer;
+	const char *scope;
+};
+
+static char *sanitize_char(char c);
+static char *parse_kvpairs_for_auth(char **input);
+static void generate_error_response(struct oauth_ctx *ctx, char **output, int *outputlen);
+static bool validate(Port *port, const char *auth, char **logdetail);
+static bool run_validator_command(Port *port, const char *token);
+static bool check_exit(FILE **fh, const char *command);
+static bool unset_cloexec(int fd);
+static bool username_ok_for_shell(const char *username);
+
+#define KVSEP 0x01
+#define AUTH_KEY "auth"
+#define BEARER_SCHEME "Bearer "
+
+static void
+oauth_get_mechanisms(Port *port, StringInfo buf)
+{
+	/* Only OAUTHBEARER is supported. */
+	appendStringInfoString(buf, OAUTHBEARER_NAME);
+	appendStringInfoChar(buf, '\0');
+}
+
+static void *
+oauth_init(Port *port, const char *selected_mech, const char *shadow_pass)
+{
+	struct oauth_ctx *ctx;
+
+	if (strcmp(selected_mech, OAUTHBEARER_NAME))
+		ereport(ERROR,
+				(errcode(ERRCODE_PROTOCOL_VIOLATION),
+				 errmsg("client selected an invalid SASL authentication mechanism")));
+
+	ctx = palloc0(sizeof(*ctx));
+
+	ctx->state = OAUTH_STATE_INIT;
+	ctx->port = port;
+
+	Assert(port->hba);
+	ctx->issuer = port->hba->oauth_issuer;
+	ctx->scope = port->hba->oauth_scope;
+
+	return ctx;
+}
+
+static int
+oauth_exchange(void *opaq, const char *input, int inputlen,
+			   char **output, int *outputlen, char **logdetail)
+{
+	char   *p;
+	char	cbind_flag;
+	char   *auth;
+
+	struct oauth_ctx *ctx = opaq;
+
+	*output = NULL;
+	*outputlen = -1;
+
+	/*
+	 * If the client didn't include an "Initial Client Response" in the
+	 * SASLInitialResponse message, send an empty challenge, to which the
+	 * client will respond with the same data that usually comes in the
+	 * Initial Client Response.
+	 */
+	if (input == NULL)
+	{
+		Assert(ctx->state == OAUTH_STATE_INIT);
+
+		*output = pstrdup("");
+		*outputlen = 0;
+		return PG_SASL_EXCHANGE_CONTINUE;
+	}
+
+	/*
+	 * Check that the input length agrees with the string length of the input.
+	 */
+	if (inputlen == 0)
+		ereport(ERROR,
+				(errcode(ERRCODE_PROTOCOL_VIOLATION),
+				 errmsg("malformed OAUTHBEARER message"),
+				 errdetail("The message is empty.")));
+	if (inputlen != strlen(input))
+		ereport(ERROR,
+				(errcode(ERRCODE_PROTOCOL_VIOLATION),
+				 errmsg("malformed OAUTHBEARER message"),
+				 errdetail("Message length does not match input length.")));
+
+	switch (ctx->state)
+	{
+		case OAUTH_STATE_INIT:
+			/* Handle this case below. */
+			break;
+
+		case OAUTH_STATE_ERROR:
+			/*
+			 * Only one response is valid for the client during authentication
+			 * failure: a single kvsep.
+			 */
+			if (inputlen != 1 || *input != KVSEP)
+				ereport(ERROR,
+						(errcode(ERRCODE_PROTOCOL_VIOLATION),
+						 errmsg("malformed OAUTHBEARER message"),
+						 errdetail("Client did not send a kvsep response.")));
+
+			/* The (failed) handshake is now complete. */
+			ctx->state = OAUTH_STATE_FINISHED;
+			return PG_SASL_EXCHANGE_FAILURE;
+
+		default:
+			elog(ERROR, "invalid OAUTHBEARER exchange state");
+			return PG_SASL_EXCHANGE_FAILURE;
+	}
+
+	/* Handle the client's initial message. */
+	p = pstrdup(input);
+
+	/*
+	 * OAUTHBEARER does not currently define a channel binding (so there is no
+	 * OAUTHBEARER-PLUS, and we do not accept a 'p' specifier). We accept a 'y'
+	 * specifier purely for the remote chance that a future specification could
+	 * define one; then future clients can still interoperate with this server
+	 * implementation. 'n' is the expected case.
+	 */
+	cbind_flag = *p;
+	switch (cbind_flag)
+	{
+		case 'p':
+			ereport(ERROR,
+					(errcode(ERRCODE_PROTOCOL_VIOLATION),
+					 errmsg("malformed OAUTHBEARER message"),
+					 errdetail("The server does not support channel binding for OAuth, but the client message includes channel binding data.")));
+			break;
+
+		case 'y': /* fall through */
+		case 'n':
+			p++;
+			if (*p != ',')
+				ereport(ERROR,
+						(errcode(ERRCODE_PROTOCOL_VIOLATION),
+						 errmsg("malformed OAUTHBEARER message"),
+						 errdetail("Comma expected, but found character %s.",
+								   sanitize_char(*p))));
+			p++;
+			break;
+
+		default:
+			ereport(ERROR,
+					(errcode(ERRCODE_PROTOCOL_VIOLATION),
+					 errmsg("malformed OAUTHBEARER message"),
+					 errdetail("Unexpected channel-binding flag %s.",
+							   sanitize_char(cbind_flag))));
+	}
+
+	/*
+	 * Forbid optional authzid (authorization identity).  We don't support it.
+	 */
+	if (*p == 'a')
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("client uses authorization identity, but it is not supported")));
+	if (*p != ',')
+		ereport(ERROR,
+				(errcode(ERRCODE_PROTOCOL_VIOLATION),
+				 errmsg("malformed OAUTHBEARER message"),
+				 errdetail("Unexpected attribute %s in client-first-message.",
+						   sanitize_char(*p))));
+	p++;
+
+	/* All remaining fields are separated by the RFC's kvsep (\x01). */
+	if (*p != KVSEP)
+		ereport(ERROR,
+				(errcode(ERRCODE_PROTOCOL_VIOLATION),
+				 errmsg("malformed OAUTHBEARER message"),
+				 errdetail("Key-value separator expected, but found character %s.",
+						   sanitize_char(*p))));
+	p++;
+
+	auth = parse_kvpairs_for_auth(&p);
+	if (!auth)
+		ereport(ERROR,
+				(errcode(ERRCODE_PROTOCOL_VIOLATION),
+				 errmsg("malformed OAUTHBEARER message"),
+				 errdetail("Message does not contain an auth value.")));
+
+	/* We should be at the end of our message. */
+	if (*p)
+		ereport(ERROR,
+				(errcode(ERRCODE_PROTOCOL_VIOLATION),
+				 errmsg("malformed OAUTHBEARER message"),
+				 errdetail("Message contains additional data after the final terminator.")));
+
+	if (!validate(ctx->port, auth, logdetail))
+	{
+		generate_error_response(ctx, output, outputlen);
+
+		ctx->state = OAUTH_STATE_ERROR;
+		return PG_SASL_EXCHANGE_CONTINUE;
+	}
+
+	ctx->state = OAUTH_STATE_FINISHED;
+	return PG_SASL_EXCHANGE_SUCCESS;
+}
+
+/*
+ * Convert an arbitrary byte to printable form.  For error messages.
+ *
+ * If it's a printable ASCII character, print it as a single character.
+ * otherwise, print it in hex.
+ *
+ * The returned pointer points to a static buffer.
+ */
+static char *
+sanitize_char(char c)
+{
+	static char buf[5];
+
+	if (c >= 0x21 && c <= 0x7E)
+		snprintf(buf, sizeof(buf), "'%c'", c);
+	else
+		snprintf(buf, sizeof(buf), "0x%02x", (unsigned char) c);
+	return buf;
+}
+
+/*
+ * Consumes all kvpairs in an OAUTHBEARER exchange message. If the "auth" key is
+ * found, its value is returned.
+ */
+static char *
+parse_kvpairs_for_auth(char **input)
+{
+	char   *pos = *input;
+	char   *auth = NULL;
+
+	/*
+	 * The relevant ABNF, from Sec. 3.1:
+	 *
+	 *     kvsep          = %x01
+	 *     key            = 1*(ALPHA)
+	 *     value          = *(VCHAR / SP / HTAB / CR / LF )
+	 *     kvpair         = key "=" value kvsep
+	 *   ;;gs2-header     = See RFC 5801
+	 *     client-resp    = (gs2-header kvsep *kvpair kvsep) / kvsep
+	 *
+	 * By the time we reach this code, the gs2-header and initial kvsep have
+	 * already been validated. We start at the beginning of the first kvpair.
+	 */
+
+	while (*pos)
+	{
+		char   *end;
+		char   *sep;
+		char   *key;
+		char   *value;
+
+		/*
+		 * Find the end of this kvpair. Note that input is null-terminated by
+		 * the SASL code, so the strchr() is bounded.
+		 */
+		end = strchr(pos, KVSEP);
+		if (!end)
+			ereport(ERROR,
+					(errcode(ERRCODE_PROTOCOL_VIOLATION),
+					 errmsg("malformed OAUTHBEARER message"),
+					 errdetail("Message contains an unterminated key/value pair.")));
+		*end = '\0';
+
+		if (pos == end)
+		{
+			/* Empty kvpair, signifying the end of the list. */
+			*input = pos + 1;
+			return auth;
+		}
+
+		/*
+		 * Find the end of the key name.
+		 *
+		 * TODO further validate the key/value grammar? empty keys, bad chars...
+		 */
+		sep = strchr(pos, '=');
+		if (!sep)
+			ereport(ERROR,
+					(errcode(ERRCODE_PROTOCOL_VIOLATION),
+					 errmsg("malformed OAUTHBEARER message"),
+					 errdetail("Message contains a key without a value.")));
+		*sep = '\0';
+
+		/* Both key and value are now safely terminated. */
+		key = pos;
+		value = sep + 1;
+
+		if (!strcmp(key, AUTH_KEY))
+		{
+			if (auth)
+				ereport(ERROR,
+						(errcode(ERRCODE_PROTOCOL_VIOLATION),
+						 errmsg("malformed OAUTHBEARER message"),
+						 errdetail("Message contains multiple auth values.")));
+
+			auth = value;
+		}
+		else
+		{
+			/*
+			 * The RFC also defines the host and port keys, but they are not
+			 * required for OAUTHBEARER and we do not use them. Also, per
+			 * Sec. 3.1, any key/value pairs we don't recognize must be ignored.
+			 */
+		}
+
+		/* Move to the next pair. */
+		pos = end + 1;
+	}
+
+	ereport(ERROR,
+			(errcode(ERRCODE_PROTOCOL_VIOLATION),
+			 errmsg("malformed OAUTHBEARER message"),
+			 errdetail("Message did not contain a final terminator.")));
+
+	return NULL; /* unreachable */
+}
+
+static void
+generate_error_response(struct oauth_ctx *ctx, char **output, int *outputlen)
+{
+	StringInfoData	buf;
+
+	/*
+	 * The admin needs to set an issuer and scope for OAuth to work. There's not
+	 * really a way to hide this from the user, either, because we can't choose
+	 * a "default" issuer, so be honest in the failure message.
+	 *
+	 * TODO: see if there's a better place to fail, earlier than this.
+	 */
+	if (!ctx->issuer || !ctx->scope)
+		ereport(FATAL,
+				(errcode(ERRCODE_INTERNAL_ERROR),
+				 errmsg("OAuth is not properly configured for this user"),
+				 errdetail_log("The issuer and scope parameters must be set in pg_hba.conf.")));
+
+
+	initStringInfo(&buf);
+
+	/*
+	 * TODO: JSON escaping
+	 */
+	appendStringInfo(&buf,
+		"{ "
+			"\"status\": \"invalid_token\", "
+			"\"openid-configuration\": \"%s/.well-known/openid-configuration\","
+			"\"scope\": \"%s\" "
+		"}",
+		ctx->issuer, ctx->scope);
+
+	*output = buf.data;
+	*outputlen = buf.len;
+}
+
+static bool
+validate(Port *port, const char *auth, char **logdetail)
+{
+	static const char * const b64_set = "abcdefghijklmnopqrstuvwxyz"
+										"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+										"0123456789-._~+/";
+
+	const char *token;
+	size_t		span;
+	int			ret;
+
+	/* TODO: handle logdetail when the test framework can check it */
+
+	/*
+	 * Only Bearer tokens are accepted. The ABNF is defined in RFC 6750, Sec.
+	 * 2.1:
+	 *
+	 *      b64token    = 1*( ALPHA / DIGIT /
+	 *                        "-" / "." / "_" / "~" / "+" / "/" ) *"="
+	 *      credentials = "Bearer" 1*SP b64token
+	 *
+	 * The "credentials" construction is what we receive in our auth value.
+	 *
+	 * Since that spec is subordinate to HTTP (i.e. the HTTP Authorization
+	 * header format; RFC 7235 Sec. 2), the "Bearer" scheme string must be
+	 * compared case-insensitively. (This is not mentioned in RFC 6750, but it's
+	 * pointed out in RFC 7628 Sec. 4.)
+	 *
+	 * TODO: handle the Authorization spec, RFC 7235 Sec. 2.1.
+	 */
+	if (strncasecmp(auth, BEARER_SCHEME, strlen(BEARER_SCHEME)))
+		return false;
+
+	/* Pull the bearer token out of the auth value. */
+	token = auth + strlen(BEARER_SCHEME);
+
+	/* Swallow any additional spaces. */
+	while (*token == ' ')
+		token++;
+
+	/*
+	 * Before invoking the validator command, sanity-check the token format to
+	 * avoid any injection attacks later in the chain. Invalid formats are
+	 * technically a protocol violation, but don't reflect any information about
+	 * the sensitive Bearer token back to the client; log at COMMERROR instead.
+	 */
+
+	/* Tokens must not be empty. */
+	if (!*token)
+	{
+		ereport(COMMERROR,
+				(errcode(ERRCODE_PROTOCOL_VIOLATION),
+				 errmsg("malformed OAUTHBEARER message"),
+				 errdetail("Bearer token is empty.")));
+		return false;
+	}
+
+	/*
+	 * Make sure the token contains only allowed characters. Tokens may end with
+	 * any number of '=' characters.
+	 */
+	span = strspn(token, b64_set);
+	while (token[span] == '=')
+		span++;
+
+	if (token[span] != '\0')
+	{
+		/*
+		 * This error message could be more helpful by printing the problematic
+		 * character(s), but that'd be a bit like printing a piece of someone's
+		 * password into the logs.
+		 */
+		ereport(COMMERROR,
+				(errcode(ERRCODE_PROTOCOL_VIOLATION),
+				 errmsg("malformed OAUTHBEARER message"),
+				 errdetail("Bearer token is not in the correct format.")));
+		return false;
+	}
+
+	/* Have the validator check the token. */
+	if (!run_validator_command(port, token))
+		return false;
+
+	if (port->hba->oauth_skip_usermap)
+	{
+		/*
+		 * If the validator is our authorization authority, we're done.
+		 * Authentication may or may not have been performed depending on the
+		 * validator implementation; all that matters is that the validator says
+		 * the user can log in with the target role.
+		 */
+		return true;
+	}
+
+	/* Make sure the validator authenticated the user. */
+	if (!port->authn_id)
+	{
+		/* TODO: use logdetail; reduce message duplication */
+		ereport(LOG,
+				(errmsg("OAuth bearer authentication failed for user \"%s\": validator provided no identity",
+						port->user_name)));
+		return false;
+	}
+
+	/* Finally, check the user map. */
+	ret = check_usermap(port->hba->usermap, port->user_name, port->authn_id,
+						false);
+	return (ret == STATUS_OK);
+}
+
+static bool
+run_validator_command(Port *port, const char *token)
+{
+	bool		success = false;
+	int			rc;
+	int			pipefd[2];
+	int			rfd = -1;
+	int			wfd = -1;
+
+	StringInfoData command = { 0 };
+	char	   *p;
+	FILE	   *fh = NULL;
+
+	ssize_t		written;
+	char	   *line = NULL;
+	size_t		size = 0;
+	ssize_t		len;
+
+	Assert(oauth_validator_command);
+
+	if (!oauth_validator_command[0])
+	{
+		ereport(COMMERROR,
+				(errmsg("oauth_validator_command is not set"),
+				 errhint("To allow OAuth authenticated connections, set "
+						 "oauth_validator_command in postgresql.conf.")));
+		return false;
+	}
+
+	/*
+	 * Since popen() is unidirectional, open up a pipe for the other direction.
+	 * Use CLOEXEC to ensure that our write end doesn't accidentally get copied
+	 * into child processes, which would prevent us from closing it cleanly.
+	 *
+	 * XXX this is ugly. We should just read from the child process's stdout,
+	 * but that's a lot more code.
+	 * XXX by bypassing the popen API, we open the potential of process
+	 * deadlock. Clearly document child process requirements (i.e. the child
+	 * MUST read all data off of the pipe before writing anything).
+	 * TODO: port to Windows using _pipe().
+	 */
+	rc = pipe2(pipefd, O_CLOEXEC);
+	if (rc < 0)
+	{
+		ereport(COMMERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not create child pipe: %m")));
+		return false;
+	}
+
+	rfd = pipefd[0];
+	wfd = pipefd[1];
+
+	/* Allow the read pipe be passed to the child. */
+	if (!unset_cloexec(rfd))
+	{
+		/* error message was already logged */
+		goto cleanup;
+	}
+
+	/*
+	 * Construct the command, substituting any recognized %-specifiers:
+	 *
+	 *   %f: the file descriptor of the input pipe
+	 *   %r: the role that the client wants to assume (port->user_name)
+	 *   %%: a literal '%'
+	 */
+	initStringInfo(&command);
+
+	for (p = oauth_validator_command; *p; p++)
+	{
+		if (p[0] == '%')
+		{
+			switch (p[1])
+			{
+				case 'f':
+					appendStringInfo(&command, "%d", rfd);
+					p++;
+					break;
+				case 'r':
+					/*
+					 * TODO: decide how this string should be escaped. The role
+					 * is controlled by the client, so if we don't escape it,
+					 * command injections are inevitable.
+					 *
+					 * This is probably an indication that the role name needs
+					 * to be communicated to the validator process in some other
+					 * way. For this proof of concept, just be incredibly strict
+					 * about the characters that are allowed in user names.
+					 */
+					if (!username_ok_for_shell(port->user_name))
+						goto cleanup;
+
+					appendStringInfoString(&command, port->user_name);
+					p++;
+					break;
+				case '%':
+					appendStringInfoChar(&command, '%');
+					p++;
+					break;
+				default:
+					appendStringInfoChar(&command, p[0]);
+			}
+		}
+		else
+			appendStringInfoChar(&command, p[0]);
+	}
+
+	/* Execute the command. */
+	fh = OpenPipeStream(command.data, "re");
+	/* TODO: handle failures */
+
+	/* We don't need the read end of the pipe anymore. */
+	close(rfd);
+	rfd = -1;
+
+	/* Give the command the token to validate. */
+	written = write(wfd, token, strlen(token));
+	if (written != strlen(token))
+	{
+		/* TODO must loop for short writes, EINTR et al */
+		ereport(COMMERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not write token to child pipe: %m")));
+		goto cleanup;
+	}
+
+	close(wfd);
+	wfd = -1;
+
+	/*
+	 * Read the command's response.
+	 *
+	 * TODO: getline() is probably too new to use, unfortunately.
+	 * TODO: loop over all lines
+	 */
+	if ((len = getline(&line, &size, fh)) >= 0)
+	{
+		/* TODO: fail if the authn_id doesn't end with a newline */
+		if (len > 0)
+			line[len - 1] = '\0';
+
+		set_authn_id(port, line);
+	}
+	else if (ferror(fh))
+	{
+		ereport(COMMERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not read from command \"%s\": %m",
+						command.data)));
+		goto cleanup;
+	}
+
+	/* Make sure the command exits cleanly. */
+	if (!check_exit(&fh, command.data))
+	{
+		/* error message already logged */
+		goto cleanup;
+	}
+
+	/* Done. */
+	success = true;
+
+cleanup:
+	if (line)
+		free(line);
+
+	/*
+	 * In the successful case, the pipe fds are already closed. For the error
+	 * case, always close out the pipe before waiting for the command, to
+	 * prevent deadlock.
+	 */
+	if (rfd >= 0)
+		close(rfd);
+	if (wfd >= 0)
+		close(wfd);
+
+	if (fh)
+	{
+		Assert(!success);
+		check_exit(&fh, command.data);
+	}
+
+	if (command.data)
+		pfree(command.data);
+
+	return success;
+}
+
+static bool
+check_exit(FILE **fh, const char *command)
+{
+	int rc;
+
+	rc = ClosePipeStream(*fh);
+	*fh = NULL;
+
+	if (rc == -1)
+	{
+		/* pclose() itself failed. */
+		ereport(COMMERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not close pipe to command \"%s\": %m",
+						command)));
+	}
+	else if (rc != 0)
+	{
+		char *reason = wait_result_to_str(rc);
+
+		ereport(COMMERROR,
+				(errmsg("failed to execute command \"%s\": %s",
+						command, reason)));
+
+		pfree(reason);
+	}
+
+	return (rc == 0);
+}
+
+static bool
+unset_cloexec(int fd)
+{
+	int			flags;
+	int			rc;
+
+	flags = fcntl(fd, F_GETFD);
+	if (flags == -1)
+	{
+		ereport(COMMERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not get fd flags for child pipe: %m")));
+		return false;
+	}
+
+	rc = fcntl(fd, F_SETFD, flags & ~FD_CLOEXEC);
+	if (rc < 0)
+	{
+		ereport(COMMERROR,
+				(errcode_for_file_access(),
+				 errmsg("could not unset FD_CLOEXEC for child pipe: %m")));
+		return false;
+	}
+
+	return true;
+}
+
+/*
+ * XXX This should go away eventually and be replaced with either a proper
+ * escape or a different strategy for communication with the validator command.
+ */
+static bool
+username_ok_for_shell(const char *username)
+{
+	/* This set is borrowed from fe_utils' appendShellStringNoError(). */
+	static const char * const allowed = "abcdefghijklmnopqrstuvwxyz"
+										"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+										"0123456789-_./:";
+	size_t	span;
+
+	Assert(username && username[0]); /* should have already been checked */
+
+	span = strspn(username, allowed);
+	if (username[span] != '\0')
+	{
+		ereport(COMMERROR,
+				(errmsg("PostgreSQL user name contains unsafe characters and cannot be passed to the OAuth validator")));
+		return false;
+	}
+
+	return true;
+}
diff --git a/src/backend/libpq/auth-sasl.c b/src/backend/libpq/auth-sasl.c
index 6cfd90fa21..f6c49a4de5 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 9df8f17837..5bb0388c01 100644
--- a/src/backend/libpq/auth-scram.c
+++ b/src/backend/libpq/auth-scram.c
@@ -117,7 +117,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/backend/libpq/auth.c b/src/backend/libpq/auth.c
index 8cc23ef7fb..fbcc2c55b4 100644
--- a/src/backend/libpq/auth.c
+++ b/src/backend/libpq/auth.c
@@ -29,6 +29,7 @@
 #include "libpq/auth.h"
 #include "libpq/crypt.h"
 #include "libpq/libpq.h"
+#include "libpq/oauth.h"
 #include "libpq/pqformat.h"
 #include "libpq/sasl.h"
 #include "libpq/scram.h"
@@ -47,7 +48,6 @@
  */
 static void auth_failed(Port *port, int status, char *logdetail);
 static char *recv_password_packet(Port *port);
-static void set_authn_id(Port *port, const char *id);
 
 
 /*----------------------------------------------------------------
@@ -205,22 +205,6 @@ static int	CheckRADIUSAuth(Port *port);
 static int	PerformRadiusTransaction(const char *server, const char *secret, const char *portstr, const char *identifier, const char *user_name, const char *passwd);
 
 
-/*
- * Maximum accepted size of GSS and SSPI authentication tokens.
- * We also use this as a limit on ordinary password packet lengths.
- *
- * Kerberos tickets are usually quite small, but the TGTs issued by Windows
- * domain controllers include an authorization field known as the Privilege
- * Attribute Certificate (PAC), which contains the user's Windows permissions
- * (group memberships etc.). The PAC is copied into all tickets obtained on
- * the basis of this TGT (even those issued by Unix realms which the Windows
- * realm trusts), and can be several kB in size. The maximum token size
- * accepted by Windows systems is determined by the MaxAuthToken Windows
- * registry setting. Microsoft recommends that it is not set higher than
- * 65535 bytes, so that seems like a reasonable limit for us as well.
- */
-#define PG_MAX_AUTH_TOKEN_LENGTH	65535
-
 /*----------------------------------------------------------------
  * Global authentication functions
  *----------------------------------------------------------------
@@ -309,6 +293,9 @@ auth_failed(Port *port, int status, char *logdetail)
 		case uaRADIUS:
 			errstr = gettext_noop("RADIUS authentication failed for user \"%s\"");
 			break;
+		case uaOAuth:
+			errstr = gettext_noop("OAuth bearer authentication failed for user \"%s\"");
+			break;
 		default:
 			errstr = gettext_noop("authentication failed for user \"%s\": invalid authentication method");
 			break;
@@ -343,7 +330,7 @@ auth_failed(Port *port, int status, char *logdetail)
  * lifetime of the Port, so it is safe to pass a string that is managed by an
  * external library.
  */
-static void
+void
 set_authn_id(Port *port, const char *id)
 {
 	Assert(id);
@@ -628,6 +615,9 @@ ClientAuthentication(Port *port)
 		case uaTrust:
 			status = STATUS_OK;
 			break;
+		case uaOAuth:
+			status = CheckSASLAuth(&pg_be_oauth_mech, port, NULL, NULL);
+			break;
 	}
 
 	if ((status == STATUS_OK && port->hba->clientcert == clientCertFull)
diff --git a/src/backend/libpq/hba.c b/src/backend/libpq/hba.c
index 3be8778d21..98147700dd 100644
--- a/src/backend/libpq/hba.c
+++ b/src/backend/libpq/hba.c
@@ -134,7 +134,8 @@ static const char *const UserAuthName[] =
 	"ldap",
 	"cert",
 	"radius",
-	"peer"
+	"peer",
+	"oauth",
 };
 
 
@@ -1399,6 +1400,8 @@ parse_hba_line(TokenizedLine *tok_line, int elevel)
 #endif
 	else if (strcmp(token->string, "radius") == 0)
 		parsedline->auth_method = uaRADIUS;
+	else if (strcmp(token->string, "oauth") == 0)
+		parsedline->auth_method = uaOAuth;
 	else
 	{
 		ereport(elevel,
@@ -1713,8 +1716,9 @@ parse_hba_auth_opt(char *name, char *val, HbaLine *hbaline,
 			hbaline->auth_method != uaPeer &&
 			hbaline->auth_method != uaGSS &&
 			hbaline->auth_method != uaSSPI &&
+			hbaline->auth_method != uaOAuth &&
 			hbaline->auth_method != uaCert)
-			INVALID_AUTH_OPTION("map", gettext_noop("ident, peer, gssapi, sspi, and cert"));
+			INVALID_AUTH_OPTION("map", gettext_noop("ident, peer, gssapi, sspi, oauth, and cert"));
 		hbaline->usermap = pstrdup(val);
 	}
 	else if (strcmp(name, "clientcert") == 0)
@@ -2098,6 +2102,27 @@ parse_hba_auth_opt(char *name, char *val, HbaLine *hbaline,
 		hbaline->radiusidentifiers = parsed_identifiers;
 		hbaline->radiusidentifiers_s = pstrdup(val);
 	}
+	else if (strcmp(name, "issuer") == 0)
+	{
+		if (hbaline->auth_method != uaOAuth)
+			INVALID_AUTH_OPTION("issuer", gettext_noop("oauth"));
+		hbaline->oauth_issuer = pstrdup(val);
+	}
+	else if (strcmp(name, "scope") == 0)
+	{
+		if (hbaline->auth_method != uaOAuth)
+			INVALID_AUTH_OPTION("scope", gettext_noop("oauth"));
+		hbaline->oauth_scope = pstrdup(val);
+	}
+	else if (strcmp(name, "trust_validator_authz") == 0)
+	{
+		if (hbaline->auth_method != uaOAuth)
+			INVALID_AUTH_OPTION("trust_validator_authz", gettext_noop("oauth"));
+		if (strcmp(val, "1") == 0)
+			hbaline->oauth_skip_usermap = true;
+		else
+			hbaline->oauth_skip_usermap = false;
+	}
 	else
 	{
 		ereport(elevel,
diff --git a/src/backend/utils/misc/guc.c b/src/backend/utils/misc/guc.c
index 467b0fd6fe..2b42862f71 100644
--- a/src/backend/utils/misc/guc.c
+++ b/src/backend/utils/misc/guc.c
@@ -56,6 +56,7 @@
 #include "libpq/auth.h"
 #include "libpq/libpq.h"
 #include "libpq/pqformat.h"
+#include "libpq/oauth.h"
 #include "miscadmin.h"
 #include "optimizer/cost.h"
 #include "optimizer/geqo.h"
@@ -4594,6 +4595,17 @@ static struct config_string ConfigureNamesString[] =
 		check_backtrace_functions, assign_backtrace_functions, NULL
 	},
 
+	{
+		{"oauth_validator_command", PGC_SIGHUP, CONN_AUTH_AUTH,
+			gettext_noop("Command to validate OAuth v2 bearer tokens."),
+			NULL,
+			GUC_SUPERUSER_ONLY
+		},
+		&oauth_validator_command,
+		"",
+		NULL, NULL, NULL
+	},
+
 	/* End-of-list marker */
 	{
 		{NULL, 0, 0, NULL, NULL}, NULL, NULL, NULL, NULL, NULL
diff --git a/src/include/libpq/auth.h b/src/include/libpq/auth.h
index 3d6734f253..1c77dcb0c1 100644
--- a/src/include/libpq/auth.h
+++ b/src/include/libpq/auth.h
@@ -16,6 +16,22 @@
 
 #include "libpq/libpq-be.h"
 
+/*
+ * Maximum accepted size of GSS and SSPI authentication tokens.
+ * We also use this as a limit on ordinary password packet lengths.
+ *
+ * Kerberos tickets are usually quite small, but the TGTs issued by Windows
+ * domain controllers include an authorization field known as the Privilege
+ * Attribute Certificate (PAC), which contains the user's Windows permissions
+ * (group memberships etc.). The PAC is copied into all tickets obtained on
+ * the basis of this TGT (even those issued by Unix realms which the Windows
+ * realm trusts), and can be several kB in size. The maximum token size
+ * accepted by Windows systems is determined by the MaxAuthToken Windows
+ * registry setting. Microsoft recommends that it is not set higher than
+ * 65535 bytes, so that seems like a reasonable limit for us as well.
+ */
+#define PG_MAX_AUTH_TOKEN_LENGTH	65535
+
 extern char *pg_krb_server_keyfile;
 extern bool pg_krb_caseins_users;
 extern char *pg_krb_realm;
@@ -23,6 +39,7 @@ extern char *pg_krb_realm;
 extern void ClientAuthentication(Port *port);
 extern void sendAuthRequest(Port *port, AuthRequest areq, const char *extradata,
 							int extralen);
+extern void set_authn_id(Port *port, const char *id);
 
 /* Hook for plugins to get control in ClientAuthentication() */
 typedef void (*ClientAuthentication_hook_type) (Port *, int);
diff --git a/src/include/libpq/hba.h b/src/include/libpq/hba.h
index 8d9f3821b1..441dd5623e 100644
--- a/src/include/libpq/hba.h
+++ b/src/include/libpq/hba.h
@@ -38,8 +38,9 @@ typedef enum UserAuth
 	uaLDAP,
 	uaCert,
 	uaRADIUS,
-	uaPeer
-#define USER_AUTH_LAST uaPeer	/* Must be last value of this enum */
+	uaPeer,
+	uaOAuth
+#define USER_AUTH_LAST uaOAuth	/* Must be last value of this enum */
 } UserAuth;
 
 /*
@@ -120,6 +121,9 @@ typedef struct HbaLine
 	char	   *radiusidentifiers_s;
 	List	   *radiusports;
 	char	   *radiusports_s;
+	char	   *oauth_issuer;
+	char	   *oauth_scope;
+	bool		oauth_skip_usermap;
 } HbaLine;
 
 typedef struct IdentLine
diff --git a/src/include/libpq/oauth.h b/src/include/libpq/oauth.h
new file mode 100644
index 0000000000..870e426af1
--- /dev/null
+++ b/src/include/libpq/oauth.h
@@ -0,0 +1,24 @@
+/*-------------------------------------------------------------------------
+ *
+ * oauth.h
+ *	  Interface to libpq/auth-oauth.c
+ *
+ * Portions Copyright (c) 1996-2021, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ * src/include/libpq/oauth.h
+ *
+ *-------------------------------------------------------------------------
+ */
+#ifndef PG_OAUTH_H
+#define PG_OAUTH_H
+
+#include "libpq/libpq-be.h"
+#include "libpq/sasl.h"
+
+extern char *oauth_validator_command;
+
+/* Implementation */
+extern const pg_be_sasl_mech pg_be_oauth_mech;
+
+#endif /* PG_OAUTH_H */
diff --git a/src/include/libpq/sasl.h b/src/include/libpq/sasl.h
index 4c611bab6b..c0a88430d5 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,
 							 char **logdetail);
+
+	/* The maximum size allowed for client SASLResponses. */
+	int			max_message_length;
 } pg_be_sasl_mech;
 
 /* Common implementation for auth.c */
-- 
2.25.1



  [text/x-patch] v2-0004-Add-a-very-simple-authn_id-extension.patch (2.8K, ../[email protected]/5-v2-0004-Add-a-very-simple-authn_id-extension.patch)
  download | inline diff:
From 7c4175f9ad87141d40dd44d6c9fe9312ce8e5b88 Mon Sep 17 00:00:00 2001
From: Jacob Champion <[email protected]>
Date: Tue, 18 May 2021 15:01:29 -0700
Subject: [PATCH v2 4/5] Add a very simple authn_id extension

...for retrieving the authn_id from the server in tests.
---
 contrib/authn_id/Makefile          | 19 +++++++++++++++++++
 contrib/authn_id/authn_id--1.0.sql |  8 ++++++++
 contrib/authn_id/authn_id.c        | 28 ++++++++++++++++++++++++++++
 contrib/authn_id/authn_id.control  |  5 +++++
 4 files changed, 60 insertions(+)
 create mode 100644 contrib/authn_id/Makefile
 create mode 100644 contrib/authn_id/authn_id--1.0.sql
 create mode 100644 contrib/authn_id/authn_id.c
 create mode 100644 contrib/authn_id/authn_id.control

diff --git a/contrib/authn_id/Makefile b/contrib/authn_id/Makefile
new file mode 100644
index 0000000000..46026358e0
--- /dev/null
+++ b/contrib/authn_id/Makefile
@@ -0,0 +1,19 @@
+# contrib/authn_id/Makefile
+
+MODULE_big = authn_id
+OBJS = authn_id.o
+
+EXTENSION = authn_id
+DATA = authn_id--1.0.sql
+PGFILEDESC = "authn_id - information about the authenticated user"
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = contrib/authn_id
+top_builddir = ../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/contrib/authn_id/authn_id--1.0.sql b/contrib/authn_id/authn_id--1.0.sql
new file mode 100644
index 0000000000..af2a4d3991
--- /dev/null
+++ b/contrib/authn_id/authn_id--1.0.sql
@@ -0,0 +1,8 @@
+/* contrib/authn_id/authn_id--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION authn_id" to load this file. \quit
+
+CREATE FUNCTION authn_id() RETURNS text
+AS 'MODULE_PATHNAME', 'authn_id'
+LANGUAGE C IMMUTABLE;
diff --git a/contrib/authn_id/authn_id.c b/contrib/authn_id/authn_id.c
new file mode 100644
index 0000000000..0fecac36a8
--- /dev/null
+++ b/contrib/authn_id/authn_id.c
@@ -0,0 +1,28 @@
+/*
+ * Extension to expose the current user's authn_id.
+ *
+ * contrib/authn_id/authn_id.c
+ */
+
+#include "postgres.h"
+
+#include "fmgr.h"
+#include "libpq/libpq-be.h"
+#include "miscadmin.h"
+#include "utils/builtins.h"
+
+PG_MODULE_MAGIC;
+
+PG_FUNCTION_INFO_V1(authn_id);
+
+/*
+ * Returns the current user's authenticated identity.
+ */
+Datum
+authn_id(PG_FUNCTION_ARGS)
+{
+	if (!MyProcPort->authn_id)
+		PG_RETURN_NULL();
+
+	PG_RETURN_TEXT_P(cstring_to_text(MyProcPort->authn_id));
+}
diff --git a/contrib/authn_id/authn_id.control b/contrib/authn_id/authn_id.control
new file mode 100644
index 0000000000..e0f9e06bed
--- /dev/null
+++ b/contrib/authn_id/authn_id.control
@@ -0,0 +1,5 @@
+# authn_id extension
+comment = 'current user identity'
+default_version = '1.0'
+module_pathname = '$libdir/authn_id'
+relocatable = true
-- 
2.25.1



  [text/x-patch] v2-0005-Add-pytest-suite-for-OAuth.patch (131.6K, ../[email protected]/6-v2-0005-Add-pytest-suite-for-OAuth.patch)
  download | inline diff:
From 0281635f35a44e0fdfd4369423f98ebe5b467ce3 Mon Sep 17 00:00:00 2001
From: Jacob Champion <[email protected]>
Date: Fri, 4 Jun 2021 09:06:38 -0700
Subject: [PATCH v2 5/5] Add pytest suite for OAuth

Requires Python 3; on the first run of `make installcheck` the
dependencies will be installed into ./venv for you. See the README for
more details.
---
 src/test/python/.gitignore                 |    2 +
 src/test/python/Makefile                   |   38 +
 src/test/python/README                     |   54 ++
 src/test/python/client/__init__.py         |    0
 src/test/python/client/conftest.py         |  126 +++
 src/test/python/client/test_client.py      |  180 ++++
 src/test/python/client/test_oauth.py       |  936 ++++++++++++++++++
 src/test/python/pq3.py                     |  727 ++++++++++++++
 src/test/python/pytest.ini                 |    4 +
 src/test/python/requirements.txt           |    7 +
 src/test/python/server/__init__.py         |    0
 src/test/python/server/conftest.py         |   45 +
 src/test/python/server/test_oauth.py       | 1012 ++++++++++++++++++++
 src/test/python/server/test_server.py      |   21 +
 src/test/python/server/validate_bearer.py  |  101 ++
 src/test/python/server/validate_reflect.py |   34 +
 src/test/python/test_internals.py          |  138 +++
 src/test/python/test_pq3.py                |  558 +++++++++++
 src/test/python/tls.py                     |  195 ++++
 19 files changed, 4178 insertions(+)
 create mode 100644 src/test/python/.gitignore
 create mode 100644 src/test/python/Makefile
 create mode 100644 src/test/python/README
 create mode 100644 src/test/python/client/__init__.py
 create mode 100644 src/test/python/client/conftest.py
 create mode 100644 src/test/python/client/test_client.py
 create mode 100644 src/test/python/client/test_oauth.py
 create mode 100644 src/test/python/pq3.py
 create mode 100644 src/test/python/pytest.ini
 create mode 100644 src/test/python/requirements.txt
 create mode 100644 src/test/python/server/__init__.py
 create mode 100644 src/test/python/server/conftest.py
 create mode 100644 src/test/python/server/test_oauth.py
 create mode 100644 src/test/python/server/test_server.py
 create mode 100755 src/test/python/server/validate_bearer.py
 create mode 100755 src/test/python/server/validate_reflect.py
 create mode 100644 src/test/python/test_internals.py
 create mode 100644 src/test/python/test_pq3.py
 create mode 100644 src/test/python/tls.py

diff --git a/src/test/python/.gitignore b/src/test/python/.gitignore
new file mode 100644
index 0000000000..0e8f027b2e
--- /dev/null
+++ b/src/test/python/.gitignore
@@ -0,0 +1,2 @@
+__pycache__/
+/venv/
diff --git a/src/test/python/Makefile b/src/test/python/Makefile
new file mode 100644
index 0000000000..b0695b6287
--- /dev/null
+++ b/src/test/python/Makefile
@@ -0,0 +1,38 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+# Only Python 3 is supported, but if it's named something different on your
+# system you can override it with the PYTHON3 variable.
+PYTHON3 := python3
+
+# All dependencies are placed into this directory. The default is .gitignored
+# for you, but you can override it if you'd like.
+VENV := ./venv
+
+override VBIN   := $(VENV)/bin
+override PIP    := $(VBIN)/pip
+override PYTEST := $(VBIN)/py.test
+override ISORT  := $(VBIN)/isort
+override BLACK  := $(VBIN)/black
+
+.PHONY: installcheck indent
+
+installcheck: $(PYTEST)
+	$(PYTEST) -v -rs
+
+indent: $(ISORT) $(BLACK)
+	$(ISORT) --profile black *.py client/*.py server/*.py
+	$(BLACK) *.py client/*.py server/*.py
+
+$(PYTEST) $(ISORT) $(BLACK) &: requirements.txt | $(PIP)
+	$(PIP) install --force-reinstall -r $<
+
+$(PIP):
+	$(PYTHON3) -m venv $(VENV)
+
+# A convenience recipe to rebuild psycopg2 against the local libpq.
+.PHONY: rebuild-psycopg2
+rebuild-psycopg2: | $(PIP)
+	$(PIP) install --force-reinstall --no-binary :all: $(shell grep psycopg2 requirements.txt)
diff --git a/src/test/python/README b/src/test/python/README
new file mode 100644
index 0000000000..0bda582c4b
--- /dev/null
+++ b/src/test/python/README
@@ -0,0 +1,54 @@
+A test suite for exercising both the libpq client and the server backend at the
+protocol level, based on pytest and Construct.
+
+The test suite currently assumes that the standard PG* environment variables
+point to the database under test and are sufficient to log in a superuser on
+that system. In other words, a bare `psql` needs to Just Work before the test
+suite can do its thing. For a newly built dev cluster, typically all that I need
+to do is a
+
+    export PGDATABASE=postgres
+
+but you can adjust as needed for your setup.
+
+## Requirements
+
+A supported version (3.6+) of Python.
+
+The first run of
+
+    make installcheck
+
+will install a local virtual environment and all needed dependencies. During
+development, if libpq changes incompatibly, you can issue
+
+    $ make rebuild-psycopg2
+
+to force a rebuild of the client library.
+
+## Hacking
+
+The code style is enforced by a _very_ opinionated autoformatter. Running the
+
+    make indent
+
+recipe will invoke it for you automatically. Don't fight the tool; part of the
+zen is in knowing that if the formatter makes your code ugly, there's probably a
+cleaner way to write your code.
+
+## Advanced Usage
+
+The Makefile is there for convenience, but you don't have to use it. Activate
+the virtualenv to be able to use pytest directly:
+
+    $ source venv/bin/activate
+    $ py.test -k oauth
+    ...
+    $ py.test ./server/test_server.py
+    ...
+    $ deactivate  # puts the PATH et al back the way it was before
+
+To make quick smoke tests possible, slow tests have been marked explicitly. You
+can skip them by saying e.g.
+
+    $ py.test -m 'not slow'
diff --git a/src/test/python/client/__init__.py b/src/test/python/client/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/test/python/client/conftest.py b/src/test/python/client/conftest.py
new file mode 100644
index 0000000000..f38da7a138
--- /dev/null
+++ b/src/test/python/client/conftest.py
@@ -0,0 +1,126 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import socket
+import sys
+import threading
+
+import psycopg2
+import pytest
+
+import pq3
+
+BLOCKING_TIMEOUT = 2  # the number of seconds to wait for blocking calls
+
+
[email protected]
+def server_socket(unused_tcp_port_factory):
+    """
+    Returns a listening socket bound to an ephemeral port.
+    """
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", unused_tcp_port_factory()))
+        s.listen(1)
+        s.settimeout(BLOCKING_TIMEOUT)
+        yield s
+
+
+class ClientHandshake(threading.Thread):
+    """
+    A thread that connects to a local Postgres server using psycopg2. Once the
+    opening handshake completes, the connection will be immediately closed.
+    """
+
+    def __init__(self, *, port, **kwargs):
+        super().__init__()
+
+        kwargs["port"] = port
+        self._kwargs = kwargs
+
+        self.exception = None
+
+    def run(self):
+        try:
+            conn = psycopg2.connect(host="127.0.0.1", **self._kwargs)
+            conn.close()
+        except Exception as e:
+            self.exception = e
+
+    def check_completed(self, timeout=BLOCKING_TIMEOUT):
+        """
+        Joins the client thread. Raises an exception if the thread could not be
+        joined, or if it threw an exception itself. (The exception will be
+        cleared, so future calls to check_completed will succeed.)
+        """
+        self.join(timeout)
+
+        if self.is_alive():
+            raise TimeoutError("client thread did not handshake within the timeout")
+        elif self.exception:
+            e = self.exception
+            self.exception = None
+            raise e
+
+
[email protected]
+def accept(server_socket):
+    """
+    Returns a factory function that, when called, returns a pair (sock, client)
+    where sock is a server socket that has accepted a connection from client,
+    and client is an instance of ClientHandshake. Clients will complete their
+    handshakes and cleanly disconnect.
+
+    The default connstring options may be extended or overridden by passing
+    arbitrary keyword arguments. Keep in mind that you generally should not
+    override the host or port, since they point to the local test server.
+
+    For situations where a client needs to connect more than once to complete a
+    handshake, the accept function may be called more than once. (The client
+    returned for subsequent calls will always be the same client that was
+    returned for the first call.)
+
+    Tests must either complete the handshake so that the client thread can be
+    automatically joined during teardown, or else call client.check_completed()
+    and manually handle any expected errors.
+    """
+    _, port = server_socket.getsockname()
+
+    client = None
+    default_opts = dict(
+        port=port,
+        user=pq3.pguser(),
+        sslmode="disable",
+    )
+
+    def factory(**kwargs):
+        nonlocal client
+
+        if client is None:
+            opts = dict(default_opts)
+            opts.update(kwargs)
+
+            # The server_socket is already listening, so the client thread can
+            # be safely started; it'll block on the connection until we accept.
+            client = ClientHandshake(**opts)
+            client.start()
+
+        sock, _ = server_socket.accept()
+        return sock, client
+
+    yield factory
+    client.check_completed()
+
+
[email protected]
+def conn(accept):
+    """
+    Returns an accepted, wrapped pq3 connection to a psycopg2 client. The socket
+    will be closed when the test finishes, and the client will be checked for a
+    cleanly completed handshake.
+    """
+    sock, client = accept()
+    with sock:
+        with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+            yield conn
diff --git a/src/test/python/client/test_client.py b/src/test/python/client/test_client.py
new file mode 100644
index 0000000000..c4c946fda4
--- /dev/null
+++ b/src/test/python/client/test_client.py
@@ -0,0 +1,180 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import base64
+import sys
+
+import psycopg2
+import pytest
+from cryptography.hazmat.primitives import hashes, hmac
+
+import pq3
+
+
+def finish_handshake(conn):
+    """
+    Sends the AuthenticationOK message and the standard opening salvo of server
+    messages, then asserts that the client immediately sends a Terminate message
+    to close the connection cleanly.
+    """
+    pq3.send(conn, pq3.types.AuthnRequest, type=pq3.authn.OK)
+    pq3.send(conn, pq3.types.ParameterStatus, name=b"client_encoding", value=b"UTF-8")
+    pq3.send(conn, pq3.types.ParameterStatus, name=b"DateStyle", value=b"ISO, MDY")
+    pq3.send(conn, pq3.types.ReadyForQuery, status=b"I")
+
+    pkt = pq3.recv1(conn)
+    assert pkt.type == pq3.types.Terminate
+
+
+def test_handshake(conn):
+    startup = pq3.recv1(conn, cls=pq3.Startup)
+    assert startup.proto == pq3.protocol(3, 0)
+
+    finish_handshake(conn)
+
+
+def test_aborted_connection(accept):
+    """
+    Make sure the client correctly reports an early close during handshakes.
+    """
+    sock, client = accept()
+    sock.close()
+
+    expected = "server closed the connection unexpectedly"
+    with pytest.raises(psycopg2.OperationalError, match=expected):
+        client.check_completed()
+
+
+#
+# SCRAM-SHA-256 (see RFC 5802: https://tools.ietf.org/html/rfc5802)
+#
+
+
[email protected]
+def password():
+    """
+    Returns a password for use by both client and server.
+    """
+    # TODO: parameterize this with passwords that require SASLprep.
+    return "secret"
+
+
[email protected]
+def pwconn(accept, password):
+    """
+    Like the conn fixture, but uses a password in the connection.
+    """
+    sock, client = accept(password=password)
+    with sock:
+        with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+            yield conn
+
+
+def sha256(data):
+    """The H(str) function from Section 2.2."""
+    digest = hashes.Hash(hashes.SHA256())
+    digest.update(data)
+    return digest.finalize()
+
+
+def hmac_256(key, data):
+    """The HMAC(key, str) function from Section 2.2."""
+    h = hmac.HMAC(key, hashes.SHA256())
+    h.update(data)
+    return h.finalize()
+
+
+def xor(a, b):
+    """The XOR operation from Section 2.2."""
+    res = bytearray(a)
+    for i, byte in enumerate(b):
+        res[i] ^= byte
+    return bytes(res)
+
+
+def h_i(data, salt, i):
+    """The Hi(str, salt, i) function from Section 2.2."""
+    assert i > 0
+
+    acc = hmac_256(data, salt + b"\x00\x00\x00\x01")
+    last = acc
+    i -= 1
+
+    while i:
+        u = hmac_256(data, last)
+        acc = xor(acc, u)
+
+        last = u
+        i -= 1
+
+    return acc
+
+
+def test_scram(pwconn, password):
+    startup = pq3.recv1(pwconn, cls=pq3.Startup)
+    assert startup.proto == pq3.protocol(3, 0)
+
+    pq3.send(
+        pwconn,
+        pq3.types.AuthnRequest,
+        type=pq3.authn.SASL,
+        body=[b"SCRAM-SHA-256", b""],
+    )
+
+    # Get the client-first-message.
+    pkt = pq3.recv1(pwconn)
+    assert pkt.type == pq3.types.PasswordMessage
+
+    initial = pq3.SASLInitialResponse.parse(pkt.payload)
+    assert initial.name == b"SCRAM-SHA-256"
+
+    c_bind, authzid, c_name, c_nonce = initial.data.split(b",")
+    assert c_bind == b"n"  # no channel bindings on a plaintext connection
+    assert authzid == b""  # we don't support authzid currently
+    assert c_name == b"n="  # libpq doesn't honor the GS2 username
+    assert c_nonce.startswith(b"r=")
+
+    # Send the server-first-message.
+    salt = b"12345"
+    iterations = 2
+
+    s_nonce = c_nonce + b"somenonce"
+    s_salt = b"s=" + base64.b64encode(salt)
+    s_iterations = b"i=%d" % iterations
+
+    msg = b",".join([s_nonce, s_salt, s_iterations])
+    pq3.send(pwconn, pq3.types.AuthnRequest, type=pq3.authn.SASLContinue, body=msg)
+
+    # Get the client-final-message.
+    pkt = pq3.recv1(pwconn)
+    assert pkt.type == pq3.types.PasswordMessage
+
+    c_bind_final, c_nonce_final, c_proof = pkt.payload.split(b",")
+    assert c_bind_final == b"c=" + base64.b64encode(c_bind + b"," + authzid + b",")
+    assert c_nonce_final == s_nonce
+
+    # Calculate what the client proof should be.
+    salted_password = h_i(password.encode("ascii"), salt, iterations)
+    client_key = hmac_256(salted_password, b"Client Key")
+    stored_key = sha256(client_key)
+
+    auth_message = b",".join(
+        [c_name, c_nonce, s_nonce, s_salt, s_iterations, c_bind_final, c_nonce_final]
+    )
+    client_signature = hmac_256(stored_key, auth_message)
+    client_proof = xor(client_key, client_signature)
+
+    expected = b"p=" + base64.b64encode(client_proof)
+    assert c_proof == expected
+
+    # Send the correct server signature.
+    server_key = hmac_256(salted_password, b"Server Key")
+    server_signature = hmac_256(server_key, auth_message)
+
+    s_verify = b"v=" + base64.b64encode(server_signature)
+    pq3.send(pwconn, pq3.types.AuthnRequest, type=pq3.authn.SASLFinal, body=s_verify)
+
+    # Done!
+    finish_handshake(pwconn)
diff --git a/src/test/python/client/test_oauth.py b/src/test/python/client/test_oauth.py
new file mode 100644
index 0000000000..a754a9c0b6
--- /dev/null
+++ b/src/test/python/client/test_oauth.py
@@ -0,0 +1,936 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import base64
+import http.server
+import json
+import secrets
+import sys
+import threading
+import time
+import urllib.parse
+
+import psycopg2
+import pytest
+
+import pq3
+
+from .conftest import BLOCKING_TIMEOUT
+
+
+def finish_handshake(conn):
+    """
+    Sends the AuthenticationOK message and the standard opening salvo of server
+    messages, then asserts that the client immediately sends a Terminate message
+    to close the connection cleanly.
+    """
+    pq3.send(conn, pq3.types.AuthnRequest, type=pq3.authn.OK)
+    pq3.send(conn, pq3.types.ParameterStatus, name=b"client_encoding", value=b"UTF-8")
+    pq3.send(conn, pq3.types.ParameterStatus, name=b"DateStyle", value=b"ISO, MDY")
+    pq3.send(conn, pq3.types.ReadyForQuery, status=b"I")
+
+    pkt = pq3.recv1(conn)
+    assert pkt.type == pq3.types.Terminate
+
+
+#
+# OAUTHBEARER (see RFC 7628: https://tools.ietf.org/html/rfc7628)
+#
+
+
+def start_oauth_handshake(conn):
+    """
+    Negotiates an OAUTHBEARER SASL challenge. Returns the client's initial
+    response data.
+    """
+    startup = pq3.recv1(conn, cls=pq3.Startup)
+    assert startup.proto == pq3.protocol(3, 0)
+
+    pq3.send(
+        conn, pq3.types.AuthnRequest, type=pq3.authn.SASL, body=[b"OAUTHBEARER", b""]
+    )
+
+    pkt = pq3.recv1(conn)
+    assert pkt.type == pq3.types.PasswordMessage
+
+    initial = pq3.SASLInitialResponse.parse(pkt.payload)
+    assert initial.name == b"OAUTHBEARER"
+
+    return initial.data
+
+
+def get_auth_value(initial):
+    """
+    Finds the auth value (e.g. "Bearer somedata..." in the client's initial SASL
+    response.
+    """
+    kvpairs = initial.split(b"\x01")
+    assert kvpairs[0] == b"n,,"  # no channel binding or authzid
+    assert kvpairs[2] == b""  # ends with an empty kvpair
+    assert kvpairs[3] == b""  # ...and there's nothing after it
+    assert len(kvpairs) == 4
+
+    key, value = kvpairs[1].split(b"=", 2)
+    assert key == b"auth"
+
+    return value
+
+
+def xtest_oauth_success(conn):  # TODO
+    initial = start_oauth_handshake(conn)
+
+    auth = get_auth_value(initial)
+    assert auth.startswith(b"Bearer ")
+
+    # Accept the token. TODO actually validate
+    pq3.send(conn, pq3.types.AuthnRequest, type=pq3.authn.SASLFinal)
+    finish_handshake(conn)
+
+
+class OpenIDProvider(threading.Thread):
+    """
+    A thread that runs a mock OpenID provider server.
+    """
+
+    def __init__(self, *, port):
+        super().__init__()
+
+        self.exception = None
+
+        addr = ("", port)
+        self.server = self._Server(addr, self._Handler)
+
+        # TODO: allow HTTPS only, somehow
+        oauth = self._OAuthState()
+        oauth.host = f"localhost:{port}"
+        oauth.issuer = f"http://localhost:{port}"
+
+        # The following endpoints are required to be advertised by providers,
+        # even though our chosen client implementation does not actually make
+        # use of them.
+        oauth.register_endpoint(
+            "authorization_endpoint", "POST", "/authorize", self._authorization_handler
+        )
+        oauth.register_endpoint("jwks_uri", "GET", "/keys", self._jwks_handler)
+
+        self.server.oauth = oauth
+
+    def run(self):
+        try:
+            self.server.serve_forever()
+        except Exception as e:
+            self.exception = e
+
+    def stop(self, timeout=BLOCKING_TIMEOUT):
+        """
+        Shuts down the server and joins its thread. Raises an exception if the
+        thread could not be joined, or if it threw an exception itself. Must
+        only be called once, after start().
+        """
+        self.server.shutdown()
+        self.join(timeout)
+
+        if self.is_alive():
+            raise TimeoutError("client thread did not handshake within the timeout")
+        elif self.exception:
+            e = self.exception
+            raise e
+
+    class _OAuthState(object):
+        def __init__(self):
+            self.endpoint_paths = {}
+            self._endpoints = {}
+
+        def register_endpoint(self, name, method, path, func):
+            if method not in self._endpoints:
+                self._endpoints[method] = {}
+
+            self._endpoints[method][path] = func
+            self.endpoint_paths[name] = path
+
+        def endpoint(self, method, path):
+            if method not in self._endpoints:
+                return None
+
+            return self._endpoints[method].get(path)
+
+    class _Server(http.server.HTTPServer):
+        def handle_error(self, request, addr):
+            self.shutdown_request(request)
+            raise
+
+    @staticmethod
+    def _jwks_handler(headers, params):
+        return 200, {"keys": []}
+
+    @staticmethod
+    def _authorization_handler(headers, params):
+        # We don't actually want this to be called during these tests -- we
+        # should be using the device authorization endpoint instead.
+        assert (
+            False
+        ), "authorization handler called instead of device authorization handler"
+
+    class _Handler(http.server.BaseHTTPRequestHandler):
+        timeout = BLOCKING_TIMEOUT
+
+        def _discovery_handler(self, headers, params):
+            oauth = self.server.oauth
+
+            doc = {
+                "issuer": oauth.issuer,
+                "response_types_supported": ["token"],
+                "subject_types_supported": ["public"],
+                "id_token_signing_alg_values_supported": ["RS256"],
+            }
+
+            for name, path in oauth.endpoint_paths.items():
+                doc[name] = oauth.issuer + path
+
+            return 200, doc
+
+        def _handle(self, *, params=None, handler=None):
+            oauth = self.server.oauth
+            assert self.headers["Host"] == oauth.host
+
+            if handler is None:
+                handler = oauth.endpoint(self.command, self.path)
+                assert (
+                    handler is not None
+                ), f"no registered endpoint for {self.command} {self.path}"
+
+            code, resp = handler(self.headers, params)
+
+            self.send_response(code)
+            self.send_header("Content-Type", "application/json")
+            self.end_headers()
+
+            resp = json.dumps(resp)
+            resp = resp.encode("utf-8")
+            self.wfile.write(resp)
+
+            self.close_connection = True
+
+        def do_GET(self):
+            if self.path == "/.well-known/openid-configuration":
+                self._handle(handler=self._discovery_handler)
+                return
+
+            self._handle()
+
+        def _request_body(self):
+            length = self.headers["Content-Length"]
+
+            # Handle only an explicit content-length.
+            assert length is not None
+            length = int(length)
+
+            return self.rfile.read(length).decode("utf-8")
+
+        def do_POST(self):
+            assert self.headers["Content-Type"] == "application/x-www-form-urlencoded"
+
+            body = self._request_body()
+            params = urllib.parse.parse_qs(body)
+
+            self._handle(params=params)
+
+
[email protected]
+def openid_provider(unused_tcp_port_factory):
+    """
+    A fixture that returns the OAuth state of a running OpenID provider server. The
+    server will be stopped when the fixture is torn down.
+    """
+    thread = OpenIDProvider(port=unused_tcp_port_factory())
+    thread.start()
+
+    try:
+        yield thread.server.oauth
+    finally:
+        thread.stop()
+
+
[email protected]("secret", [None, "", "hunter2"])
[email protected]("scope", [None, "", "openid email"])
[email protected]("retries", [0, 1])
+def test_oauth_with_explicit_issuer(
+    capfd, accept, openid_provider, retries, scope, secret
+):
+    client_id = secrets.token_hex()
+
+    sock, client = accept(
+        oauth_issuer=openid_provider.issuer,
+        oauth_client_id=client_id,
+        oauth_client_secret=secret,
+        oauth_scope=scope,
+    )
+
+    device_code = secrets.token_hex()
+    user_code = f"{secrets.token_hex(2)}-{secrets.token_hex(2)}"
+    verification_url = "https://example.com/device"
+
+    access_token = secrets.token_urlsafe()
+
+    def check_client_authn(headers, params):
+        if not secret:
+            assert params["client_id"] == [client_id]
+            return
+
+        # Require the client to use Basic authn; request-body credentials are
+        # NOT RECOMMENDED (RFC 6749, Sec. 2.3.1).
+        assert "Authorization" in headers
+
+        method, creds = headers["Authorization"].split()
+        assert method == "Basic"
+
+        expected = f"{client_id}:{secret}"
+        assert base64.b64decode(creds) == expected.encode("ascii")
+
+    # Set up our provider callbacks.
+    # NOTE that these callbacks will be called on a background thread. Don't do
+    # any unprotected state mutation here.
+
+    def authorization_endpoint(headers, params):
+        check_client_authn(headers, params)
+
+        if scope:
+            assert params["scope"] == [scope]
+        else:
+            assert "scope" not in params
+
+        resp = {
+            "device_code": device_code,
+            "user_code": user_code,
+            "interval": 0,
+            "verification_uri": verification_url,
+            "expires_in": 5,
+        }
+
+        return 200, resp
+
+    openid_provider.register_endpoint(
+        "device_authorization_endpoint", "POST", "/device", authorization_endpoint
+    )
+
+    attempts = 0
+    retry_lock = threading.Lock()
+
+    def token_endpoint(headers, params):
+        check_client_authn(headers, params)
+
+        assert params["grant_type"] == ["urn:ietf:params:oauth:grant-type:device_code"]
+        assert params["device_code"] == [device_code]
+
+        now = time.monotonic()
+
+        with retry_lock:
+            nonlocal attempts
+
+            # If the test wants to force the client to retry, return an
+            # authorization_pending response and decrement the retry count.
+            if attempts < retries:
+                attempts += 1
+                return 400, {"error": "authorization_pending"}
+
+        # Successfully finish the request by sending the access bearer token.
+        resp = {
+            "access_token": access_token,
+            "token_type": "bearer",
+        }
+
+        return 200, resp
+
+    openid_provider.register_endpoint(
+        "token_endpoint", "POST", "/token", token_endpoint
+    )
+
+    with sock:
+        with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+            # Initiate a handshake, which should result in the above endpoints
+            # being called.
+            initial = start_oauth_handshake(conn)
+
+            # Validate and accept the token.
+            auth = get_auth_value(initial)
+            assert auth == f"Bearer {access_token}".encode("ascii")
+
+            pq3.send(conn, pq3.types.AuthnRequest, type=pq3.authn.SASLFinal)
+            finish_handshake(conn)
+
+    if retries:
+        # Finally, make sure that the client prompted the user with the expected
+        # authorization URL and user code.
+        expected = f"Visit {verification_url} and enter the code: {user_code}"
+        _, stderr = capfd.readouterr()
+        assert expected in stderr
+
+
+def test_oauth_requires_client_id(accept, openid_provider):
+    sock, client = accept(
+        oauth_issuer=openid_provider.issuer,
+        # Do not set a client ID; this should cause a client error after the
+        # server asks for OAUTHBEARER and the client tries to contact the
+        # issuer.
+    )
+
+    with sock:
+        with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+            # Initiate a handshake.
+            startup = pq3.recv1(conn, cls=pq3.Startup)
+            assert startup.proto == pq3.protocol(3, 0)
+
+            pq3.send(
+                conn,
+                pq3.types.AuthnRequest,
+                type=pq3.authn.SASL,
+                body=[b"OAUTHBEARER", b""],
+            )
+
+            # The client should disconnect at this point.
+            assert not conn.read()
+
+    expected_error = "no oauth_client_id is set"
+    with pytest.raises(psycopg2.OperationalError, match=expected_error):
+        client.check_completed()
+
+
[email protected]
[email protected]("error_code", ["authorization_pending", "slow_down"])
[email protected]("retries", [1, 2])
+def test_oauth_retry_interval(accept, openid_provider, retries, error_code):
+    sock, client = accept(
+        oauth_issuer=openid_provider.issuer,
+        oauth_client_id="some-id",
+    )
+
+    expected_retry_interval = 1
+    access_token = secrets.token_urlsafe()
+
+    # Set up our provider callbacks.
+    # NOTE that these callbacks will be called on a background thread. Don't do
+    # any unprotected state mutation here.
+
+    def authorization_endpoint(headers, params):
+        resp = {
+            "device_code": "my-device-code",
+            "user_code": "my-user-code",
+            "interval": expected_retry_interval,
+            "verification_uri": "https://example.com",
+            "expires_in": 5,
+        }
+
+        return 200, resp
+
+    openid_provider.register_endpoint(
+        "device_authorization_endpoint", "POST", "/device", authorization_endpoint
+    )
+
+    attempts = 0
+    last_retry = None
+    retry_lock = threading.Lock()
+
+    def token_endpoint(headers, params):
+        now = time.monotonic()
+
+        with retry_lock:
+            nonlocal attempts, last_retry, expected_retry_interval
+
+            # Make sure the retry interval is being respected by the client.
+            if last_retry is not None:
+                interval = now - last_retry
+                assert interval >= expected_retry_interval
+
+            last_retry = now
+
+            # If the test wants to force the client to retry, return the desired
+            # error response and decrement the retry count.
+            if attempts < retries:
+                attempts += 1
+
+                # A slow_down code requires the client to additionally increase
+                # its interval by five seconds.
+                if error_code == "slow_down":
+                    expected_retry_interval += 5
+
+                return 400, {"error": error_code}
+
+        # Successfully finish the request by sending the access bearer token.
+        resp = {
+            "access_token": access_token,
+            "token_type": "bearer",
+        }
+
+        return 200, resp
+
+    openid_provider.register_endpoint(
+        "token_endpoint", "POST", "/token", token_endpoint
+    )
+
+    with sock:
+        with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+            # Initiate a handshake, which should result in the above endpoints
+            # being called.
+            initial = start_oauth_handshake(conn)
+
+            # Validate and accept the token.
+            auth = get_auth_value(initial)
+            assert auth == f"Bearer {access_token}".encode("ascii")
+
+            pq3.send(conn, pq3.types.AuthnRequest, type=pq3.authn.SASLFinal)
+            finish_handshake(conn)
+
+
[email protected](
+    "failure_mode, error_pattern",
+    [
+        pytest.param(
+            {
+                "error": "invalid_client",
+                "error_description": "client authentication failed",
+            },
+            r"client authentication failed \(invalid_client\)",
+            id="authentication failure with description",
+        ),
+        pytest.param(
+            {"error": "invalid_request"},
+            r"\(invalid_request\)",
+            id="invalid request without description",
+        ),
+        pytest.param(
+            {},
+            r"failed to obtain device authorization",
+            id="broken error response",
+        ),
+    ],
+)
+def test_oauth_device_authorization_failures(
+    accept, openid_provider, failure_mode, error_pattern
+):
+    client_id = secrets.token_hex()
+
+    sock, client = accept(
+        oauth_issuer=openid_provider.issuer,
+        oauth_client_id=client_id,
+    )
+
+    # Set up our provider callbacks.
+    # NOTE that these callbacks will be called on a background thread. Don't do
+    # any unprotected state mutation here.
+
+    def authorization_endpoint(headers, params):
+        return 400, failure_mode
+
+    openid_provider.register_endpoint(
+        "device_authorization_endpoint", "POST", "/device", authorization_endpoint
+    )
+
+    def token_endpoint(headers, params):
+        assert False, "token endpoint was invoked unexpectedly"
+
+    openid_provider.register_endpoint(
+        "token_endpoint", "POST", "/token", token_endpoint
+    )
+
+    with sock:
+        with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+            # Initiate a handshake, which should result in the above endpoints
+            # being called.
+            startup = pq3.recv1(conn, cls=pq3.Startup)
+            assert startup.proto == pq3.protocol(3, 0)
+
+            pq3.send(
+                conn,
+                pq3.types.AuthnRequest,
+                type=pq3.authn.SASL,
+                body=[b"OAUTHBEARER", b""],
+            )
+
+            # The client should not continue the connection due to the hardcoded
+            # provider failure; we disconnect here.
+
+    # Now make sure the client correctly failed.
+    with pytest.raises(psycopg2.OperationalError, match=error_pattern):
+        client.check_completed()
+
+
[email protected](
+    "failure_mode, error_pattern",
+    [
+        pytest.param(
+            {
+                "error": "expired_token",
+                "error_description": "the device code has expired",
+            },
+            r"the device code has expired \(expired_token\)",
+            id="expired token with description",
+        ),
+        pytest.param(
+            {"error": "access_denied"},
+            r"\(access_denied\)",
+            id="access denied without description",
+        ),
+        pytest.param(
+            {},
+            r"OAuth token retrieval failed",
+            id="broken error response",
+        ),
+    ],
+)
[email protected]("retries", [0, 1])
+def test_oauth_token_failures(
+    accept, openid_provider, retries, failure_mode, error_pattern
+):
+    client_id = secrets.token_hex()
+
+    sock, client = accept(
+        oauth_issuer=openid_provider.issuer,
+        oauth_client_id=client_id,
+    )
+
+    device_code = secrets.token_hex()
+    user_code = f"{secrets.token_hex(2)}-{secrets.token_hex(2)}"
+
+    # Set up our provider callbacks.
+    # NOTE that these callbacks will be called on a background thread. Don't do
+    # any unprotected state mutation here.
+
+    def authorization_endpoint(headers, params):
+        assert params["client_id"] == [client_id]
+
+        resp = {
+            "device_code": device_code,
+            "user_code": user_code,
+            "interval": 0,
+            "verification_uri": "https://example.com/device",
+            "expires_in": 5,
+        }
+
+        return 200, resp
+
+    openid_provider.register_endpoint(
+        "device_authorization_endpoint", "POST", "/device", authorization_endpoint
+    )
+
+    retry_lock = threading.Lock()
+
+    def token_endpoint(headers, params):
+        with retry_lock:
+            nonlocal retries
+
+            # If the test wants to force the client to retry, return an
+            # authorization_pending response and decrement the retry count.
+            if retries > 0:
+                retries -= 1
+                return 400, {"error": "authorization_pending"}
+
+        return 400, failure_mode
+
+    openid_provider.register_endpoint(
+        "token_endpoint", "POST", "/token", token_endpoint
+    )
+
+    with sock:
+        with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+            # Initiate a handshake, which should result in the above endpoints
+            # being called.
+            startup = pq3.recv1(conn, cls=pq3.Startup)
+            assert startup.proto == pq3.protocol(3, 0)
+
+            pq3.send(
+                conn,
+                pq3.types.AuthnRequest,
+                type=pq3.authn.SASL,
+                body=[b"OAUTHBEARER", b""],
+            )
+
+            # The client should not continue the connection due to the hardcoded
+            # provider failure; we disconnect here.
+
+    # Now make sure the client correctly failed.
+    with pytest.raises(psycopg2.OperationalError, match=error_pattern):
+        client.check_completed()
+
+
[email protected]("scope", [None, "openid email"])
[email protected](
+    "base_response",
+    [
+        {"status": "invalid_token"},
+        {"extra_object": {"key": "value"}, "status": "invalid_token"},
+        {"extra_object": {"status": 1}, "status": "invalid_token"},
+    ],
+)
+def test_oauth_discovery(accept, openid_provider, base_response, scope):
+    sock, client = accept(oauth_client_id=secrets.token_hex())
+
+    device_code = secrets.token_hex()
+    user_code = f"{secrets.token_hex(2)}-{secrets.token_hex(2)}"
+    verification_url = "https://example.com/device"
+
+    access_token = secrets.token_urlsafe()
+
+    # Set up our provider callbacks.
+    # NOTE that these callbacks will be called on a background thread. Don't do
+    # any unprotected state mutation here.
+
+    def authorization_endpoint(headers, params):
+        if scope:
+            assert params["scope"] == [scope]
+        else:
+            assert "scope" not in params
+
+        resp = {
+            "device_code": device_code,
+            "user_code": user_code,
+            "interval": 0,
+            "verification_uri": verification_url,
+            "expires_in": 5,
+        }
+
+        return 200, resp
+
+    openid_provider.register_endpoint(
+        "device_authorization_endpoint", "POST", "/device", authorization_endpoint
+    )
+
+    def token_endpoint(headers, params):
+        assert params["grant_type"] == ["urn:ietf:params:oauth:grant-type:device_code"]
+        assert params["device_code"] == [device_code]
+
+        # Successfully finish the request by sending the access bearer token.
+        resp = {
+            "access_token": access_token,
+            "token_type": "bearer",
+        }
+
+        return 200, resp
+
+    openid_provider.register_endpoint(
+        "token_endpoint", "POST", "/token", token_endpoint
+    )
+
+    with sock:
+        with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+            initial = start_oauth_handshake(conn)
+
+            # For discovery, the client should send an empty auth header. See
+            # RFC 7628, Sec. 4.3.
+            auth = get_auth_value(initial)
+            assert auth == b""
+
+            # We will fail the first SASL exchange. First return a link to the
+            # discovery document, pointing to the test provider server.
+            resp = dict(base_response)
+
+            discovery_uri = f"{openid_provider.issuer}/.well-known/openid-configuration"
+            resp["openid-configuration"] = discovery_uri
+
+            if scope:
+                resp["scope"] = scope
+
+            resp = json.dumps(resp)
+
+            pq3.send(
+                conn,
+                pq3.types.AuthnRequest,
+                type=pq3.authn.SASLContinue,
+                body=resp.encode("ascii"),
+            )
+
+            # Per RFC, the client is required to send a dummy ^A response.
+            pkt = pq3.recv1(conn)
+            assert pkt.type == pq3.types.PasswordMessage
+            assert pkt.payload == b"\x01"
+
+            # Now fail the SASL exchange.
+            pq3.send(
+                conn,
+                pq3.types.ErrorResponse,
+                fields=[
+                    b"SFATAL",
+                    b"C28000",
+                    b"Mdoesn't matter",
+                    b"",
+                ],
+            )
+
+    # The client will connect to us a second time, using the parameters we sent
+    # it.
+    sock, _ = accept()
+
+    with sock:
+        with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+            initial = start_oauth_handshake(conn)
+
+            # Validate and accept the token.
+            auth = get_auth_value(initial)
+            assert auth == f"Bearer {access_token}".encode("ascii")
+
+            pq3.send(conn, pq3.types.AuthnRequest, type=pq3.authn.SASLFinal)
+            finish_handshake(conn)
+
+
[email protected](
+    "response,expected_error",
+    [
+        pytest.param(
+            "abcde",
+            'Token "abcde" is invalid',
+            id="bad JSON: invalid syntax",
+        ),
+        pytest.param(
+            '"abcde"',
+            "top-level element must be an object",
+            id="bad JSON: top-level element is a string",
+        ),
+        pytest.param(
+            "[]",
+            "top-level element must be an object",
+            id="bad JSON: top-level element is an array",
+        ),
+        pytest.param(
+            "{}",
+            "server sent error response without a status",
+            id="bad JSON: no status member",
+        ),
+        pytest.param(
+            '{ "status": null }',
+            'field "status" must be a string',
+            id="bad JSON: null status member",
+        ),
+        pytest.param(
+            '{ "status": 0 }',
+            'field "status" must be a string',
+            id="bad JSON: int status member",
+        ),
+        pytest.param(
+            '{ "status": [ "bad" ] }',
+            'field "status" must be a string',
+            id="bad JSON: array status member",
+        ),
+        pytest.param(
+            '{ "status": { "bad": "bad" } }',
+            'field "status" must be a string',
+            id="bad JSON: object status member",
+        ),
+        pytest.param(
+            '{ "nested": { "status": "bad" } }',
+            "server sent error response without a status",
+            id="bad JSON: nested status",
+        ),
+        pytest.param(
+            '{ "status": "invalid_token" ',
+            "The input string ended unexpectedly",
+            id="bad JSON: unterminated object",
+        ),
+        pytest.param(
+            '{ "status": "invalid_token" } { }',
+            'Expected end of input, but found "{"',
+            id="bad JSON: trailing data",
+        ),
+        pytest.param(
+            '{ "status": "invalid_token", "openid-configuration": 1 }',
+            'field "openid-configuration" must be a string',
+            id="bad JSON: int openid-configuration member",
+        ),
+        pytest.param(
+            '{ "status": "invalid_token", "openid-configuration": 1 }',
+            'field "openid-configuration" must be a string',
+            id="bad JSON: int openid-configuration member",
+        ),
+        pytest.param(
+            '{ "status": "invalid_token", "scope": 1 }',
+            'field "scope" must be a string',
+            id="bad JSON: int scope member",
+        ),
+    ],
+)
+def test_oauth_discovery_server_error(accept, response, expected_error):
+    sock, client = accept(oauth_client_id=secrets.token_hex())
+
+    with sock:
+        with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+            initial = start_oauth_handshake(conn)
+
+            # Fail the SASL exchange with an invalid JSON response.
+            pq3.send(
+                conn,
+                pq3.types.AuthnRequest,
+                type=pq3.authn.SASLContinue,
+                body=response.encode("utf-8"),
+            )
+
+            # The client should disconnect, so the socket is closed here. (If
+            # the client doesn't disconnect, it will report a different error
+            # below and the test will fail.)
+
+    with pytest.raises(psycopg2.OperationalError, match=expected_error):
+        client.check_completed()
+
+
[email protected](
+    "sasl_err,resp_type,resp_payload,expected_error",
+    [
+        pytest.param(
+            {"status": "invalid_request"},
+            pq3.types.ErrorResponse,
+            dict(
+                fields=[b"SFATAL", b"C28000", b"Mexpected error message", b""],
+            ),
+            "expected error message",
+            id="standard server error: invalid_request",
+        ),
+        pytest.param(
+            {"status": "invalid_token"},
+            pq3.types.ErrorResponse,
+            dict(
+                fields=[b"SFATAL", b"C28000", b"Mexpected error message", b""],
+            ),
+            "expected error message",
+            id="standard server error: invalid_token without discovery URI",
+        ),
+        pytest.param(
+            {"status": "invalid_request"},
+            pq3.types.AuthnRequest,
+            dict(type=pq3.authn.SASLContinue, body=b""),
+            "server sent additional OAuth data",
+            id="broken server: additional challenge after error",
+        ),
+        pytest.param(
+            {"status": "invalid_request"},
+            pq3.types.AuthnRequest,
+            dict(type=pq3.authn.SASLFinal),
+            "server sent additional OAuth data",
+            id="broken server: SASL success after error",
+        ),
+    ],
+)
+def test_oauth_server_error(accept, sasl_err, resp_type, resp_payload, expected_error):
+    sock, client = accept()
+
+    with sock:
+        with pq3.wrap(sock, debug_stream=sys.stdout) as conn:
+            start_oauth_handshake(conn)
+
+            # Ignore the client data. Return an error "challenge".
+            resp = json.dumps(sasl_err)
+            resp = resp.encode("utf-8")
+
+            pq3.send(
+                conn, pq3.types.AuthnRequest, type=pq3.authn.SASLContinue, body=resp
+            )
+
+            # Per RFC, the client is required to send a dummy ^A response.
+            pkt = pq3.recv1(conn)
+            assert pkt.type == pq3.types.PasswordMessage
+            assert pkt.payload == b"\x01"
+
+            # Now fail the SASL exchange (in either a valid way, or an invalid
+            # one, depending on the test).
+            pq3.send(conn, resp_type, **resp_payload)
+
+    with pytest.raises(psycopg2.OperationalError, match=expected_error):
+        client.check_completed()
diff --git a/src/test/python/pq3.py b/src/test/python/pq3.py
new file mode 100644
index 0000000000..3a22dad0b6
--- /dev/null
+++ b/src/test/python/pq3.py
@@ -0,0 +1,727 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import contextlib
+import getpass
+import io
+import os
+import ssl
+import sys
+import textwrap
+
+from construct import *
+
+import tls
+
+
+def protocol(major, minor):
+    """
+    Returns the protocol version, in integer format, corresponding to the given
+    major and minor version numbers.
+    """
+    return (major << 16) | minor
+
+
+# Startup
+
+StringList = GreedyRange(NullTerminated(GreedyBytes))
+
+
+class KeyValueAdapter(Adapter):
+    """
+    Turns a key-value store into a null-terminated list of null-terminated
+    strings, as presented on the wire in the startup packet.
+    """
+
+    def _encode(self, obj, context, path):
+        if isinstance(obj, list):
+            return obj
+
+        l = []
+
+        for k, v in obj.items():
+            if isinstance(k, str):
+                k = k.encode("utf-8")
+            l.append(k)
+
+            if isinstance(v, str):
+                v = v.encode("utf-8")
+            l.append(v)
+
+        l.append(b"")
+        return l
+
+    def _decode(self, obj, context, path):
+        # TODO: turn a list back into a dict
+        return obj
+
+
+KeyValues = KeyValueAdapter(StringList)
+
+_startup_payload = Switch(
+    this.proto,
+    {
+        protocol(3, 0): KeyValues,
+    },
+    default=GreedyBytes,
+)
+
+
+def _default_protocol(this):
+    try:
+        if isinstance(this.payload, (list, dict)):
+            return protocol(3, 0)
+    except AttributeError:
+        pass  # no payload passed during build
+
+    return 0
+
+
+def _startup_payload_len(this):
+    """
+    The payload field has a fixed size based on the length of the packet. But
+    if the caller hasn't supplied an explicit length at build time, we have to
+    build the payload to figure out how long it is, which requires us to know
+    the length first... This function exists solely to break the cycle.
+    """
+    assert this._building, "_startup_payload_len() cannot be called during parsing"
+
+    try:
+        payload = this.payload
+    except AttributeError:
+        return 0  # no payload
+
+    if isinstance(payload, bytes):
+        # already serialized; just use the given length
+        return len(payload)
+
+    try:
+        proto = this.proto
+    except AttributeError:
+        proto = _default_protocol(this)
+
+    data = _startup_payload.build(payload, proto=proto)
+    return len(data)
+
+
+Startup = Struct(
+    "len" / Default(Int32sb, lambda this: _startup_payload_len(this) + 8),
+    "proto" / Default(Hex(Int32sb), _default_protocol),
+    "payload" / FixedSized(this.len - 8, Default(_startup_payload, b"")),
+)
+
+# Pq3
+
+# Adapted from construct.core.EnumIntegerString
+class EnumNamedByte:
+    def __init__(self, val, name):
+        self._val = val
+        self._name = name
+
+    def __int__(self):
+        return ord(self._val)
+
+    def __str__(self):
+        return "(enum) %s %r" % (self._name, self._val)
+
+    def __repr__(self):
+        return "EnumNamedByte(%r)" % self._val
+
+    def __eq__(self, other):
+        if isinstance(other, EnumNamedByte):
+            other = other._val
+        if not isinstance(other, bytes):
+            return NotImplemented
+
+        return self._val == other
+
+    def __hash__(self):
+        return hash(self._val)
+
+
+# Adapted from construct.core.Enum
+class ByteEnum(Adapter):
+    def __init__(self, **mapping):
+        super(ByteEnum, self).__init__(Byte)
+        self.namemapping = {k: EnumNamedByte(v, k) for k, v in mapping.items()}
+        self.decmapping = {v: EnumNamedByte(v, k) for k, v in mapping.items()}
+
+    def __getattr__(self, name):
+        if name in self.namemapping:
+            return self.decmapping[self.namemapping[name]]
+        raise AttributeError
+
+    def _decode(self, obj, context, path):
+        b = bytes([obj])
+        try:
+            return self.decmapping[b]
+        except KeyError:
+            return EnumNamedByte(b, "(unknown)")
+
+    def _encode(self, obj, context, path):
+        if isinstance(obj, int):
+            return obj
+        elif isinstance(obj, bytes):
+            return ord(obj)
+        return int(obj)
+
+
+types = ByteEnum(
+    ErrorResponse=b"E",
+    ReadyForQuery=b"Z",
+    Query=b"Q",
+    EmptyQueryResponse=b"I",
+    AuthnRequest=b"R",
+    PasswordMessage=b"p",
+    BackendKeyData=b"K",
+    CommandComplete=b"C",
+    ParameterStatus=b"S",
+    DataRow=b"D",
+    Terminate=b"X",
+)
+
+
+authn = Enum(
+    Int32ub,
+    OK=0,
+    SASL=10,
+    SASLContinue=11,
+    SASLFinal=12,
+)
+
+
+_authn_body = Switch(
+    this.type,
+    {
+        authn.OK: Terminated,
+        authn.SASL: StringList,
+    },
+    default=GreedyBytes,
+)
+
+
+def _data_len(this):
+    assert this._building, "_data_len() cannot be called during parsing"
+
+    if not hasattr(this, "data") or this.data is None:
+        return -1
+
+    return len(this.data)
+
+
+# The protocol reuses the PasswordMessage for several authentication response
+# types, and there's no good way to figure out which is which without keeping
+# state for the entire stream. So this is a separate Construct that can be
+# explicitly parsed/built by code that knows it's needed.
+SASLInitialResponse = Struct(
+    "name" / NullTerminated(GreedyBytes),
+    "len" / Default(Int32sb, lambda this: _data_len(this)),
+    "data"
+    / IfThenElse(
+        # Allow tests to explicitly pass an incorrect length during testing, by
+        # not enforcing a FixedSized during build. (The len calculation above
+        # defaults to the correct size.)
+        this._building,
+        Optional(GreedyBytes),
+        If(this.len != -1, Default(FixedSized(this.len, GreedyBytes), b"")),
+    ),
+    Terminated,  # make sure the entire response is consumed
+)
+
+
+_column = FocusedSeq(
+    "data",
+    "len" / Default(Int32sb, lambda this: _data_len(this)),
+    "data" / If(this.len != -1, FixedSized(this.len, GreedyBytes)),
+)
+
+
+_payload_map = {
+    types.ErrorResponse: Struct("fields" / StringList),
+    types.ReadyForQuery: Struct("status" / Bytes(1)),
+    types.Query: Struct("query" / NullTerminated(GreedyBytes)),
+    types.EmptyQueryResponse: Terminated,
+    types.AuthnRequest: Struct("type" / authn, "body" / Default(_authn_body, b"")),
+    types.BackendKeyData: Struct("pid" / Int32ub, "key" / Hex(Int32ub)),
+    types.CommandComplete: Struct("tag" / NullTerminated(GreedyBytes)),
+    types.ParameterStatus: Struct(
+        "name" / NullTerminated(GreedyBytes), "value" / NullTerminated(GreedyBytes)
+    ),
+    types.DataRow: Struct("columns" / Default(PrefixedArray(Int16sb, _column), b"")),
+    types.Terminate: Terminated,
+}
+
+
+_payload = FocusedSeq(
+    "_payload",
+    "_payload"
+    / Switch(
+        this._.type,
+        _payload_map,
+        default=GreedyBytes,
+    ),
+    Terminated,  # make sure every payload consumes the entire packet
+)
+
+
+def _payload_len(this):
+    """
+    See _startup_payload_len() for an explanation.
+    """
+    assert this._building, "_payload_len() cannot be called during parsing"
+
+    try:
+        payload = this.payload
+    except AttributeError:
+        return 0  # no payload
+
+    if isinstance(payload, bytes):
+        # already serialized; just use the given length
+        return len(payload)
+
+    data = _payload.build(payload, type=this.type)
+    return len(data)
+
+
+Pq3 = Struct(
+    "type" / types,
+    "len" / Default(Int32ub, lambda this: _payload_len(this) + 4),
+    "payload" / FixedSized(this.len - 4, Default(_payload, b"")),
+)
+
+
+# Environment
+
+
+def pghost():
+    return os.environ.get("PGHOST", default="localhost")
+
+
+def pgport():
+    return int(os.environ.get("PGPORT", default=5432))
+
+
+def pguser():
+    try:
+        return os.environ["PGUSER"]
+    except KeyError:
+        return getpass.getuser()
+
+
+def pgdatabase():
+    return os.environ.get("PGDATABASE", default="postgres")
+
+
+# Connections
+
+
+def _hexdump_translation_map():
+    """
+    For hexdumps. Translates any unprintable or non-ASCII bytes into '.'.
+    """
+    input = bytearray()
+
+    for i in range(128):
+        c = chr(i)
+
+        if not c.isprintable():
+            input += bytes([i])
+
+    input += bytes(range(128, 256))
+
+    return bytes.maketrans(input, b"." * len(input))
+
+
+class _DebugStream(object):
+    """
+    Wraps a file-like object and adds hexdumps of the read and write data. Call
+    end_packet() on a _DebugStream to write the accumulated hexdumps to the
+    output stream, along with the packet that was sent.
+    """
+
+    _translation_map = _hexdump_translation_map()
+
+    def __init__(self, stream, out=sys.stdout):
+        """
+        Creates a new _DebugStream wrapping the given stream (which must have
+        been created by wrap()). All attributes not provided by the _DebugStream
+        are delegated to the wrapped stream. out is the text stream to which
+        hexdumps are written.
+        """
+        self.raw = stream
+        self._out = out
+        self._rbuf = io.BytesIO()
+        self._wbuf = io.BytesIO()
+
+    def __getattr__(self, name):
+        return getattr(self.raw, name)
+
+    def __setattr__(self, name, value):
+        if name in ("raw", "_out", "_rbuf", "_wbuf"):
+            return object.__setattr__(self, name, value)
+
+        setattr(self.raw, name, value)
+
+    def read(self, *args, **kwargs):
+        buf = self.raw.read(*args, **kwargs)
+
+        self._rbuf.write(buf)
+        return buf
+
+    def write(self, b):
+        self._wbuf.write(b)
+        return self.raw.write(b)
+
+    def recv(self, *args):
+        buf = self.raw.recv(*args)
+
+        self._rbuf.write(buf)
+        return buf
+
+    def _flush(self, buf, prefix):
+        width = 16
+        hexwidth = width * 3 - 1
+
+        count = 0
+        buf.seek(0)
+
+        while True:
+            line = buf.read(16)
+
+            if not line:
+                if count:
+                    self._out.write("\n")  # separate the output block with a newline
+                return
+
+            self._out.write("%s %04X:\t" % (prefix, count))
+            self._out.write("%*s\t" % (-hexwidth, line.hex(" ")))
+            self._out.write(line.translate(self._translation_map).decode("ascii"))
+            self._out.write("\n")
+
+            count += 16
+
+    def print_debug(self, obj, *, prefix=""):
+        contents = ""
+        if obj is not None:
+            contents = str(obj)
+
+        for line in contents.splitlines():
+            self._out.write("%s%s\n" % (prefix, line))
+
+        self._out.write("\n")
+
+    def flush_debug(self, *, prefix=""):
+        self._flush(self._rbuf, prefix + "<")
+        self._rbuf = io.BytesIO()
+
+        self._flush(self._wbuf, prefix + ">")
+        self._wbuf = io.BytesIO()
+
+    def end_packet(self, pkt, *, read=False, prefix="", indent="  "):
+        """
+        Marks the end of a logical "packet" of data. A string representation of
+        pkt will be printed, and the debug buffers will be flushed with an
+        indent. All lines can be optionally prefixed.
+
+        If read is True, the packet representation is written after the debug
+        buffers; otherwise the default of False (meaning write) causes the
+        packet representation to be dumped first. This is meant to capture the
+        logical flow of layer translation.
+        """
+        write = not read
+
+        if write:
+            self.print_debug(pkt, prefix=prefix + "> ")
+
+        self.flush_debug(prefix=prefix + indent)
+
+        if read:
+            self.print_debug(pkt, prefix=prefix + "< ")
+
+
[email protected]
+def wrap(socket, *, debug_stream=None):
+    """
+    Transforms a raw socket into a connection that can be used for Construct
+    building and parsing. The return value is a context manager and can be used
+    in a with statement.
+    """
+    # It is critical that buffering be disabled here, so that we can still
+    # manipulate the raw socket without desyncing the stream.
+    with socket.makefile("rwb", buffering=0) as sfile:
+        # Expose the original socket's recv() on the SocketIO object we return.
+        def recv(self, *args):
+            return socket.recv(*args)
+
+        sfile.recv = recv.__get__(sfile)
+
+        conn = sfile
+        if debug_stream:
+            conn = _DebugStream(conn, debug_stream)
+
+        try:
+            yield conn
+        finally:
+            if debug_stream:
+                conn.flush_debug(prefix="? ")
+
+
+def _send(stream, cls, obj):
+    debugging = hasattr(stream, "flush_debug")
+    out = io.BytesIO()
+
+    # Ideally we would build directly to the passed stream, but because we need
+    # to reparse the generated output for the debugging case, build to an
+    # intermediate BytesIO and send it instead.
+    cls.build_stream(obj, out)
+    buf = out.getvalue()
+
+    stream.write(buf)
+    if debugging:
+        pkt = cls.parse(buf)
+        stream.end_packet(pkt)
+
+    stream.flush()
+
+
+def send(stream, packet_type, payload_data=None, **payloadkw):
+    """
+    Sends a packet on the given pq3 connection. type is the pq3.types member
+    that should be assigned to the packet. If payload_data is given, it will be
+    used as the packet payload; otherwise the key/value pairs in payloadkw will
+    be the payload contents.
+    """
+    data = payloadkw
+
+    if payload_data is not None:
+        if payloadkw:
+            raise ValueError(
+                "payload_data and payload keywords may not be used simultaneously"
+            )
+
+        data = payload_data
+
+    _send(stream, Pq3, dict(type=packet_type, payload=data))
+
+
+def send_startup(stream, proto=None, **kwargs):
+    """
+    Sends a startup packet on the given pq3 connection. In most cases you should
+    use the handshake functions instead, which will do this for you.
+
+    By default, a protocol version 3 packet will be sent. This can be overridden
+    with the proto parameter.
+    """
+    pkt = {}
+
+    if proto is not None:
+        pkt["proto"] = proto
+    if kwargs:
+        pkt["payload"] = kwargs
+
+    _send(stream, Startup, pkt)
+
+
+def recv1(stream, *, cls=Pq3):
+    """
+    Receives a single pq3 packet from the given stream and returns it.
+    """
+    resp = cls.parse_stream(stream)
+
+    debugging = hasattr(stream, "flush_debug")
+    if debugging:
+        stream.end_packet(resp, read=True)
+
+    return resp
+
+
+def handshake(stream, **kwargs):
+    """
+    Performs a libpq v3 startup handshake. kwargs should contain the key/value
+    parameters to send to the server in the startup packet.
+    """
+    # Send our startup parameters.
+    send_startup(stream, **kwargs)
+
+    # Receive and dump packets until the server indicates it's ready for our
+    # first query.
+    while True:
+        resp = recv1(stream)
+        if resp is None:
+            raise RuntimeError("server closed connection during handshake")
+
+        if resp.type == types.ReadyForQuery:
+            return
+        elif resp.type == types.ErrorResponse:
+            raise RuntimeError(
+                f"received error response from peer: {resp.payload.fields!r}"
+            )
+
+
+# TLS
+
+
+class _TLSStream(object):
+    """
+    A file-like object that performs TLS encryption/decryption on a wrapped
+    stream. Differs from ssl.SSLSocket in that we have full visibility and
+    control over the TLS layer.
+    """
+
+    def __init__(self, stream, context):
+        self._stream = stream
+        self._debugging = hasattr(stream, "flush_debug")
+
+        self._in = ssl.MemoryBIO()
+        self._out = ssl.MemoryBIO()
+        self._ssl = context.wrap_bio(self._in, self._out)
+
+    def handshake(self):
+        try:
+            self._pump(lambda: self._ssl.do_handshake())
+        finally:
+            self._flush_debug(prefix="? ")
+
+    def read(self, *args):
+        return self._pump(lambda: self._ssl.read(*args))
+
+    def write(self, *args):
+        return self._pump(lambda: self._ssl.write(*args))
+
+    def _decode(self, buf):
+        """
+        Attempts to decode a buffer of TLS data into a packet representation
+        that can be printed.
+
+        TODO: handle buffers (and record fragments) that don't align with packet
+        boundaries.
+        """
+        end = len(buf)
+        bio = io.BytesIO(buf)
+
+        ret = io.StringIO()
+
+        while bio.tell() < end:
+            record = tls.Plaintext.parse_stream(bio)
+
+            if ret.tell() > 0:
+                ret.write("\n")
+            ret.write("[Record] ")
+            ret.write(str(record))
+            ret.write("\n")
+
+            if record.type == tls.ContentType.handshake:
+                record_cls = tls.Handshake
+            else:
+                continue
+
+            innerlen = len(record.fragment)
+            inner = io.BytesIO(record.fragment)
+
+            while inner.tell() < innerlen:
+                msg = record_cls.parse_stream(inner)
+
+                indented = "[Message] " + str(msg)
+                indented = textwrap.indent(indented, "    ")
+
+                ret.write("\n")
+                ret.write(indented)
+                ret.write("\n")
+
+        return ret.getvalue()
+
+    def flush(self):
+        if not self._out.pending:
+            self._stream.flush()
+            return
+
+        buf = self._out.read()
+        self._stream.write(buf)
+
+        if self._debugging:
+            pkt = self._decode(buf)
+            self._stream.end_packet(pkt, prefix="  ")
+
+        self._stream.flush()
+
+    def _pump(self, operation):
+        while True:
+            try:
+                return operation()
+            except (ssl.SSLWantReadError, ssl.SSLWantWriteError) as e:
+                want = e
+            self._read_write(want)
+
+    def _recv(self, maxsize):
+        buf = self._stream.recv(4096)
+        if not buf:
+            self._in.write_eof()
+            return
+
+        self._in.write(buf)
+
+        if not self._debugging:
+            return
+
+        pkt = self._decode(buf)
+        self._stream.end_packet(pkt, read=True, prefix="  ")
+
+    def _read_write(self, want):
+        # XXX This needs work. So many corner cases yet to handle. For one,
+        # doing blocking writes in flush may lead to distributed deadlock if the
+        # peer is already blocking on its writes.
+
+        if isinstance(want, ssl.SSLWantWriteError):
+            assert self._out.pending, "SSL backend wants write without data"
+
+        self.flush()
+
+        if isinstance(want, ssl.SSLWantReadError):
+            self._recv(4096)
+
+    def _flush_debug(self, prefix):
+        if not self._debugging:
+            return
+
+        self._stream.flush_debug(prefix=prefix)
+
+
[email protected]
+def tls_handshake(stream, context):
+    """
+    Performs a TLS handshake over the given stream (which must have been created
+    via a call to wrap()), and returns a new stream which transparently tunnels
+    data over the TLS connection.
+
+    If the passed stream has debugging enabled, the returned stream will also
+    have debugging, using the same output IO.
+    """
+    debugging = hasattr(stream, "flush_debug")
+
+    # Send our startup parameters.
+    send_startup(stream, proto=protocol(1234, 5679))
+
+    # Look at the SSL response.
+    resp = stream.read(1)
+    if debugging:
+        stream.flush_debug(prefix="  ")
+
+    if resp == b"N":
+        raise RuntimeError("server does not support SSLRequest")
+    if resp != b"S":
+        raise RuntimeError(f"unexpected response of type {resp!r} during TLS startup")
+
+    tls = _TLSStream(stream, context)
+    tls.handshake()
+
+    if debugging:
+        tls = _DebugStream(tls, stream._out)
+
+    try:
+        yield tls
+        # TODO: teardown/unwrap the connection?
+    finally:
+        if debugging:
+            tls.flush_debug(prefix="? ")
diff --git a/src/test/python/pytest.ini b/src/test/python/pytest.ini
new file mode 100644
index 0000000000..ab7a6e7fb9
--- /dev/null
+++ b/src/test/python/pytest.ini
@@ -0,0 +1,4 @@
+[pytest]
+
+markers =
+    slow: mark test as slow
diff --git a/src/test/python/requirements.txt b/src/test/python/requirements.txt
new file mode 100644
index 0000000000..32f105ea84
--- /dev/null
+++ b/src/test/python/requirements.txt
@@ -0,0 +1,7 @@
+black
+cryptography~=3.4.6
+construct~=2.10.61
+isort~=5.6
+psycopg2~=2.8.6
+pytest~=6.1
+pytest-asyncio~=0.14.0
diff --git a/src/test/python/server/__init__.py b/src/test/python/server/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/src/test/python/server/conftest.py b/src/test/python/server/conftest.py
new file mode 100644
index 0000000000..ba7342a453
--- /dev/null
+++ b/src/test/python/server/conftest.py
@@ -0,0 +1,45 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import contextlib
+import socket
+import sys
+
+import pytest
+
+import pq3
+
+
[email protected]
+def connect():
+    """
+    A factory fixture that, when called, returns a socket connected to a
+    Postgres server, wrapped in a pq3 connection. The calling test will be
+    skipped automatically if a server is not running at PGHOST:PGPORT, so it's
+    best to connect as soon as possible after the test case begins, to avoid
+    doing unnecessary work.
+    """
+    # Set up an ExitStack to handle safe cleanup of all of the moving pieces.
+    with contextlib.ExitStack() as stack:
+
+        def conn_factory():
+            addr = (pq3.pghost(), pq3.pgport())
+
+            try:
+                sock = socket.create_connection(addr, timeout=2)
+            except ConnectionError as e:
+                pytest.skip(f"unable to connect to {addr}: {e}")
+
+            # Have ExitStack close our socket.
+            stack.enter_context(sock)
+
+            # Wrap the connection in a pq3 layer and have ExitStack clean it up
+            # too.
+            wrap_ctx = pq3.wrap(sock, debug_stream=sys.stdout)
+            conn = stack.enter_context(wrap_ctx)
+
+            return conn
+
+        yield conn_factory
diff --git a/src/test/python/server/test_oauth.py b/src/test/python/server/test_oauth.py
new file mode 100644
index 0000000000..cb5ca7fa23
--- /dev/null
+++ b/src/test/python/server/test_oauth.py
@@ -0,0 +1,1012 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import base64
+import contextlib
+import json
+import os
+import pathlib
+import secrets
+import shlex
+import shutil
+import socket
+import struct
+from multiprocessing import shared_memory
+
+import psycopg2
+import pytest
+from psycopg2 import sql
+
+import pq3
+
+MAX_SASL_MESSAGE_LENGTH = 65535
+
+INVALID_AUTHORIZATION_ERRCODE = b"28000"
+PROTOCOL_VIOLATION_ERRCODE = b"08P01"
+FEATURE_NOT_SUPPORTED_ERRCODE = b"0A000"
+
+SHARED_MEM_NAME = "oauth-pytest"
+MAX_TOKEN_SIZE = 4096
+MAX_UINT16 = 2 ** 16 - 1
+
+
+def skip_if_no_postgres():
+    """
+    Used by the oauth_ctx fixture to skip this test module if no Postgres server
+    is running.
+
+    This logic is nearly duplicated with the conn fixture. Ideally oauth_ctx
+    would depend on that, but a module-scope fixture can't depend on a
+    test-scope fixture, and we haven't reached the rule of three yet.
+    """
+    addr = (pq3.pghost(), pq3.pgport())
+
+    try:
+        with socket.create_connection(addr, timeout=2):
+            pass
+    except ConnectionError as e:
+        pytest.skip(f"unable to connect to {addr}: {e}")
+
+
[email protected]
+def prepend_file(path, lines):
+    """
+    A context manager that prepends a file on disk with the desired lines of
+    text. When the context manager is exited, the file will be restored to its
+    original contents.
+    """
+    # First make a backup of the original file.
+    bak = path + ".bak"
+    shutil.copy2(path, bak)
+
+    try:
+        # Write the new lines, followed by the original file content.
+        with open(path, "w") as new, open(bak, "r") as orig:
+            new.writelines(lines)
+            shutil.copyfileobj(orig, new)
+
+        # Return control to the calling code.
+        yield
+
+    finally:
+        # Put the backup back into place.
+        os.replace(bak, path)
+
+
[email protected](scope="module")
+def oauth_ctx():
+    """
+    Creates a database and user that use the oauth auth method. The context
+    object contains the dbname and user attributes as strings to be used during
+    connection, as well as the issuer and scope that have been set in the HBA
+    configuration.
+
+    This fixture assumes that the standard PG* environment variables point to a
+    server running on a local machine, and that the PGUSER has rights to create
+    databases and roles.
+    """
+    skip_if_no_postgres()  # don't bother running these tests without a server
+
+    id = secrets.token_hex(4)
+
+    class Context:
+        dbname = "oauth_test_" + id
+
+        user = "oauth_user_" + id
+        map_user = "oauth_map_user_" + id
+        authz_user = "oauth_authz_user_" + id
+
+        issuer = "https://example.com/" + id
+        scope = "openid " + id
+
+    ctx = Context()
+    hba_lines = (
+        f'host {ctx.dbname} {ctx.map_user}   samehost oauth issuer="{ctx.issuer}" scope="{ctx.scope}" map=oauth\n',
+        f'host {ctx.dbname} {ctx.authz_user} samehost oauth issuer="{ctx.issuer}" scope="{ctx.scope}" trust_validator_authz=1\n',
+        f'host {ctx.dbname} all              samehost oauth issuer="{ctx.issuer}" scope="{ctx.scope}"\n',
+    )
+    ident_lines = (r"oauth /^(.*)@example\.com$ \1",)
+
+    conn = psycopg2.connect("")
+    conn.autocommit = True
+
+    with contextlib.closing(conn):
+        c = conn.cursor()
+
+        # Create our roles and database.
+        user = sql.Identifier(ctx.user)
+        map_user = sql.Identifier(ctx.map_user)
+        authz_user = sql.Identifier(ctx.authz_user)
+        dbname = sql.Identifier(ctx.dbname)
+
+        c.execute(sql.SQL("CREATE ROLE {} LOGIN;").format(user))
+        c.execute(sql.SQL("CREATE ROLE {} LOGIN;").format(map_user))
+        c.execute(sql.SQL("CREATE ROLE {} LOGIN;").format(authz_user))
+        c.execute(sql.SQL("CREATE DATABASE {};").format(dbname))
+
+        # Make this test script the server's oauth_validator.
+        path = pathlib.Path(__file__).parent / "validate_bearer.py"
+        path = str(path.absolute())
+
+        cmd = f"{shlex.quote(path)} {SHARED_MEM_NAME} <&%f"
+        c.execute("ALTER SYSTEM SET oauth_validator_command TO %s;", (cmd,))
+
+        # Replace pg_hba and pg_ident.
+        c.execute("SHOW hba_file;")
+        hba = c.fetchone()[0]
+
+        c.execute("SHOW ident_file;")
+        ident = c.fetchone()[0]
+
+        with prepend_file(hba, hba_lines), prepend_file(ident, ident_lines):
+            c.execute("SELECT pg_reload_conf();")
+
+            # Use the new database and user.
+            yield ctx
+
+        # Put things back the way they were.
+        c.execute("SELECT pg_reload_conf();")
+
+        c.execute("ALTER SYSTEM RESET oauth_validator_command;")
+        c.execute(sql.SQL("DROP DATABASE {};").format(dbname))
+        c.execute(sql.SQL("DROP ROLE {};").format(authz_user))
+        c.execute(sql.SQL("DROP ROLE {};").format(map_user))
+        c.execute(sql.SQL("DROP ROLE {};").format(user))
+
+
[email protected]()
+def conn(oauth_ctx, connect):
+    """
+    A convenience wrapper for connect(). The main purpose of this fixture is to
+    make sure oauth_ctx runs its setup code before the connection is made.
+    """
+    return connect()
+
+
[email protected](scope="module", autouse=True)
+def authn_id_extension(oauth_ctx):
+    """
+    Performs a `CREATE EXTENSION authn_id` in the test database. This fixture is
+    autoused, so tests don't need to rely on it.
+    """
+    conn = psycopg2.connect(database=oauth_ctx.dbname)
+    conn.autocommit = True
+
+    with contextlib.closing(conn):
+        c = conn.cursor()
+        c.execute("CREATE EXTENSION authn_id;")
+
+
[email protected](scope="session")
+def shared_mem():
+    """
+    Yields a shared memory segment that can be used for communication between
+    the bearer_token fixture and ./validate_bearer.py.
+    """
+    size = MAX_TOKEN_SIZE + 2  # two byte length prefix
+    mem = shared_memory.SharedMemory(SHARED_MEM_NAME, create=True, size=size)
+
+    try:
+        with contextlib.closing(mem):
+            yield mem
+    finally:
+        mem.unlink()
+
+
[email protected]()
+def bearer_token(shared_mem):
+    """
+    Returns a factory function that, when called, will store a Bearer token in
+    shared_mem. If token is None (the default), a new token will be generated
+    using secrets.token_urlsafe() and returned; otherwise the passed token will
+    be used as-is.
+
+    When token is None, the generated token size in bytes may be specified as an
+    argument; if unset, a small 16-byte token will be generated. The token size
+    may not exceed MAX_TOKEN_SIZE in any case.
+
+    The return value is the token, converted to a bytes object.
+
+    As a special case for testing failure modes, accept_any may be set to True.
+    This signals to the validator command that any bearer token should be
+    accepted. The returned token in this case may be used or discarded as needed
+    by the test.
+    """
+
+    def set_token(token=None, *, size=16, accept_any=False):
+        if token is not None:
+            size = len(token)
+
+        if size > MAX_TOKEN_SIZE:
+            raise ValueError(f"token size {size} exceeds maximum size {MAX_TOKEN_SIZE}")
+
+        if token is None:
+            if size % 4:
+                raise ValueError(f"requested token size {size} is not a multiple of 4")
+
+            token = secrets.token_urlsafe(size // 4 * 3)
+            assert len(token) == size
+
+        try:
+            token = token.encode("ascii")
+        except AttributeError:
+            pass  # already encoded
+
+        if accept_any:
+            # Two-byte magic value.
+            shared_mem.buf[:2] = struct.pack("H", MAX_UINT16)
+        else:
+            # Two-byte length prefix, then the token data.
+            shared_mem.buf[:2] = struct.pack("H", len(token))
+            shared_mem.buf[2 : size + 2] = token
+
+        return token
+
+    return set_token
+
+
+def begin_oauth_handshake(conn, oauth_ctx, *, user=None):
+    if user is None:
+        user = oauth_ctx.authz_user
+
+    pq3.send_startup(conn, user=user, database=oauth_ctx.dbname)
+
+    resp = pq3.recv1(conn)
+    assert resp.type == pq3.types.AuthnRequest
+
+    # The server should advertise exactly one mechanism.
+    assert resp.payload.type == pq3.authn.SASL
+    assert resp.payload.body == [b"OAUTHBEARER", b""]
+
+
+def send_initial_response(conn, *, auth=None, bearer=None):
+    """
+    Sends the OAUTHBEARER initial response on the connection, using the given
+    bearer token. Alternatively to a bearer token, the initial response's auth
+    field may be explicitly specified to test corner cases.
+    """
+    if bearer is not None and auth is not None:
+        raise ValueError("exactly one of the auth and bearer kwargs must be set")
+
+    if bearer is not None:
+        auth = b"Bearer " + bearer
+
+    if auth is None:
+        raise ValueError("exactly one of the auth and bearer kwargs must be set")
+
+    initial = pq3.SASLInitialResponse.build(
+        dict(
+            name=b"OAUTHBEARER",
+            data=b"n,,\x01auth=" + auth + b"\x01\x01",
+        )
+    )
+    pq3.send(conn, pq3.types.PasswordMessage, initial)
+
+
+def expect_handshake_success(conn):
+    """
+    Validates that the server responds with an AuthnOK message, and then drains
+    the connection until a ReadyForQuery message is received.
+    """
+    resp = pq3.recv1(conn)
+
+    assert resp.type == pq3.types.AuthnRequest
+    assert resp.payload.type == pq3.authn.OK
+    assert not resp.payload.body
+
+    receive_until(conn, pq3.types.ReadyForQuery)
+
+
+def expect_handshake_failure(conn, oauth_ctx):
+    """
+    Performs the OAUTHBEARER SASL failure "handshake" and validates the server's
+    side of the conversation, including the final ErrorResponse.
+    """
+
+    # We expect a discovery "challenge" back from the server before the authn
+    # failure message.
+    resp = pq3.recv1(conn)
+    assert resp.type == pq3.types.AuthnRequest
+
+    req = resp.payload
+    assert req.type == pq3.authn.SASLContinue
+
+    body = json.loads(req.body)
+    assert body["status"] == "invalid_token"
+    assert body["scope"] == oauth_ctx.scope
+
+    expected_config = oauth_ctx.issuer + "/.well-known/openid-configuration"
+    assert body["openid-configuration"] == expected_config
+
+    # Send the dummy response to complete the failed handshake.
+    pq3.send(conn, pq3.types.PasswordMessage, b"\x01")
+    resp = pq3.recv1(conn)
+
+    err = ExpectedError(INVALID_AUTHORIZATION_ERRCODE, "bearer authentication failed")
+    err.match(resp)
+
+
+def receive_until(conn, type):
+    """
+    receive_until pulls packets off the pq3 connection until a packet with the
+    desired type is found, or an error response is received.
+    """
+    while True:
+        pkt = pq3.recv1(conn)
+
+        if pkt.type == type:
+            return pkt
+        elif pkt.type == pq3.types.ErrorResponse:
+            raise RuntimeError(
+                f"received error response from peer: {pkt.payload.fields!r}"
+            )
+
+
[email protected]("token_len", [16, 1024, 4096])
[email protected](
+    "auth_prefix",
+    [
+        b"Bearer ",
+        b"bearer ",
+        b"Bearer    ",
+    ],
+)
+def test_oauth(conn, oauth_ctx, bearer_token, auth_prefix, token_len):
+    begin_oauth_handshake(conn, oauth_ctx)
+
+    # Generate our bearer token with the desired length.
+    token = bearer_token(size=token_len)
+    auth = auth_prefix + token
+
+    send_initial_response(conn, auth=auth)
+    expect_handshake_success(conn)
+
+    # Make sure that the server has not set an authenticated ID.
+    pq3.send(conn, pq3.types.Query, query=b"SELECT authn_id();")
+    resp = receive_until(conn, pq3.types.DataRow)
+
+    row = resp.payload
+    assert row.columns == [None]
+
+
[email protected](
+    "token_value",
+    [
+        "abcdzA==",
+        "123456M=",
+        "x-._~+/x",
+    ],
+)
+def test_oauth_bearer_corner_cases(conn, oauth_ctx, bearer_token, token_value):
+    begin_oauth_handshake(conn, oauth_ctx)
+
+    send_initial_response(conn, bearer=bearer_token(token_value))
+
+    expect_handshake_success(conn)
+
+
[email protected](
+    "user,authn_id,should_succeed",
+    [
+        pytest.param(
+            lambda ctx: ctx.user,
+            lambda ctx: ctx.user,
+            True,
+            id="validator authn: succeeds when authn_id == username",
+        ),
+        pytest.param(
+            lambda ctx: ctx.user,
+            lambda ctx: None,
+            False,
+            id="validator authn: fails when authn_id is not set",
+        ),
+        pytest.param(
+            lambda ctx: ctx.user,
+            lambda ctx: ctx.authz_user,
+            False,
+            id="validator authn: fails when authn_id != username",
+        ),
+        pytest.param(
+            lambda ctx: ctx.map_user,
+            lambda ctx: ctx.map_user + "@example.com",
+            True,
+            id="validator with map: succeeds when authn_id matches map",
+        ),
+        pytest.param(
+            lambda ctx: ctx.map_user,
+            lambda ctx: None,
+            False,
+            id="validator with map: fails when authn_id is not set",
+        ),
+        pytest.param(
+            lambda ctx: ctx.map_user,
+            lambda ctx: ctx.map_user + "@example.net",
+            False,
+            id="validator with map: fails when authn_id doesn't match map",
+        ),
+        pytest.param(
+            lambda ctx: ctx.authz_user,
+            lambda ctx: None,
+            True,
+            id="validator authz: succeeds with no authn_id",
+        ),
+        pytest.param(
+            lambda ctx: ctx.authz_user,
+            lambda ctx: "",
+            True,
+            id="validator authz: succeeds with empty authn_id",
+        ),
+        pytest.param(
+            lambda ctx: ctx.authz_user,
+            lambda ctx: "postgres",
+            True,
+            id="validator authz: succeeds with basic username",
+        ),
+        pytest.param(
+            lambda ctx: ctx.authz_user,
+            lambda ctx: "[email protected]",
+            True,
+            id="validator authz: succeeds with email address",
+        ),
+    ],
+)
+def test_oauth_authn_id(conn, oauth_ctx, bearer_token, user, authn_id, should_succeed):
+    token = None
+
+    authn_id = authn_id(oauth_ctx)
+    if authn_id is not None:
+        authn_id = authn_id.encode("ascii")
+
+        # As a hack to get the validator to reflect arbitrary output from this
+        # test, encode the desired output as a base64 token. The validator will
+        # key on the leading "output=" to differentiate this from the random
+        # tokens generated by secrets.token_urlsafe().
+        output = b"output=" + authn_id + b"\n"
+        token = base64.urlsafe_b64encode(output)
+
+    token = bearer_token(token)
+    username = user(oauth_ctx)
+
+    begin_oauth_handshake(conn, oauth_ctx, user=username)
+    send_initial_response(conn, bearer=token)
+
+    if not should_succeed:
+        expect_handshake_failure(conn, oauth_ctx)
+        return
+
+    expect_handshake_success(conn)
+
+    # Check the reported authn_id.
+    pq3.send(conn, pq3.types.Query, query=b"SELECT authn_id();")
+    resp = receive_until(conn, pq3.types.DataRow)
+
+    row = resp.payload
+    assert row.columns == [authn_id]
+
+
+class ExpectedError(object):
+    def __init__(self, code, msg=None, detail=None):
+        self.code = code
+        self.msg = msg
+        self.detail = detail
+
+        # Protect against the footgun of an accidental empty string, which will
+        # "match" anything. If you don't want to match message or detail, just
+        # don't pass them.
+        if self.msg == "":
+            raise ValueError("msg must be non-empty or None")
+        if self.detail == "":
+            raise ValueError("detail must be non-empty or None")
+
+    def _getfield(self, resp, type):
+        """
+        Searches an ErrorResponse for a single field of the given type (e.g.
+        "M", "C", "D") and returns its value. Asserts if it doesn't find exactly
+        one field.
+        """
+        prefix = type.encode("ascii")
+        fields = [f for f in resp.payload.fields if f.startswith(prefix)]
+
+        assert len(fields) == 1
+        return fields[0][1:]  # strip off the type byte
+
+    def match(self, resp):
+        """
+        Checks that the given response matches the expected code, message, and
+        detail (if given). The error code must match exactly. The expected
+        message and detail must be contained within the actual strings.
+        """
+        assert resp.type == pq3.types.ErrorResponse
+
+        code = self._getfield(resp, "C")
+        assert code == self.code
+
+        if self.msg:
+            msg = self._getfield(resp, "M")
+            expected = self.msg.encode("utf-8")
+            assert expected in msg
+
+        if self.detail:
+            detail = self._getfield(resp, "D")
+            expected = self.detail.encode("utf-8")
+            assert expected in detail
+
+
+def test_oauth_rejected_bearer(conn, oauth_ctx, bearer_token):
+    # Generate a new bearer token, which we will proceed not to use.
+    _ = bearer_token()
+
+    begin_oauth_handshake(conn, oauth_ctx)
+
+    # Send a bearer token that doesn't match what the validator expects. It
+    # should fail the connection.
+    send_initial_response(conn, bearer=b"xxxxxx")
+
+    expect_handshake_failure(conn, oauth_ctx)
+
+
[email protected](
+    "bad_bearer",
+    [
+        b"Bearer    ",
+        b"Bearer a===b",
+        b"Bearer hello!",
+        b"Bearer [email protected]",
+        b'OAuth realm="Example"',
+        b"",
+    ],
+)
+def test_oauth_invalid_bearer(conn, oauth_ctx, bearer_token, bad_bearer):
+    # Tell the validator to accept any token. This ensures that the invalid
+    # bearer tokens are rejected before the validation step.
+    _ = bearer_token(accept_any=True)
+
+    begin_oauth_handshake(conn, oauth_ctx)
+    send_initial_response(conn, auth=bad_bearer)
+
+    expect_handshake_failure(conn, oauth_ctx)
+
+
[email protected]
[email protected](
+    "resp_type,resp,err",
+    [
+        pytest.param(
+            None,
+            None,
+            None,
+            marks=pytest.mark.slow,
+            id="no response (expect timeout)",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            b"hello",
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "did not send a kvsep response",
+            ),
+            id="bad dummy response",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            b"\x01\x01",
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "did not send a kvsep response",
+            ),
+            id="multiple kvseps",
+        ),
+        pytest.param(
+            pq3.types.Query,
+            dict(query=b""),
+            ExpectedError(PROTOCOL_VIOLATION_ERRCODE, "expected SASL response"),
+            id="bad response message type",
+        ),
+    ],
+)
+def test_oauth_bad_response_to_error_challenge(conn, oauth_ctx, resp_type, resp, err):
+    begin_oauth_handshake(conn, oauth_ctx)
+
+    # Send an empty auth initial response, which will force an authn failure.
+    send_initial_response(conn, auth=b"")
+
+    # We expect a discovery "challenge" back from the server before the authn
+    # failure message.
+    pkt = pq3.recv1(conn)
+    assert pkt.type == pq3.types.AuthnRequest
+
+    req = pkt.payload
+    assert req.type == pq3.authn.SASLContinue
+
+    body = json.loads(req.body)
+    assert body["status"] == "invalid_token"
+
+    if resp_type is None:
+        # Do not send the dummy response. We should time out and not get a
+        # response from the server.
+        with pytest.raises(socket.timeout):
+            conn.read(1)
+
+        # Done with the test.
+        return
+
+    # Send the bad response.
+    pq3.send(conn, resp_type, resp)
+
+    # Make sure the server fails the connection correctly.
+    pkt = pq3.recv1(conn)
+    err.match(pkt)
+
+
[email protected](
+    "type,payload,err",
+    [
+        pytest.param(
+            pq3.types.ErrorResponse,
+            dict(fields=[b""]),
+            ExpectedError(PROTOCOL_VIOLATION_ERRCODE, "expected SASL response"),
+            id="error response in initial message",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            b"x" * (MAX_SASL_MESSAGE_LENGTH + 1),
+            ExpectedError(
+                INVALID_AUTHORIZATION_ERRCODE, "bearer authentication failed"
+            ),
+            id="overlong initial response data",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(dict(name=b"SCRAM-SHA-256")),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE, "invalid SASL authentication mechanism"
+            ),
+            id="bad SASL mechanism selection",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", len=2, data=b"x")),
+            ExpectedError(PROTOCOL_VIOLATION_ERRCODE, "insufficient data"),
+            id="SASL data underflow",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", len=0, data=b"x")),
+            ExpectedError(PROTOCOL_VIOLATION_ERRCODE, "invalid message format"),
+            id="SASL data overflow",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"")),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "message is empty",
+            ),
+            id="empty",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(
+                dict(name=b"OAUTHBEARER", data=b"n,,\x01auth=\x01\x01\0")
+            ),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "length does not match input length",
+            ),
+            id="contains null byte",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"\x01")),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "Unexpected channel-binding flag",  # XXX this is a bit strange
+            ),
+            id="initial error response",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(
+                dict(name=b"OAUTHBEARER", data=b"p=tls-server-end-point,,\x01")
+            ),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "server does not support channel binding",
+            ),
+            id="uses channel binding",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"x,,\x01")),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "Unexpected channel-binding flag",
+            ),
+            id="invalid channel binding specifier",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y")),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "Comma expected",
+            ),
+            id="bad GS2 header: missing channel binding terminator",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y,a")),
+            ExpectedError(
+                FEATURE_NOT_SUPPORTED_ERRCODE,
+                "client uses authorization identity",
+            ),
+            id="bad GS2 header: authzid in use",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y,b,")),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "Unexpected attribute",
+            ),
+            id="bad GS2 header: extra attribute",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y,")),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "Unexpected attribute 0x00",  # XXX this is a bit strange
+            ),
+            id="bad GS2 header: missing authzid terminator",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y,,")),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "Key-value separator expected",
+            ),
+            id="missing initial kvsep",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y,,")),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "Key-value separator expected",
+            ),
+            id="missing initial kvsep",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(
+                dict(name=b"OAUTHBEARER", data=b"y,,\x01\x01")
+            ),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "does not contain an auth value",
+            ),
+            id="missing auth value: empty key-value list",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(
+                dict(name=b"OAUTHBEARER", data=b"y,,\x01host=example.com\x01\x01")
+            ),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "does not contain an auth value",
+            ),
+            id="missing auth value: other keys present",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(
+                dict(name=b"OAUTHBEARER", data=b"y,,\x01host=example.com")
+            ),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "unterminated key/value pair",
+            ),
+            id="missing value terminator",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER", data=b"y,,\x01")),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "did not contain a final terminator",
+            ),
+            id="missing list terminator: empty list",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(
+                dict(name=b"OAUTHBEARER", data=b"y,,\x01auth=Bearer 0\x01")
+            ),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "did not contain a final terminator",
+            ),
+            id="missing list terminator: with auth value",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(
+                dict(name=b"OAUTHBEARER", data=b"y,,\x01auth=Bearer 0\x01\x01blah")
+            ),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "additional data after the final terminator",
+            ),
+            id="additional key after terminator",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(
+                dict(name=b"OAUTHBEARER", data=b"y,,\x01key\x01\x01")
+            ),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "key without a value",
+            ),
+            id="key without value",
+        ),
+        pytest.param(
+            pq3.types.PasswordMessage,
+            pq3.SASLInitialResponse.build(
+                dict(
+                    name=b"OAUTHBEARER",
+                    data=b"y,,\x01auth=Bearer 0\x01auth=Bearer 1\x01\x01",
+                )
+            ),
+            ExpectedError(
+                PROTOCOL_VIOLATION_ERRCODE,
+                "malformed OAUTHBEARER message",
+                "contains multiple auth values",
+            ),
+            id="multiple auth values",
+        ),
+    ],
+)
+def test_oauth_bad_initial_response(conn, oauth_ctx, type, payload, err):
+    begin_oauth_handshake(conn, oauth_ctx)
+
+    # The server expects a SASL response; give it something else instead.
+    if not isinstance(payload, dict):
+        payload = dict(payload_data=payload)
+    pq3.send(conn, type, **payload)
+
+    resp = pq3.recv1(conn)
+    err.match(resp)
+
+
+def test_oauth_empty_initial_response(conn, oauth_ctx, bearer_token):
+    begin_oauth_handshake(conn, oauth_ctx)
+
+    # Send an initial response without data.
+    initial = pq3.SASLInitialResponse.build(dict(name=b"OAUTHBEARER"))
+    pq3.send(conn, pq3.types.PasswordMessage, initial)
+
+    # The server should respond with an empty challenge so we can send the data
+    # it wants.
+    pkt = pq3.recv1(conn)
+
+    assert pkt.type == pq3.types.AuthnRequest
+    assert pkt.payload.type == pq3.authn.SASLContinue
+    assert not pkt.payload.body
+
+    # Now send the initial data.
+    data = b"n,,\x01auth=Bearer " + bearer_token() + b"\x01\x01"
+    pq3.send(conn, pq3.types.PasswordMessage, data)
+
+    # Server should now complete the handshake.
+    expect_handshake_success(conn)
+
+
[email protected]()
+def set_validator():
+    """
+    A per-test fixture that allows a test to override the setting of
+    oauth_validator_command for the cluster. The setting will be reverted during
+    teardown.
+
+    Passing None will perform an ALTER SYSTEM RESET.
+    """
+    conn = psycopg2.connect("")
+    conn.autocommit = True
+
+    with contextlib.closing(conn):
+        c = conn.cursor()
+
+        # Save the previous value.
+        c.execute("SHOW oauth_validator_command;")
+        prev_cmd = c.fetchone()[0]
+
+        def setter(cmd):
+            c.execute("ALTER SYSTEM SET oauth_validator_command TO %s;", (cmd,))
+            c.execute("SELECT pg_reload_conf();")
+
+        yield setter
+
+        # Restore the previous value.
+        c.execute("ALTER SYSTEM SET oauth_validator_command TO %s;", (prev_cmd,))
+        c.execute("SELECT pg_reload_conf();")
+
+
+def test_oauth_no_validator(oauth_ctx, set_validator, connect, bearer_token):
+    # Clear out our validator command, then establish a new connection.
+    set_validator("")
+    conn = connect()
+
+    begin_oauth_handshake(conn, oauth_ctx)
+    send_initial_response(conn, bearer=bearer_token())
+
+    # The server should fail the connection.
+    expect_handshake_failure(conn, oauth_ctx)
+
+
+def test_oauth_validator_role(oauth_ctx, set_validator, connect):
+    # Switch the validator implementation. This validator will reflect the
+    # PGUSER as the authenticated identity.
+    path = pathlib.Path(__file__).parent / "validate_reflect.py"
+    path = str(path.absolute())
+
+    set_validator(f"{shlex.quote(path)} '%r' <&%f")
+    conn = connect()
+
+    # Log in. Note that the reflection validator ignores the bearer token.
+    begin_oauth_handshake(conn, oauth_ctx, user=oauth_ctx.user)
+    send_initial_response(conn, bearer=b"dontcare")
+    expect_handshake_success(conn)
+
+    # Check the user identity.
+    pq3.send(conn, pq3.types.Query, query=b"SELECT authn_id();")
+    resp = receive_until(conn, pq3.types.DataRow)
+
+    row = resp.payload
+    expected = oauth_ctx.user.encode("utf-8")
+    assert row.columns == [expected]
+
+
+def test_oauth_role_with_shell_unsafe_characters(oauth_ctx, set_validator, connect):
+    """
+    XXX This test pins undesirable behavior. We should be able to handle any
+    valid Postgres role name.
+    """
+    # Switch the validator implementation. This validator will reflect the
+    # PGUSER as the authenticated identity.
+    path = pathlib.Path(__file__).parent / "validate_reflect.py"
+    path = str(path.absolute())
+
+    set_validator(f"{shlex.quote(path)} '%r' <&%f")
+    conn = connect()
+
+    unsafe_username = "hello'there"
+    begin_oauth_handshake(conn, oauth_ctx, user=unsafe_username)
+
+    # The server should reject the handshake.
+    send_initial_response(conn, bearer=b"dontcare")
+    expect_handshake_failure(conn, oauth_ctx)
diff --git a/src/test/python/server/test_server.py b/src/test/python/server/test_server.py
new file mode 100644
index 0000000000..02126dba79
--- /dev/null
+++ b/src/test/python/server/test_server.py
@@ -0,0 +1,21 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import pq3
+
+
+def test_handshake(connect):
+    """Basic sanity check."""
+    conn = connect()
+
+    pq3.handshake(conn, user=pq3.pguser(), database=pq3.pgdatabase())
+
+    pq3.send(conn, pq3.types.Query, query=b"")
+
+    resp = pq3.recv1(conn)
+    assert resp.type == pq3.types.EmptyQueryResponse
+
+    resp = pq3.recv1(conn)
+    assert resp.type == pq3.types.ReadyForQuery
diff --git a/src/test/python/server/validate_bearer.py b/src/test/python/server/validate_bearer.py
new file mode 100755
index 0000000000..2cc73ff154
--- /dev/null
+++ b/src/test/python/server/validate_bearer.py
@@ -0,0 +1,101 @@
+#! /usr/bin/env python3
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+# DO NOT USE THIS OAUTH VALIDATOR IN PRODUCTION. It doesn't actually validate
+# anything, and it logs the bearer token data, which is sensitive.
+#
+# This executable is used as an oauth_validator_command in concert with
+# test_oauth.py. Memory is shared and communicated from that test module's
+# bearer_token() fixture.
+#
+# This script must run under the Postgres server environment; keep the
+# dependency list fairly standard.
+
+import base64
+import binascii
+import contextlib
+import struct
+import sys
+from multiprocessing import shared_memory
+
+MAX_UINT16 = 2 ** 16 - 1
+
+
+def remove_shm_from_resource_tracker():
+    """
+    Monkey-patch multiprocessing.resource_tracker so SharedMemory won't be
+    tracked. Pulled from this thread, where there are more details:
+
+        https://bugs.python.org/issue38119
+
+    TL;DR: all clients of shared memory segments automatically destroy them on
+    process exit, which makes shared memory segments much less useful. This
+    monkeypatch removes that behavior so that we can defer to the test to manage
+    the segment lifetime.
+
+    Ideally a future Python patch will pull in this fix and then the entire
+    function can go away.
+    """
+    from multiprocessing import resource_tracker
+
+    def fix_register(name, rtype):
+        if rtype == "shared_memory":
+            return
+        return resource_tracker._resource_tracker.register(self, name, rtype)
+
+    resource_tracker.register = fix_register
+
+    def fix_unregister(name, rtype):
+        if rtype == "shared_memory":
+            return
+        return resource_tracker._resource_tracker.unregister(self, name, rtype)
+
+    resource_tracker.unregister = fix_unregister
+
+    if "shared_memory" in resource_tracker._CLEANUP_FUNCS:
+        del resource_tracker._CLEANUP_FUNCS["shared_memory"]
+
+
+def main(args):
+    remove_shm_from_resource_tracker()  # XXX remove some day
+
+    # Get the expected token from the currently running test.
+    shared_mem_name = args[0]
+
+    mem = shared_memory.SharedMemory(shared_mem_name)
+    with contextlib.closing(mem):
+        # First two bytes are the token length.
+        size = struct.unpack("H", mem.buf[:2])[0]
+
+        if size == MAX_UINT16:
+            # Special case: the test wants us to accept any token.
+            sys.stderr.write("accepting token without validation\n")
+            return
+
+        # The remainder of the buffer contains the expected token.
+        assert size <= (mem.size - 2)
+        expected_token = mem.buf[2 : size + 2].tobytes()
+
+        mem.buf[:] = b"\0" * mem.size  # scribble over the token
+
+    token = sys.stdin.buffer.read()
+    if token != expected_token:
+        sys.exit(f"failed to match Bearer token ({token!r} != {expected_token!r})")
+
+    # See if the test wants us to print anything. If so, it will have encoded
+    # the desired output in the token with an "output=" prefix.
+    try:
+        # altchars="-_" corresponds to the urlsafe alphabet.
+        data = base64.b64decode(token, altchars="-_", validate=True)
+
+        if data.startswith(b"output="):
+            sys.stdout.buffer.write(data[7:])
+
+    except binascii.Error:
+        pass
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])
diff --git a/src/test/python/server/validate_reflect.py b/src/test/python/server/validate_reflect.py
new file mode 100755
index 0000000000..24c3a7e715
--- /dev/null
+++ b/src/test/python/server/validate_reflect.py
@@ -0,0 +1,34 @@
+#! /usr/bin/env python3
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+# DO NOT USE THIS OAUTH VALIDATOR IN PRODUCTION. It ignores the bearer token
+# entirely and automatically logs the user in.
+#
+# This executable is used as an oauth_validator_command in concert with
+# test_oauth.py. It expects the user's desired role name as an argument; the
+# actual token will be discarded and the user will be logged in with the role
+# name as the authenticated identity.
+#
+# This script must run under the Postgres server environment; keep the
+# dependency list fairly standard.
+
+import sys
+
+
+def main(args):
+    # We have to read the entire token as our first action to unblock the
+    # server, but we won't actually use it.
+    _ = sys.stdin.buffer.read()
+
+    if len(args) != 1:
+        sys.exit("usage: ./validate_reflect.py ROLE")
+
+    # Log the user in as the provided role.
+    role = args[0]
+    print(role)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])
diff --git a/src/test/python/test_internals.py b/src/test/python/test_internals.py
new file mode 100644
index 0000000000..dee4855fc0
--- /dev/null
+++ b/src/test/python/test_internals.py
@@ -0,0 +1,138 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import io
+
+from pq3 import _DebugStream
+
+
+def test_DebugStream_read():
+    under = io.BytesIO(b"abcdefghijklmnopqrstuvwxyz")
+    out = io.StringIO()
+
+    stream = _DebugStream(under, out)
+
+    res = stream.read(5)
+    assert res == b"abcde"
+
+    res = stream.read(16)
+    assert res == b"fghijklmnopqrstu"
+
+    stream.flush_debug()
+
+    res = stream.read()
+    assert res == b"vwxyz"
+
+    stream.flush_debug()
+
+    expected = (
+        "< 0000:\t61 62 63 64 65 66 67 68 69 6a 6b 6c 6d 6e 6f 70\tabcdefghijklmnop\n"
+        "< 0010:\t71 72 73 74 75                                 \tqrstu\n"
+        "\n"
+        "< 0000:\t76 77 78 79 7a                                 \tvwxyz\n"
+        "\n"
+    )
+    assert out.getvalue() == expected
+
+
+def test_DebugStream_write():
+    under = io.BytesIO()
+    out = io.StringIO()
+
+    stream = _DebugStream(under, out)
+
+    stream.write(b"\x00\x01\x02")
+    stream.flush()
+
+    assert under.getvalue() == b"\x00\x01\x02"
+
+    stream.write(b"\xc0\xc1\xc2")
+    stream.flush()
+
+    assert under.getvalue() == b"\x00\x01\x02\xc0\xc1\xc2"
+
+    stream.flush_debug()
+
+    expected = "> 0000:\t00 01 02 c0 c1 c2                              \t......\n\n"
+    assert out.getvalue() == expected
+
+
+def test_DebugStream_read_write():
+    under = io.BytesIO(b"abcdefghijklmnopqrstuvwxyz")
+    out = io.StringIO()
+    stream = _DebugStream(under, out)
+
+    res = stream.read(5)
+    assert res == b"abcde"
+
+    stream.write(b"xxxxx")
+    stream.flush()
+
+    assert under.getvalue() == b"abcdexxxxxklmnopqrstuvwxyz"
+
+    res = stream.read(5)
+    assert res == b"klmno"
+
+    stream.write(b"xxxxx")
+    stream.flush()
+
+    assert under.getvalue() == b"abcdexxxxxklmnoxxxxxuvwxyz"
+
+    stream.flush_debug()
+
+    expected = (
+        "< 0000:\t61 62 63 64 65 6b 6c 6d 6e 6f                  \tabcdeklmno\n"
+        "\n"
+        "> 0000:\t78 78 78 78 78 78 78 78 78 78                  \txxxxxxxxxx\n"
+        "\n"
+    )
+    assert out.getvalue() == expected
+
+
+def test_DebugStream_end_packet():
+    under = io.BytesIO(b"abcdefghijklmnopqrstuvwxyz")
+    out = io.StringIO()
+    stream = _DebugStream(under, out)
+
+    stream.read(5)
+    stream.end_packet("read description", read=True, indent=" ")
+
+    stream.write(b"xxxxx")
+    stream.flush()
+    stream.end_packet("write description", indent=" ")
+
+    stream.read(5)
+    stream.write(b"xxxxx")
+    stream.flush()
+    stream.end_packet("read/write combo for read", read=True, indent=" ")
+
+    stream.read(5)
+    stream.write(b"xxxxx")
+    stream.flush()
+    stream.end_packet("read/write combo for write", indent=" ")
+
+    expected = (
+        " < 0000:\t61 62 63 64 65                                 \tabcde\n"
+        "\n"
+        "< read description\n"
+        "\n"
+        "> write description\n"
+        "\n"
+        " > 0000:\t78 78 78 78 78                                 \txxxxx\n"
+        "\n"
+        " < 0000:\t6b 6c 6d 6e 6f                                 \tklmno\n"
+        "\n"
+        " > 0000:\t78 78 78 78 78                                 \txxxxx\n"
+        "\n"
+        "< read/write combo for read\n"
+        "\n"
+        "> read/write combo for write\n"
+        "\n"
+        " < 0000:\t75 76 77 78 79                                 \tuvwxy\n"
+        "\n"
+        " > 0000:\t78 78 78 78 78                                 \txxxxx\n"
+        "\n"
+    )
+    assert out.getvalue() == expected
diff --git a/src/test/python/test_pq3.py b/src/test/python/test_pq3.py
new file mode 100644
index 0000000000..e0c0e0568d
--- /dev/null
+++ b/src/test/python/test_pq3.py
@@ -0,0 +1,558 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+import contextlib
+import getpass
+import io
+import struct
+import sys
+
+import pytest
+from construct import Container, PaddingError, StreamError, TerminatedError
+
+import pq3
+
+
[email protected](
+    "raw,expected,extra",
+    [
+        pytest.param(
+            b"\x00\x00\x00\x10\x00\x04\x00\x00abcdefgh",
+            Container(len=16, proto=0x40000, payload=b"abcdefgh"),
+            b"",
+            id="8-byte payload",
+        ),
+        pytest.param(
+            b"\x00\x00\x00\x08\x00\x04\x00\x00",
+            Container(len=8, proto=0x40000, payload=b""),
+            b"",
+            id="no payload",
+        ),
+        pytest.param(
+            b"\x00\x00\x00\x09\x00\x04\x00\x00abcde",
+            Container(len=9, proto=0x40000, payload=b"a"),
+            b"bcde",
+            id="1-byte payload and extra padding",
+        ),
+        pytest.param(
+            b"\x00\x00\x00\x0B\x00\x03\x00\x00hi\x00",
+            Container(len=11, proto=pq3.protocol(3, 0), payload=[b"hi"]),
+            b"",
+            id="implied parameter list when using proto version 3.0",
+        ),
+    ],
+)
+def test_Startup_parse(raw, expected, extra):
+    with io.BytesIO(raw) as stream:
+        actual = pq3.Startup.parse_stream(stream)
+
+        assert actual == expected
+        assert stream.read() == extra
+
+
[email protected](
+    "packet,expected_bytes",
+    [
+        pytest.param(
+            dict(),
+            b"\x00\x00\x00\x08\x00\x00\x00\x00",
+            id="nothing set",
+        ),
+        pytest.param(
+            dict(len=10, proto=0x12345678),
+            b"\x00\x00\x00\x0A\x12\x34\x56\x78\x00\x00",
+            id="len and proto set explicitly",
+        ),
+        pytest.param(
+            dict(proto=0x12345678),
+            b"\x00\x00\x00\x08\x12\x34\x56\x78",
+            id="implied len with no payload",
+        ),
+        pytest.param(
+            dict(proto=0x12345678, payload=b"abcd"),
+            b"\x00\x00\x00\x0C\x12\x34\x56\x78abcd",
+            id="implied len with payload",
+        ),
+        pytest.param(
+            dict(payload=[b""]),
+            b"\x00\x00\x00\x09\x00\x03\x00\x00\x00",
+            id="implied proto version 3 when sending parameters",
+        ),
+        pytest.param(
+            dict(payload=[b"hi", b""]),
+            b"\x00\x00\x00\x0C\x00\x03\x00\x00hi\x00\x00",
+            id="implied proto version 3 and len when sending more than one parameter",
+        ),
+        pytest.param(
+            dict(payload=dict(user="jsmith", database="postgres")),
+            b"\x00\x00\x00\x27\x00\x03\x00\x00user\x00jsmith\x00database\x00postgres\x00\x00",
+            id="auto-serialization of dict parameters",
+        ),
+    ],
+)
+def test_Startup_build(packet, expected_bytes):
+    actual = pq3.Startup.build(packet)
+    assert actual == expected_bytes
+
+
[email protected](
+    "raw,expected,extra",
+    [
+        pytest.param(
+            b"*\x00\x00\x00\x08abcd",
+            dict(type=b"*", len=8, payload=b"abcd"),
+            b"",
+            id="4-byte payload",
+        ),
+        pytest.param(
+            b"*\x00\x00\x00\x04",
+            dict(type=b"*", len=4, payload=b""),
+            b"",
+            id="no payload",
+        ),
+        pytest.param(
+            b"*\x00\x00\x00\x05xabcd",
+            dict(type=b"*", len=5, payload=b"x"),
+            b"abcd",
+            id="1-byte payload with extra padding",
+        ),
+        pytest.param(
+            b"R\x00\x00\x00\x08\x00\x00\x00\x00",
+            dict(
+                type=pq3.types.AuthnRequest,
+                len=8,
+                payload=dict(type=pq3.authn.OK, body=None),
+            ),
+            b"",
+            id="AuthenticationOk",
+        ),
+        pytest.param(
+            b"R\x00\x00\x00\x12\x00\x00\x00\x0AEXTERNAL\x00\x00",
+            dict(
+                type=pq3.types.AuthnRequest,
+                len=18,
+                payload=dict(type=pq3.authn.SASL, body=[b"EXTERNAL", b""]),
+            ),
+            b"",
+            id="AuthenticationSASL",
+        ),
+        pytest.param(
+            b"R\x00\x00\x00\x0D\x00\x00\x00\x0B12345",
+            dict(
+                type=pq3.types.AuthnRequest,
+                len=13,
+                payload=dict(type=pq3.authn.SASLContinue, body=b"12345"),
+            ),
+            b"",
+            id="AuthenticationSASLContinue",
+        ),
+        pytest.param(
+            b"R\x00\x00\x00\x0D\x00\x00\x00\x0C12345",
+            dict(
+                type=pq3.types.AuthnRequest,
+                len=13,
+                payload=dict(type=pq3.authn.SASLFinal, body=b"12345"),
+            ),
+            b"",
+            id="AuthenticationSASLFinal",
+        ),
+        pytest.param(
+            b"p\x00\x00\x00\x0Bhunter2",
+            dict(
+                type=pq3.types.PasswordMessage,
+                len=11,
+                payload=b"hunter2",
+            ),
+            b"",
+            id="PasswordMessage",
+        ),
+        pytest.param(
+            b"K\x00\x00\x00\x0C\x00\x00\x00\x00\x12\x34\x56\x78",
+            dict(
+                type=pq3.types.BackendKeyData,
+                len=12,
+                payload=dict(pid=0, key=0x12345678),
+            ),
+            b"",
+            id="BackendKeyData",
+        ),
+        pytest.param(
+            b"C\x00\x00\x00\x08SET\x00",
+            dict(
+                type=pq3.types.CommandComplete,
+                len=8,
+                payload=dict(tag=b"SET"),
+            ),
+            b"",
+            id="CommandComplete",
+        ),
+        pytest.param(
+            b"E\x00\x00\x00\x11Mbad!\x00Mdog!\x00\x00",
+            dict(type=b"E", len=17, payload=dict(fields=[b"Mbad!", b"Mdog!", b""])),
+            b"",
+            id="ErrorResponse",
+        ),
+        pytest.param(
+            b"S\x00\x00\x00\x08a\x00b\x00",
+            dict(
+                type=pq3.types.ParameterStatus,
+                len=8,
+                payload=dict(name=b"a", value=b"b"),
+            ),
+            b"",
+            id="ParameterStatus",
+        ),
+        pytest.param(
+            b"Z\x00\x00\x00\x05x",
+            dict(type=b"Z", len=5, payload=dict(status=b"x")),
+            b"",
+            id="ReadyForQuery",
+        ),
+        pytest.param(
+            b"Q\x00\x00\x00\x06!\x00",
+            dict(type=pq3.types.Query, len=6, payload=dict(query=b"!")),
+            b"",
+            id="Query",
+        ),
+        pytest.param(
+            b"D\x00\x00\x00\x0B\x00\x01\x00\x00\x00\x01!",
+            dict(type=pq3.types.DataRow, len=11, payload=dict(columns=[b"!"])),
+            b"",
+            id="DataRow",
+        ),
+        pytest.param(
+            b"D\x00\x00\x00\x06\x00\x00extra",
+            dict(type=pq3.types.DataRow, len=6, payload=dict(columns=[])),
+            b"extra",
+            id="DataRow with extra data",
+        ),
+        pytest.param(
+            b"I\x00\x00\x00\x04",
+            dict(type=pq3.types.EmptyQueryResponse, len=4, payload=None),
+            b"",
+            id="EmptyQueryResponse",
+        ),
+        pytest.param(
+            b"I\x00\x00\x00\x04\xFF",
+            dict(type=b"I", len=4, payload=None),
+            b"\xFF",
+            id="EmptyQueryResponse with extra bytes",
+        ),
+        pytest.param(
+            b"X\x00\x00\x00\x04",
+            dict(type=pq3.types.Terminate, len=4, payload=None),
+            b"",
+            id="Terminate",
+        ),
+    ],
+)
+def test_Pq3_parse(raw, expected, extra):
+    with io.BytesIO(raw) as stream:
+        actual = pq3.Pq3.parse_stream(stream)
+
+        assert actual == expected
+        assert stream.read() == extra
+
+
[email protected](
+    "fields,expected",
+    [
+        pytest.param(
+            dict(type=b"*", len=5),
+            b"*\x00\x00\x00\x05\x00",
+            id="type and len set explicitly",
+        ),
+        pytest.param(
+            dict(type=b"*"),
+            b"*\x00\x00\x00\x04",
+            id="implied len with no payload",
+        ),
+        pytest.param(
+            dict(type=b"*", payload=b"1234"),
+            b"*\x00\x00\x00\x081234",
+            id="implied len with payload",
+        ),
+        pytest.param(
+            dict(type=pq3.types.AuthnRequest, payload=dict(type=pq3.authn.OK)),
+            b"R\x00\x00\x00\x08\x00\x00\x00\x00",
+            id="implied len/type for AuthenticationOK",
+        ),
+        pytest.param(
+            dict(
+                type=pq3.types.AuthnRequest,
+                payload=dict(
+                    type=pq3.authn.SASL,
+                    body=[b"SCRAM-SHA-256-PLUS", b"SCRAM-SHA-256", b""],
+                ),
+            ),
+            b"R\x00\x00\x00\x2A\x00\x00\x00\x0ASCRAM-SHA-256-PLUS\x00SCRAM-SHA-256\x00\x00",
+            id="implied len/type for AuthenticationSASL",
+        ),
+        pytest.param(
+            dict(
+                type=pq3.types.AuthnRequest,
+                payload=dict(type=pq3.authn.SASLContinue, body=b"12345"),
+            ),
+            b"R\x00\x00\x00\x0D\x00\x00\x00\x0B12345",
+            id="implied len/type for AuthenticationSASLContinue",
+        ),
+        pytest.param(
+            dict(
+                type=pq3.types.AuthnRequest,
+                payload=dict(type=pq3.authn.SASLFinal, body=b"12345"),
+            ),
+            b"R\x00\x00\x00\x0D\x00\x00\x00\x0C12345",
+            id="implied len/type for AuthenticationSASLFinal",
+        ),
+        pytest.param(
+            dict(
+                type=pq3.types.PasswordMessage,
+                payload=b"hunter2",
+            ),
+            b"p\x00\x00\x00\x0Bhunter2",
+            id="implied len/type for PasswordMessage",
+        ),
+        pytest.param(
+            dict(type=pq3.types.BackendKeyData, payload=dict(pid=1, key=7)),
+            b"K\x00\x00\x00\x0C\x00\x00\x00\x01\x00\x00\x00\x07",
+            id="implied len/type for BackendKeyData",
+        ),
+        pytest.param(
+            dict(type=pq3.types.CommandComplete, payload=dict(tag=b"SET")),
+            b"C\x00\x00\x00\x08SET\x00",
+            id="implied len/type for CommandComplete",
+        ),
+        pytest.param(
+            dict(type=pq3.types.ErrorResponse, payload=dict(fields=[b"error", b""])),
+            b"E\x00\x00\x00\x0Berror\x00\x00",
+            id="implied len/type for ErrorResponse",
+        ),
+        pytest.param(
+            dict(type=pq3.types.ParameterStatus, payload=dict(name=b"a", value=b"b")),
+            b"S\x00\x00\x00\x08a\x00b\x00",
+            id="implied len/type for ParameterStatus",
+        ),
+        pytest.param(
+            dict(type=pq3.types.ReadyForQuery, payload=dict(status=b"I")),
+            b"Z\x00\x00\x00\x05I",
+            id="implied len/type for ReadyForQuery",
+        ),
+        pytest.param(
+            dict(type=pq3.types.Query, payload=dict(query=b"SELECT 1;")),
+            b"Q\x00\x00\x00\x0eSELECT 1;\x00",
+            id="implied len/type for Query",
+        ),
+        pytest.param(
+            dict(type=pq3.types.DataRow, payload=dict(columns=[b"abcd"])),
+            b"D\x00\x00\x00\x0E\x00\x01\x00\x00\x00\x04abcd",
+            id="implied len/type for DataRow",
+        ),
+        pytest.param(
+            dict(type=pq3.types.EmptyQueryResponse),
+            b"I\x00\x00\x00\x04",
+            id="implied len for EmptyQueryResponse",
+        ),
+        pytest.param(
+            dict(type=pq3.types.Terminate),
+            b"X\x00\x00\x00\x04",
+            id="implied len for Terminate",
+        ),
+    ],
+)
+def test_Pq3_build(fields, expected):
+    actual = pq3.Pq3.build(fields)
+    assert actual == expected
+
+
[email protected](
+    "raw,expected,extra",
+    [
+        pytest.param(
+            b"\x00\x00",
+            dict(columns=[]),
+            b"",
+            id="no columns",
+        ),
+        pytest.param(
+            b"\x00\x01\x00\x00\x00\x04abcd",
+            dict(columns=[b"abcd"]),
+            b"",
+            id="one column",
+        ),
+        pytest.param(
+            b"\x00\x02\x00\x00\x00\x04abcd\x00\x00\x00\x01x",
+            dict(columns=[b"abcd", b"x"]),
+            b"",
+            id="multiple columns",
+        ),
+        pytest.param(
+            b"\x00\x02\x00\x00\x00\x00\x00\x00\x00\x01x",
+            dict(columns=[b"", b"x"]),
+            b"",
+            id="empty column value",
+        ),
+        pytest.param(
+            b"\x00\x02\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF",
+            dict(columns=[None, None]),
+            b"",
+            id="null columns",
+        ),
+    ],
+)
+def test_DataRow_parse(raw, expected, extra):
+    pkt = b"D" + struct.pack("!i", len(raw) + 4) + raw
+    with io.BytesIO(pkt) as stream:
+        actual = pq3.Pq3.parse_stream(stream)
+
+        assert actual.type == pq3.types.DataRow
+        assert actual.payload == expected
+        assert stream.read() == extra
+
+
[email protected](
+    "fields,expected",
+    [
+        pytest.param(
+            dict(),
+            b"\x00\x00",
+            id="no columns",
+        ),
+        pytest.param(
+            dict(columns=[None, None]),
+            b"\x00\x02\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF",
+            id="null columns",
+        ),
+    ],
+)
+def test_DataRow_build(fields, expected):
+    actual = pq3.Pq3.build(dict(type=pq3.types.DataRow, payload=fields))
+
+    expected = b"D" + struct.pack("!i", len(expected) + 4) + expected
+    assert actual == expected
+
+
[email protected](
+    "raw,expected,exception",
+    [
+        pytest.param(
+            b"EXTERNAL\x00\xFF\xFF\xFF\xFF",
+            dict(name=b"EXTERNAL", len=-1, data=None),
+            None,
+            id="no initial response",
+        ),
+        pytest.param(
+            b"EXTERNAL\x00\x00\x00\x00\x02me",
+            dict(name=b"EXTERNAL", len=2, data=b"me"),
+            None,
+            id="initial response",
+        ),
+        pytest.param(
+            b"EXTERNAL\x00\x00\x00\x00\x02meextra",
+            None,
+            TerminatedError,
+            id="extra data",
+        ),
+        pytest.param(
+            b"EXTERNAL\x00\x00\x00\x00\xFFme",
+            None,
+            StreamError,
+            id="underflow",
+        ),
+    ],
+)
+def test_SASLInitialResponse_parse(raw, expected, exception):
+    ctx = contextlib.nullcontext()
+    if exception:
+        ctx = pytest.raises(exception)
+
+    with ctx:
+        actual = pq3.SASLInitialResponse.parse(raw)
+        assert actual == expected
+
+
[email protected](
+    "fields,expected",
+    [
+        pytest.param(
+            dict(name=b"EXTERNAL"),
+            b"EXTERNAL\x00\xFF\xFF\xFF\xFF",
+            id="no initial response",
+        ),
+        pytest.param(
+            dict(name=b"EXTERNAL", data=None),
+            b"EXTERNAL\x00\xFF\xFF\xFF\xFF",
+            id="no initial response (explicit None)",
+        ),
+        pytest.param(
+            dict(name=b"EXTERNAL", data=b""),
+            b"EXTERNAL\x00\x00\x00\x00\x00",
+            id="empty response",
+        ),
+        pytest.param(
+            dict(name=b"EXTERNAL", data=b"[email protected]"),
+            b"EXTERNAL\x00\x00\x00\x00\[email protected]",
+            id="initial response",
+        ),
+        pytest.param(
+            dict(name=b"EXTERNAL", len=2, data=b"[email protected]"),
+            b"EXTERNAL\x00\x00\x00\x00\[email protected]",
+            id="data overflow",
+        ),
+        pytest.param(
+            dict(name=b"EXTERNAL", len=14, data=b"me"),
+            b"EXTERNAL\x00\x00\x00\x00\x0Eme",
+            id="data underflow",
+        ),
+    ],
+)
+def test_SASLInitialResponse_build(fields, expected):
+    actual = pq3.SASLInitialResponse.build(fields)
+    assert actual == expected
+
+
[email protected](
+    "version,expected_bytes",
+    [
+        pytest.param((3, 0), b"\x00\x03\x00\x00", id="version 3"),
+        pytest.param((1234, 5679), b"\x04\xd2\x16\x2f", id="SSLRequest"),
+    ],
+)
+def test_protocol(version, expected_bytes):
+    # Make sure the integer returned by protocol is correctly serialized on the
+    # wire.
+    assert struct.pack("!i", pq3.protocol(*version)) == expected_bytes
+
+
[email protected](
+    "envvar,func,expected",
+    [
+        ("PGHOST", pq3.pghost, "localhost"),
+        ("PGPORT", pq3.pgport, 5432),
+        ("PGUSER", pq3.pguser, getpass.getuser()),
+        ("PGDATABASE", pq3.pgdatabase, "postgres"),
+    ],
+)
+def test_env_defaults(monkeypatch, envvar, func, expected):
+    monkeypatch.delenv(envvar, raising=False)
+
+    actual = func()
+    assert actual == expected
+
+
[email protected](
+    "envvars,func,expected",
+    [
+        (dict(PGHOST="otherhost"), pq3.pghost, "otherhost"),
+        (dict(PGPORT="6789"), pq3.pgport, 6789),
+        (dict(PGUSER="postgres"), pq3.pguser, "postgres"),
+        (dict(PGDATABASE="template1"), pq3.pgdatabase, "template1"),
+    ],
+)
+def test_env(monkeypatch, envvars, func, expected):
+    for k, v in envvars.items():
+        monkeypatch.setenv(k, v)
+
+    actual = func()
+    assert actual == expected
diff --git a/src/test/python/tls.py b/src/test/python/tls.py
new file mode 100644
index 0000000000..075c02c1ca
--- /dev/null
+++ b/src/test/python/tls.py
@@ -0,0 +1,195 @@
+#
+# Copyright 2021 VMware, Inc.
+# SPDX-License-Identifier: PostgreSQL
+#
+
+from construct import *
+
+#
+# TLS 1.3
+#
+# Most of the types below are transcribed from RFC 8446:
+#
+#     https://tools.ietf.org/html/rfc8446
+#
+
+
+def _Vector(size_field, element):
+    return Prefixed(size_field, GreedyRange(element))
+
+
+# Alerts
+
+AlertLevel = Enum(
+    Byte,
+    warning=1,
+    fatal=2,
+)
+
+AlertDescription = Enum(
+    Byte,
+    close_notify=0,
+    unexpected_message=10,
+    bad_record_mac=20,
+    decryption_failed_RESERVED=21,
+    record_overflow=22,
+    decompression_failure=30,
+    handshake_failure=40,
+    no_certificate_RESERVED=41,
+    bad_certificate=42,
+    unsupported_certificate=43,
+    certificate_revoked=44,
+    certificate_expired=45,
+    certificate_unknown=46,
+    illegal_parameter=47,
+    unknown_ca=48,
+    access_denied=49,
+    decode_error=50,
+    decrypt_error=51,
+    export_restriction_RESERVED=60,
+    protocol_version=70,
+    insufficient_security=71,
+    internal_error=80,
+    user_canceled=90,
+    no_renegotiation=100,
+    unsupported_extension=110,
+)
+
+Alert = Struct(
+    "level" / AlertLevel,
+    "description" / AlertDescription,
+)
+
+
+# Extensions
+
+ExtensionType = Enum(
+    Int16ub,
+    server_name=0,
+    max_fragment_length=1,
+    status_request=5,
+    supported_groups=10,
+    signature_algorithms=13,
+    use_srtp=14,
+    heartbeat=15,
+    application_layer_protocol_negotiation=16,
+    signed_certificate_timestamp=18,
+    client_certificate_type=19,
+    server_certificate_type=20,
+    padding=21,
+    pre_shared_key=41,
+    early_data=42,
+    supported_versions=43,
+    cookie=44,
+    psk_key_exchange_modes=45,
+    certificate_authorities=47,
+    oid_filters=48,
+    post_handshake_auth=49,
+    signature_algorithms_cert=50,
+    key_share=51,
+)
+
+Extension = Struct(
+    "extension_type" / ExtensionType,
+    "extension_data" / Prefixed(Int16ub, GreedyBytes),
+)
+
+
+# ClientHello
+
+
+class _CipherSuiteAdapter(Adapter):
+    class _hextuple(tuple):
+        def __repr__(self):
+            return f"(0x{self[0]:02X}, 0x{self[1]:02X})"
+
+    def _encode(self, obj, context, path):
+        return bytes(obj)
+
+    def _decode(self, obj, context, path):
+        assert len(obj) == 2
+        return self._hextuple(obj)
+
+
+ProtocolVersion = Hex(Int16ub)
+
+Random = Hex(Bytes(32))
+
+CipherSuite = _CipherSuiteAdapter(Byte[2])
+
+ClientHello = Struct(
+    "legacy_version" / ProtocolVersion,
+    "random" / Random,
+    "legacy_session_id" / Prefixed(Byte, Hex(GreedyBytes)),
+    "cipher_suites" / _Vector(Int16ub, CipherSuite),
+    "legacy_compression_methods" / Prefixed(Byte, GreedyBytes),
+    "extensions" / _Vector(Int16ub, Extension),
+)
+
+# ServerHello
+
+ServerHello = Struct(
+    "legacy_version" / ProtocolVersion,
+    "random" / Random,
+    "legacy_session_id_echo" / Prefixed(Byte, Hex(GreedyBytes)),
+    "cipher_suite" / CipherSuite,
+    "legacy_compression_method" / Hex(Byte),
+    "extensions" / _Vector(Int16ub, Extension),
+)
+
+# Handshake
+
+HandshakeType = Enum(
+    Byte,
+    client_hello=1,
+    server_hello=2,
+    new_session_ticket=4,
+    end_of_early_data=5,
+    encrypted_extensions=8,
+    certificate=11,
+    certificate_request=13,
+    certificate_verify=15,
+    finished=20,
+    key_update=24,
+    message_hash=254,
+)
+
+Handshake = Struct(
+    "msg_type" / HandshakeType,
+    "length" / Int24ub,
+    "payload"
+    / Switch(
+        this.msg_type,
+        {
+            HandshakeType.client_hello: ClientHello,
+            HandshakeType.server_hello: ServerHello,
+            # HandshakeType.end_of_early_data: EndOfEarlyData,
+            # HandshakeType.encrypted_extensions: EncryptedExtensions,
+            # HandshakeType.certificate_request: CertificateRequest,
+            # HandshakeType.certificate: Certificate,
+            # HandshakeType.certificate_verify: CertificateVerify,
+            # HandshakeType.finished: Finished,
+            # HandshakeType.new_session_ticket: NewSessionTicket,
+            # HandshakeType.key_update: KeyUpdate,
+        },
+        default=FixedSized(this.length, GreedyBytes),
+    ),
+)
+
+# Records
+
+ContentType = Enum(
+    Byte,
+    invalid=0,
+    change_cipher_spec=20,
+    alert=21,
+    handshake=22,
+    application_data=23,
+)
+
+Plaintext = Struct(
+    "type" / ContentType,
+    "legacy_record_version" / ProtocolVersion,
+    "length" / Int16ub,
+    "fragment" / FixedSized(this.length, GreedyBytes),
+)
-- 
2.25.1



view thread (191+ messages)  latest in thread

reply

Reply instructions:

You may reply publicly to this message via plain-text email
using any one of the following methods:

* Reply to all the recipients using the --to and --cc options:
  reply via email

  To: [email protected]
  Cc: [email protected], [email protected], [email protected]
  Subject: Re: [PoC] Federated Authn/z with OAUTHBEARER
  In-Reply-To: <[email protected]>

* Save the following mbox file, import it into your mail client,
  and reply-to-all from there: mbox

This inbox is served by agora; see mirroring instructions
for how to clone and mirror all data and code used for this inbox