> On 20 Dec 2024, at 02:00, Jacob Champion <jacob.champ...@enterprisedb.com> 
> wrote:

Thanks for the new version, I was doing a v39 review but I'll roll that over
into a v40 review now.

As I was reading I was trying to identify parts can be broken out and committed
ahead of time.  This not only to trim down size, but mostly to shape the final
commit into a coherent single commit that brings a single functionality
utilizing existing APIs.  Basically I think we should keep generic
functionality out of the final commit and keep that focused on OAuth and the
required APIs and infra.

The async auth support seemed like a candidate to go in before the rest.  While
there won't be any consumers of it, it's also not limited to OAuth.  What do
you think about slicing that off and get in ahead of time?  I took a small stab
at separating out the generic bits (it includes the PG_MAX_AUTH_TOKEN_LENGTH
move as well which is unrelated, but could also be committed ahead of time)
along with some small tweaks on it.

--
Daniel Gustafsson

diff --git a/src/include/libpq/auth.h b/src/include/libpq/auth.h
index 227b41daf6..01bc3c7250 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 PGDLLIMPORT char *pg_krb_server_keyfile;
 extern PGDLLIMPORT bool pg_krb_caseins_users;
 extern PGDLLIMPORT bool pg_gss_accept_delegation;
diff --git a/src/interfaces/libpq/fe-auth-sasl.h 
b/src/interfaces/libpq/fe-auth-sasl.h
index 258bfd0564..b47011d077 100644
--- a/src/interfaces/libpq/fe-auth-sasl.h
+++ b/src/interfaces/libpq/fe-auth-sasl.h
@@ -30,6 +30,7 @@ typedef enum
        SASL_COMPLETE = 0,
        SASL_FAILED,
        SASL_CONTINUE,
+       SASL_ASYNC,
 } SASLStatus;
 
 /*
@@ -77,6 +78,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
@@ -101,12 +104,17 @@ typedef struct pg_fe_sasl_mech
         *
         *      SASL_CONTINUE:  The output buffer is filled with a client 
response.
         *                                      Additional server challenge is 
expected
+        *      SASL_ASYNC:             Some asynchronous processing external 
to the
+        *                                      connection needs to be done 
before a response can be
+        *                                      generated. The mechanism is 
responsible for setting up
+        *                                      conn->async_auth appropriately 
before returning.
         *      SASL_COMPLETE:  The SASL exchange has completed successfully.
         *      SASL_FAILED:    The exchange has failed and the connection 
should be
         *                                      dropped.
         *--------
         */
-       SASLStatus      (*exchange) (void *state, char *input, int inputlen,
+       SASLStatus      (*exchange) (void *state, bool final,
+                                                        char *input, int 
inputlen,
                                                         char **output, int 
*outputlen);
 
        /*--------
diff --git a/src/interfaces/libpq/fe-auth-scram.c 
b/src/interfaces/libpq/fe-auth-scram.c
index 0bb820e0d9..da168eb2f5 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 SASLStatus scram_exchange(void *opaq, char *input, int inputlen,
+static SASLStatus scram_exchange(void *opaq, bool final,
+                                                                char *input, 
int inputlen,
                                                                 char **output, 
int *outputlen);
 static bool scram_channel_bound(void *opaq);
 static void scram_free(void *opaq);
@@ -202,7 +203,8 @@ scram_free(void *opaq)
  * Exchange a SCRAM message with backend.
  */
 static SASLStatus
-scram_exchange(void *opaq, char *input, int inputlen,
+scram_exchange(void *opaq, bool final,
+                          char *input, int inputlen,
                           char **output, int *outputlen)
 {
        fe_scram_state *state = (fe_scram_state *) opaq;
diff --git a/src/interfaces/libpq/fe-auth.c b/src/interfaces/libpq/fe-auth.c
index 20d3427e94..8aa972fdfb 100644
--- a/src/interfaces/libpq/fe-auth.c
+++ b/src/interfaces/libpq/fe-auth.c
@@ -430,7 +430,7 @@ pg_SSPI_startup(PGconn *conn, int use_negotiate, int 
payloadlen)
  * Initialize SASL authentication exchange.
  */
 static int
-pg_SASL_init(PGconn *conn, int payloadlen)
+pg_SASL_init(PGconn *conn, int payloadlen, bool *async)
 {
        char       *initialresponse = NULL;
        int                     initialresponselen;
@@ -448,7 +448,7 @@ pg_SASL_init(PGconn *conn, int payloadlen)
                goto error;
        }
 
-       if (conn->sasl_state)
+       if (conn->sasl_state && !conn->async_auth)
        {
                libpq_append_conn_error(conn, "duplicate SASL authentication 
request");
                goto error;
@@ -578,26 +578,48 @@ pg_SASL_init(PGconn *conn, int payloadlen)
 
        Assert(conn->sasl);
 
-       /*
-        * Initialize the SASL state information with all the information 
gathered
-        * during the initial exchange.
-        *
-        * Note: Only tls-unique is supported for the moment.
-        */
-       conn->sasl_state = conn->sasl->init(conn,
-                                                                               
password,
-                                                                               
selected_mechanism);
        if (!conn->sasl_state)
-               goto oom_error;
+       {
+               /*
+                * Initialize the SASL state information with all the 
information
+                * gathered during the initial exchange.
+                *
+                * Note: Only tls-unique is supported for the moment.
+                */
+               conn->sasl_state = conn->sasl->init(conn,
+                                                                               
        password,
+                                                                               
        selected_mechanism);
+               if (!conn->sasl_state)
+                       goto oom_error;
+       }
+       else
+       {
+               /*
+                * This is only possible if we're returning from an async loop.
+                * Disconnect it now.
+                */
+               Assert(conn->async_auth);
+               conn->async_auth = NULL;
+       }
 
        /* Get the mechanism-specific Initial Client Response, if any */
-       status = conn->sasl->exchange(conn->sasl_state,
+       status = conn->sasl->exchange(conn->sasl_state, false,
                                                                  NULL, -1,
                                                                  
&initialresponse, &initialresponselen);
 
        if (status == SASL_FAILED)
                goto error;
 
+       if (status == SASL_ASYNC)
+       {
+               /*
+                * The mechanism should have set up the necessary callbacks; 
all we
+                * need to do is signal the caller.
+                */
+               *async = true;
+               return STATUS_OK;
+       }
+
        /*
         * Build a SASLInitialResponse message, and send it.
         */
@@ -642,7 +664,7 @@ oom_error:
  * the protocol.
  */
 static int
-pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
+pg_SASL_continue(PGconn *conn, int payloadlen, bool final, bool *async)
 {
        char       *output;
        int                     outputlen;
@@ -672,11 +694,21 @@ pg_SASL_continue(PGconn *conn, int payloadlen, bool final)
        /* For safety and convenience, ensure the buffer is NULL-terminated. */
        challenge[payloadlen] = '\0';
 
-       status = conn->sasl->exchange(conn->sasl_state,
+       status = conn->sasl->exchange(conn->sasl_state, final,
                                                                  challenge, 
payloadlen,
                                                                  &output, 
&outputlen);
        free(challenge);                        /* don't need the input anymore 
*/
 
+       if (status == SASL_ASYNC)
+       {
+               /*
+                * The mechanism should have set up the necessary callbacks; 
all we
+                * need to do is signal the caller.
+                */
+               *async = true;
+               return STATUS_OK;
+       }
+
        if (final && status == SASL_CONTINUE)
        {
                if (outputlen != 0)
@@ -984,12 +1016,18 @@ check_expected_areq(AuthRequest areq, PGconn *conn)
  * it. We are responsible for reading any remaining extra data, specific
  * to the authentication method. 'payloadlen' is the remaining length in
  * the message.
+ *
+ * If *async is set to true on return, the client doesn't yet have enough
+ * information to respond, and the caller must temporarily switch to
+ * conn->async_auth() to continue driving the exchange.
  */
 int
-pg_fe_sendauth(AuthRequest areq, int payloadlen, PGconn *conn)
+pg_fe_sendauth(AuthRequest areq, int payloadlen, PGconn *conn, bool *async)
 {
        int                     oldmsglen;
 
+       *async = false;
+
        if (!check_expected_areq(areq, conn))
                return STATUS_ERROR;
 
@@ -1147,7 +1185,7 @@ pg_fe_sendauth(AuthRequest areq, int payloadlen, PGconn 
*conn)
                         * The request contains the name (as assigned by IANA) 
of the
                         * authentication mechanism.
                         */
-                       if (pg_SASL_init(conn, payloadlen) != STATUS_OK)
+                       if (pg_SASL_init(conn, payloadlen, async) != STATUS_OK)
                        {
                                /* pg_SASL_init already set the error message */
                                return STATUS_ERROR;
@@ -1156,23 +1194,33 @@ pg_fe_sendauth(AuthRequest areq, int payloadlen, PGconn 
*conn)
 
                case AUTH_REQ_SASL_CONT:
                case AUTH_REQ_SASL_FIN:
-                       if (conn->sasl_state == NULL)
-                       {
-                               appendPQExpBufferStr(&conn->errorMessage,
-                                                                        
"fe_sendauth: invalid authentication request from server: AUTH_REQ_SASL_CONT 
without AUTH_REQ_SASL\n");
-                               return STATUS_ERROR;
-                       }
-                       oldmsglen = conn->errorMessage.len;
-                       if (pg_SASL_continue(conn, payloadlen,
-                                                                (areq == 
AUTH_REQ_SASL_FIN)) != STATUS_OK)
                        {
-                               /* Use this message if pg_SASL_continue didn't 
supply one */
-                               if (conn->errorMessage.len == oldmsglen)
+                               bool            final = false;
+
+                               if (conn->sasl_state == NULL)
+                               {
                                        
appendPQExpBufferStr(&conn->errorMessage,
-                                                                               
 "fe_sendauth: error in SASL authentication\n");
-                               return STATUS_ERROR;
+                                                                               
 "fe_sendauth: invalid authentication request from server: AUTH_REQ_SASL_CONT 
without AUTH_REQ_SASL\n");
+                                       return STATUS_ERROR;
+                               }
+                               oldmsglen = conn->errorMessage.len;
+
+                               if (areq == AUTH_REQ_SASL_FIN)
+                                       final = true;
+
+                               if (pg_SASL_continue(conn, payloadlen, final, 
async) != STATUS_OK)
+                               {
+                                       /*
+                                        * Append a generic error message 
unless pg_SASL_continue
+                                        * did set a more specific one already.
+                                        */
+                                       if (conn->errorMessage.len == oldmsglen)
+                                               
appendPQExpBufferStr(&conn->errorMessage,
+                                                                               
         "fe_sendauth: error in SASL authentication\n");
+                                       return STATUS_ERROR;
+                               }
+                               break;
                        }
-                       break;
 
                default:
                        libpq_append_conn_error(conn, "authentication method %u 
not supported", areq);
diff --git a/src/interfaces/libpq/fe-auth.h b/src/interfaces/libpq/fe-auth.h
index a18c508688..286b25c941 100644
--- a/src/interfaces/libpq/fe-auth.h
+++ b/src/interfaces/libpq/fe-auth.h
@@ -19,7 +19,8 @@
 
 
 /* Prototypes for functions in fe-auth.c */
-extern int     pg_fe_sendauth(AuthRequest areq, int payloadlen, PGconn *conn);
+extern int     pg_fe_sendauth(AuthRequest areq, int payloadlen, PGconn *conn,
+                                                  bool *async);
 extern char *pg_fe_getusername(uid_t user_id, PQExpBuffer errorMessage);
 extern char *pg_fe_getauthname(PQExpBuffer errorMessage);
 
diff --git a/src/interfaces/libpq/fe-connect.c 
b/src/interfaces/libpq/fe-connect.c
index aaf87e8e88..fb786a2feb 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -479,6 +479,14 @@ pqDropConnection(PGconn *conn, bool flushInput)
                closesocket(conn->sock);
        conn->sock = PGINVALID_SOCKET;
 
+       /*
+        * The altsock used for asynchronous authentication should be closed in
+        * case it's still open.
+        */
+       if (conn->altsock != PGINVALID_SOCKET)
+               closesocket(conn->altsock);
+       conn->altsock = PGINVALID_SOCKET;
+
        /* Optionally discard any unread data */
        if (flushInput)
                conn->inStart = conn->inCursor = conn->inEnd = 0;
@@ -2645,6 +2653,7 @@ PQconnectPoll(PGconn *conn)
                case CONNECTION_NEEDED:
                case CONNECTION_GSS_STARTUP:
                case CONNECTION_CHECK_TARGET:
+               case CONNECTION_AUTHENTICATING:
                        break;
 
                default:
@@ -3680,6 +3689,7 @@ keep_going:                                               
/* We will come back to here until there is
                                int                     avail;
                                AuthRequest areq;
                                int                     res;
+                               bool            async;
 
                                /*
                                 * Scan the message from current point (note 
that if we find
@@ -3868,7 +3878,17 @@ keep_going:                                              
/* We will come back to here until there is
                                 * Note that conn->pghost must be non-NULL if 
we are going to
                                 * avoid the Kerberos code doing a hostname 
look-up.
                                 */
-                               res = pg_fe_sendauth(areq, msgLength, conn);
+                               res = pg_fe_sendauth(areq, msgLength, conn, 
&async);
+
+                               if (async && (res == STATUS_OK))
+                               {
+                                       /*
+                                        * We'll come back later once we're 
ready to respond.
+                                        * Don't consume the request yet.
+                                        */
+                                       conn->status = 
CONNECTION_AUTHENTICATING;
+                                       goto keep_going;
+                               }
 
                                /*
                                 * OK, we have processed the message; mark data 
consumed.  We
@@ -3893,6 +3913,14 @@ keep_going:                                              
/* We will come back to here until there is
                                        /* We are done with authentication 
exchange */
                                        conn->status = CONNECTION_AUTH_OK;
 
+                                       /*
+                                        * Close the altsock which may have 
been used in case if
+                                        * this was an asynchronous 
authentication.
+                                        */
+                                       if (conn->altsock != PGINVALID_SOCKET)
+                                               closesocket(conn->altsock);
+                                       conn->altsock = PGINVALID_SOCKET;
+
                                        /*
                                         * Set asyncStatus so that PQgetResult 
will think that
                                         * what comes back next is the result 
of a query.  See
@@ -3905,6 +3933,43 @@ keep_going:                                              
/* We will come back to here until there is
                                goto keep_going;
                        }
 
+               case CONNECTION_AUTHENTICATING:
+                       {
+                               PostgresPollingStatusType status;
+                               pgsocket        altsock = PGINVALID_SOCKET;
+
+                               if (!conn->async_auth)
+                               {
+                                       /* programmer error; should not happen 
*/
+                                       libpq_append_conn_error(conn, "async 
authentication has no handler");
+                                       goto error_return;
+                               }
+
+                               status = conn->async_auth(conn, &altsock);
+
+                               if (status == PGRES_POLLING_FAILED)
+                                       goto error_return;
+
+                               if (status == PGRES_POLLING_OK)
+                               {
+                                       /*
+                                        * Reenter the authentication exchange 
with the server. We
+                                        * didn't consume the message that 
started external
+                                        * authentication, so it'll be 
reprocessed as if we just
+                                        * received it.
+                                        */
+                                       conn->status = 
CONNECTION_AWAITING_RESPONSE;
+                                       if (conn->altsock != PGINVALID_SOCKET)
+                                               closesocket(conn->altsock);
+                                       conn->altsock = PGINVALID_SOCKET;
+
+                                       goto keep_going;
+                               }
+
+                               conn->altsock = altsock;
+                               return status;
+                       }
+
                case CONNECTION_AUTH_OK:
                        {
                                /*
@@ -4586,6 +4651,7 @@ pqMakeEmptyPGconn(void)
        conn->verbosity = PQERRORS_DEFAULT;
        conn->show_context = PQSHOW_CONTEXT_ERRORS;
        conn->sock = PGINVALID_SOCKET;
+       conn->altsock = PGINVALID_SOCKET;
        conn->Pfdebug = NULL;
 
        /*
@@ -7227,6 +7293,8 @@ PQsocket(const PGconn *conn)
 {
        if (!conn)
                return -1;
+       if (conn->altsock != PGINVALID_SOCKET)
+               return conn->altsock;
        return (conn->sock != PGINVALID_SOCKET) ? conn->sock : -1;
 }
 
diff --git a/src/interfaces/libpq/fe-misc.c b/src/interfaces/libpq/fe-misc.c
index 928a47162d..e2ba483ea8 100644
--- a/src/interfaces/libpq/fe-misc.c
+++ b/src/interfaces/libpq/fe-misc.c
@@ -1056,10 +1056,13 @@ static int
 pqSocketCheck(PGconn *conn, int forRead, int forWrite, pg_usec_time_t end_time)
 {
        int                     result;
+       pgsocket        sock;
 
        if (!conn)
                return -1;
-       if (conn->sock == PGINVALID_SOCKET)
+
+       sock = (conn->altsock != PGINVALID_SOCKET) ? conn->altsock : conn->sock;
+       if (sock == PGINVALID_SOCKET)
        {
                libpq_append_conn_error(conn, "invalid socket");
                return -1;
@@ -1076,7 +1079,7 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, 
pg_usec_time_t end_time)
 
        /* We will retry as long as we get EINTR */
        do
-               result = PQsocketPoll(conn->sock, forRead, forWrite, end_time);
+               result = PQsocketPoll(sock, forRead, forWrite, end_time);
        while (result < 0 && SOCK_ERRNO == EINTR);
 
        if (result < 0)
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index 15012c770c..8d4c4ffd4f 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -103,6 +103,8 @@ typedef enum
        CONNECTION_CHECK_STANDBY,       /* Checking if server is in standby 
mode. */
        CONNECTION_ALLOCATED,           /* Waiting for connection attempt to be
                                                                 * started.  */
+       CONNECTION_AUTHENTICATING,      /* Authentication is in progress with 
some
+                                                                * external 
system. */
 } ConnStatusType;
 
 typedef enum
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 08cc391cbd..427b538c0b 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -506,6 +506,11 @@ struct pg_conn
                                                                                
 * know which auth response we're
                                                                                
 * sending */
 
+       /* Callback for external async authentication */
+       PostgresPollingStatusType (*async_auth) (PGconn *conn, pgsocket 
*altsock);
+       pgsocket        altsock;                /* alternative socket for 
client to poll */
+
+
        /* Transient state needed while establishing connection */
        PGTargetServerType target_server_type;  /* desired session properties */
        PGLoadBalanceType load_balance_type;    /* desired load balancing

Reply via email to