oauth_hook_client.c

Go to the documentation of this file.

/*-------------------------------------------------------------------------
 *
 * oauth_hook_client.c
 *      Test driver for t/002_client.pl, which verifies OAuth hook
 *      functionality in libpq.
 *
 * Portions Copyright (c) 1996-2026, PostgreSQL Global Development Group
 * Portions Copyright (c) 1994, Regents of the University of California
 *
 *
 * IDENTIFICATION
 *      src/test/modules/oauth_validator/oauth_hook_client.c
 *
 *-------------------------------------------------------------------------
 */
 
#include "postgres_fe.h"
 
#include <sys/socket.h>
 
#include "getopt_long.h"
#include "libpq-fe.h"
 
static int  handle_auth_data(PGauthData type, PGconn *conn, void *data);
static PostgresPollingStatusType async_cb(PGconn *conn,
                                          PGoauthBearerRequest *req,
                                          pgsocket *altsock);
static PostgresPollingStatusType misbehave_cb(PGconn *conn,
                                              PGoauthBearerRequest *req,
                                              pgsocket *altsock);
 
static void
usage(char *argv[])
{
    printf("usage: %s [flags] CONNINFO\n\n", argv[0]);
 
    printf("recognized flags:\n");
    printf("  -h, --help              show this message\n");
    printf("  --expected-scope SCOPE  fail if received scopes do not match SCOPE\n");
    printf("  --expected-uri URI      fail if received configuration link does not match URI\n");
    printf("  --misbehave=MODE        have the hook fail required postconditions\n"
           "                          (MODEs: no-hook, fail-async, no-token, no-socket)\n");
    printf("  --no-hook               don't install OAuth hooks\n");
    printf("  --hang-forever          don't ever return a token (combine with connect_timeout)\n");
    printf("  --token TOKEN           use the provided TOKEN value\n");
    printf("  --stress-async          busy-loop on PQconnectPoll rather than polling\n");
}
 
/* --options */
static bool no_hook = false;
static bool hang_forever = false;
static bool stress_async = false;
static const char *expected_uri = NULL;
static const char *expected_scope = NULL;
static const char *misbehave_mode = NULL;
static char *token = NULL;
 
int
main(int argc, char *argv[])
{
    static const struct option long_options[] = {
        {"help", no_argument, NULL, 'h'},
 
        {"expected-scope", required_argument, NULL, 1000},
        {"expected-uri", required_argument, NULL, 1001},
        {"no-hook", no_argument, NULL, 1002},
        {"token", required_argument, NULL, 1003},
        {"hang-forever", no_argument, NULL, 1004},
        {"misbehave", required_argument, NULL, 1005},
        {"stress-async", no_argument, NULL, 1006},
        {0}
    };
 
    const char *conninfo;
    PGconn     *conn;
    int         c;
 
    while ((c = getopt_long(argc, argv, "h", long_options, NULL)) != -1)
    {
        switch (c)
        {
            case 'h':
                usage(argv);
                return 0;
 
            case 1000:          /* --expected-scope */
                expected_scope = optarg;
                break;
 
            case 1001:          /* --expected-uri */
                expected_uri = optarg;
                break;
 
            case 1002:          /* --no-hook */
                no_hook = true;
                break;
 
            case 1003:          /* --token */
                token = optarg;
                break;
 
            case 1004:          /* --hang-forever */
                hang_forever = true;
                break;
 
            case 1005:          /* --misbehave */
                misbehave_mode = optarg;
                break;
 
            case 1006:          /* --stress-async */
                stress_async = true;
                break;
 
            default:
                usage(argv);
                return 1;
        }
    }
 
    if (argc != optind + 1)
    {
        usage(argv);
        return 1;
    }
 
    conninfo = argv[optind];
 
    /* Set up our OAuth hooks. */
    PQsetAuthDataHook(handle_auth_data);
 
    /* Connect. (All the actual work is in the hook.) */
    if (stress_async)
    {
        /*
         * Perform an asynchronous connection, busy-looping on PQconnectPoll()
         * without actually waiting on socket events. This stresses code paths
         * that rely on asynchronous work to be done before continuing with
         * the next step in the flow.
         */
        PostgresPollingStatusType res;
 
        conn = PQconnectStart(conninfo);
 
        do
        {
            res = PQconnectPoll(conn);
        } while (res != PGRES_POLLING_FAILED && res != PGRES_POLLING_OK);
    }
    else
    {
        /* Perform a standard synchronous connection. */
        conn = PQconnectdb(conninfo);
    }
 
    if (PQstatus(conn) != CONNECTION_OK)
    {
        fprintf(stderr, "connection to database failed: %s\n",
                PQerrorMessage(conn));
        PQfinish(conn);
        return 1;
    }
 
    printf("connection succeeded\n");
    PQfinish(conn);
    return 0;
}
 
/*
 * PQauthDataHook implementation. Replaces the default client flow by handling
 * PQAUTHDATA_OAUTH_BEARER_TOKEN.
 */
static int
handle_auth_data(PGauthData type, PGconn *conn, void *data)
{
    PGoauthBearerRequest *req = data;
 
    if (no_hook || (type != PQAUTHDATA_OAUTH_BEARER_TOKEN))
        return 0;
 
    if (hang_forever)
    {
        /* Start asynchronous processing. */
        req->async = async_cb;
        return 1;
    }
 
    if (misbehave_mode)
    {
        if (strcmp(misbehave_mode, "no-hook") != 0)
            req->async = misbehave_cb;
        return 1;
    }
 
    if (expected_uri)
    {
        if (!req->openid_configuration)
        {
            fprintf(stderr, "expected URI \"%s\", got NULL\n", expected_uri);
            return -1;
        }
 
        if (strcmp(expected_uri, req->openid_configuration) != 0)
        {
            fprintf(stderr, "expected URI \"%s\", got \"%s\"\n", expected_uri, req->openid_configuration);
            return -1;
        }
    }
 
    if (expected_scope)
    {
        if (!req->scope)
        {
            fprintf(stderr, "expected scope \"%s\", got NULL\n", expected_scope);
            return -1;
        }
 
        if (strcmp(expected_scope, req->scope) != 0)
        {
            fprintf(stderr, "expected scope \"%s\", got \"%s\"\n", expected_scope, req->scope);
            return -1;
        }
    }
 
    req->token = token;
    return 1;
}
 
static PostgresPollingStatusType
async_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock)
{
    if (hang_forever)
    {
        /*
         * This code tests that nothing is interfering with libpq's handling
         * of connect_timeout.
         */
        static pgsocket sock = PGINVALID_SOCKET;
 
        if (sock == PGINVALID_SOCKET)
        {
            /* First call. Create an unbound socket to wait on. */
#ifdef WIN32
            WSADATA     wsaData;
            int         err;
 
            err = WSAStartup(MAKEWORD(2, 2), &wsaData);
            if (err)
            {
                perror("WSAStartup failed");
                return PGRES_POLLING_FAILED;
            }
#endif
            sock = socket(AF_INET, SOCK_DGRAM, 0);
            if (sock == PGINVALID_SOCKET)
            {
                perror("failed to create datagram socket");
                return PGRES_POLLING_FAILED;
            }
        }
 
        /* Make libpq wait on the (unreadable) socket. */
        *altsock = sock;
        return PGRES_POLLING_READING;
    }
 
    req->token = token;
    return PGRES_POLLING_OK;
}
 
static PostgresPollingStatusType
misbehave_cb(PGconn *conn, PGoauthBearerRequest *req, pgsocket *altsock)
{
    if (strcmp(misbehave_mode, "fail-async") == 0)
    {
        /* Just fail "normally". */
        return PGRES_POLLING_FAILED;
    }
    else if (strcmp(misbehave_mode, "no-token") == 0)
    {
        /* Callbacks must assign req->token before returning OK. */
        return PGRES_POLLING_OK;
    }
    else if (strcmp(misbehave_mode, "no-socket") == 0)
    {
        /* Callbacks must assign *altsock before asking for polling. */
        return PGRES_POLLING_READING;
    }
    else
    {
        fprintf(stderr, "unrecognized --misbehave mode: %s\n", misbehave_mode);
        exit(1);
    }
}