From 1cf48a8f83505a0cce1f94f8b0b563d4dcdd547a Mon Sep 17 00:00:00 2001
From: Jacob Champion <jacob.champion@enterprisedb.com>
Date: Wed, 5 Feb 2025 15:49:01 -0800
Subject: [PATCH v47 08/11] fixup! Add OAUTHBEARER SASL mechanism

---
 doc/src/sgml/libpq.sgml                       | 13 +++
 src/interfaces/libpq/fe-auth-oauth-curl.c     | 86 ++++++++++++++-----
 src/interfaces/libpq/libpq-fe.h               |  3 +
 .../modules/oauth_validator/t/oauth_server.py |  2 +-
 4 files changed, 80 insertions(+), 24 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index 9a69ffbc5b3..ddfc2a27c50 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -10228,6 +10228,9 @@ typedef struct _PGpromptOAuthDevice
 {
     const char *verification_uri;   /* verification URI to visit */
     const char *user_code;          /* user code to enter */
+    const char *verification_uri_complete;  /* optional combination of URI and
+                                             * code, or NULL */
+    int         expires_in;         /* seconds until user code expires */
 } PGpromptOAuthDevice;
 </synopsis>
         </para>
@@ -10246,6 +10249,16 @@ typedef struct _PGpromptOAuthDevice
          <link linkend="libpq-oauth-authdata-oauth-bearer-token">custom OAuth
          flow</link>, this authdata type will not be used.
         </para>
+        <para>
+         If a non-NULL <structfield>verification_uri_complete</structfield> is
+         provided, it may optionally be used for non-textual verification (for
+         example, by displaying a QR code). The URL and user code should still
+         be displayed to the end user in this case, because the code will be
+         manually confirmed by the provider, and the URL lets users continue
+         even if they can't use the non-textual method. Review the RFC's
+         <ulink url="https://datatracker.ietf.org/doc/html/rfc8628#section-3.3.1">notes
+         on non-textual verification</ulink>.
+        </para>
        </listitem>
       </varlistentry>
 
diff --git a/src/interfaces/libpq/fe-auth-oauth-curl.c b/src/interfaces/libpq/fe-auth-oauth-curl.c
index 993ca3bdab9..02c5b50afcd 100644
--- a/src/interfaces/libpq/fe-auth-oauth-curl.c
+++ b/src/interfaces/libpq/fe-auth-oauth-curl.c
@@ -76,9 +76,12 @@ struct device_authz
 	char	   *device_code;
 	char	   *user_code;
 	char	   *verification_uri;
+	char	   *verification_uri_complete;
+	char	   *expires_in_str;
 	char	   *interval_str;
 
 	/* Fields below are parsed from the corresponding string above. */
+	int			expires_in;
 	int			interval;
 };
 
@@ -88,6 +91,8 @@ free_device_authz(struct device_authz *authz)
 	free(authz->device_code);
 	free(authz->user_code);
 	free(authz->verification_uri);
+	free(authz->verification_uri_complete);
+	free(authz->expires_in_str);
 	free(authz->interval_str);
 }
 
@@ -853,20 +858,12 @@ parse_provider(struct async_ctx *actx, struct provider *provider)
 }
 
 /*
- * Parses the "interval" JSON number, corresponding to the number of seconds to
- * wait between token endpoint requests.
- *
- * RFC 8628 is pretty silent on sanity checks for the interval. As a matter of
- * practicality, round any fractional intervals up to the next second, and clamp
- * the result at a minimum of one. (Zero-second intervals would result in an
- * expensive network polling loop.) Tests may remove the lower bound with
- * PGOAUTHDEBUG, for improved performance.
- *
- * TODO: maybe clamp the upper bound too, based on the libpq timeout and/or the
- * code expiration time?
+ * Parses a valid JSON number into a double. The input must have come from
+ * pg_parse_json(), so that we know the lexer has validated it; there's no
+ * in-band signal for invalid formats.
  */
-static int
-parse_interval(struct async_ctx *actx, const char *interval_str)
+static double
+parse_json_number(const char *s)
 {
 	double		parsed;
 	int			cnt;
@@ -875,7 +872,7 @@ parse_interval(struct async_ctx *actx, const char *interval_str)
 	 * The JSON lexer has already validated the number, which is stricter than
 	 * the %f format, so we should be good to use sscanf().
 	 */
-	cnt = sscanf(interval_str, "%lf", &parsed);
+	cnt = sscanf(s, "%lf", &parsed);
 
 	if (cnt != 1)
 	{
@@ -884,9 +881,28 @@ parse_interval(struct async_ctx *actx, const char *interval_str)
 		 * either way a developer needs to take a look.
 		 */
 		Assert(cnt == 1);
-		return 1;				/* don't fall through in release builds */
+		return 0;
 	}
 
+	return parsed;
+}
+
+/*
+ * Parses the "interval" JSON number, corresponding to the number of seconds to
+ * wait between token endpoint requests.
+ *
+ * RFC 8628 is pretty silent on sanity checks for the interval. As a matter of
+ * practicality, round any fractional intervals up to the next second, and clamp
+ * the result at a minimum of one. (Zero-second intervals would result in an
+ * expensive network polling loop.) Tests may remove the lower bound with
+ * PGOAUTHDEBUG, for improved performance.
+ */
+static int
+parse_interval(struct async_ctx *actx, const char *interval_str)
+{
+	double		parsed;
+
+	parsed = parse_json_number(interval_str);
 	parsed = ceil(parsed);
 
 	if (parsed < 1)
@@ -898,6 +914,31 @@ parse_interval(struct async_ctx *actx, const char *interval_str)
 	return parsed;
 }
 
+/*
+ * Parses the "expires_in" JSON number, corresponding to the number of seconds
+ * remaining in the lifetime of the device code request.
+ *
+ * Similar to parse_interval, but we have even fewer requirements for reasonable
+ * values since we don't use the expiration time directly (it's passed to the
+ * PQAUTHDATA_PROMPT_OAUTH_DEVICE hook, in case the application wants to do
+ * something with it). We simply round and clamp to int range.
+ */
+static int
+parse_expires_in(struct async_ctx *actx, const char *expires_in_str)
+{
+	double		parsed;
+
+	parsed = parse_json_number(expires_in_str);
+	parsed = round(parsed);
+
+	if (INT_MAX <= parsed)
+		return INT_MAX;
+	else if (parsed <= INT_MIN)
+		return INT_MIN;
+
+	return parsed;
+}
+
 /*
  * Parses the Device Authorization Response (RFC 8628, Sec. 3.2).
  */
@@ -908,6 +949,7 @@ parse_device_authz(struct async_ctx *actx, struct device_authz *authz)
 		{"device_code", JSON_TOKEN_STRING, {&authz->device_code}, REQUIRED},
 		{"user_code", JSON_TOKEN_STRING, {&authz->user_code}, REQUIRED},
 		{"verification_uri", JSON_TOKEN_STRING, {&authz->verification_uri}, REQUIRED},
+		{"expires_in", JSON_TOKEN_NUMBER, {&authz->expires_in_str}, REQUIRED},
 
 		/*
 		 * Some services (Google, Azure) spell verification_uri differently.
@@ -915,13 +957,7 @@ parse_device_authz(struct async_ctx *actx, struct device_authz *authz)
 		 */
 		{"verification_url", JSON_TOKEN_STRING, {&authz->verification_uri}, REQUIRED},
 
-		/*
-		 * The following fields are technically REQUIRED, but we don't use
-		 * them anywhere yet:
-		 *
-		 * - expires_in
-		 */
-
+		{"verification_uri_complete", JSON_TOKEN_STRING, {&authz->verification_uri_complete}, OPTIONAL},
 		{"interval", JSON_TOKEN_NUMBER, {&authz->interval_str}, OPTIONAL},
 
 		{0},
@@ -945,6 +981,9 @@ parse_device_authz(struct async_ctx *actx, struct device_authz *authz)
 		authz->interval = 5;
 	}
 
+	Assert(authz->expires_in_str);	/* ensured by parse_oauth_json() */
+	authz->expires_in = parse_expires_in(actx, authz->expires_in_str);
+
 	return true;
 }
 
@@ -2301,7 +2340,8 @@ prompt_user(struct async_ctx *actx, PGconn *conn)
 	PGpromptOAuthDevice prompt = {
 		.verification_uri = actx->authz.verification_uri,
 		.user_code = actx->authz.user_code,
-		/* TODO: optional fields */
+		.verification_uri_complete = actx->authz.verification_uri_complete,
+		.expires_in = actx->authz.expires_in,
 	};
 
 	res = PQauthDataHook(PQAUTHDATA_PROMPT_OAUTH_DEVICE, conn, &prompt);
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 5f8d608261e..b7399dee58e 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -733,6 +733,9 @@ typedef struct _PGpromptOAuthDevice
 {
 	const char *verification_uri;	/* verification URI to visit */
 	const char *user_code;		/* user code to enter */
+	const char *verification_uri_complete;	/* optional combination of URI and
+											 * code, or NULL */
+	int			expires_in;		/* seconds until user code expires */
 } PGpromptOAuthDevice;
 
 /* for PGoauthBearerRequest.async() */
diff --git a/src/test/modules/oauth_validator/t/oauth_server.py b/src/test/modules/oauth_validator/t/oauth_server.py
index 8ec09102027..4faf3323d38 100755
--- a/src/test/modules/oauth_validator/t/oauth_server.py
+++ b/src/test/modules/oauth_validator/t/oauth_server.py
@@ -298,7 +298,7 @@ class OAuthHandler(http.server.BaseHTTPRequestHandler):
             "device_code": "postgres",
             "user_code": "postgresuser",
             self._uri_spelling: uri,
-            "expires-in": 5,
+            "expires_in": 5,
             **self._response_padding,
         }
 
-- 
2.34.1

