/*-------------------------------------------------------------------------
 *
 * validator.c
 *   Test module for serverside OAuth token validation callbacks
 *
 * Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
 * Portions Copyright (c) 1994, Regents of the University of California
 *
 * src/test/modules/oauth_validator/validator.c
 *
 *-------------------------------------------------------------------------
 */

#include "postgres.h"

#include "fmgr.h"
#include "libpq/oauth.h"
#include "miscadmin.h"
#include "utils/guc.h"
#include "utils/memutils.h"

PG_MODULE_MAGIC;

static void validator_startup(ValidatorModuleState *state);
static void validator_shutdown(ValidatorModuleState *state);
static bool validate_token(const ValidatorModuleState *state,
         const char *token,
         const char *role,
         ValidatorModuleResult *res);

/* Callback implementations (exercise all three) */
static const OAuthValidatorCallbacks validator_callbacks = {
 PG_OAUTH_VALIDATOR_MAGIC,

 .startup_cb = validator_startup,
 .shutdown_cb = validator_shutdown,
 .validate_cb = validate_token
};

/* GUCs */
static char *authn_id = NULL;
static bool authorize_tokens = true;

/*---
 * Extension entry point. Sets up GUCs for use by tests:
 *
 * - oauth_validator.authn_id Sets the user identifier to return during token
 *        validation. Defaults to the username in the
 *        startup packet.
 *
 * - oauth_validator.authorize_tokens
 *        Sets whether to successfully validate incoming
 *        tokens. Defaults to true.
 */
void
_PG_init(void)
{
 DefineCustomStringVariable("oauth_validator.authn_id",
          "Authenticated identity to use for future connections",
          NULL,
          &authn_id,
          NULL,
          PGC_SIGHUP,
          0,
          NULL, NULL, NULL);
 DefineCustomBoolVariable("oauth_validator.authorize_tokens",
        "Should tokens be marked valid?",
        NULL,
        &authorize_tokens,
        true,
        PGC_SIGHUP,
        0,
        NULL, NULL, NULL);

 MarkGUCPrefixReserved("oauth_validator");
}

/*
 * Validator module entry point.
 */
const OAuthValidatorCallbacks *
_PG_oauth_validator_module_init(void)
{
 return &validator_callbacks;
}

#define PRIVATE_COOKIE ((void *) 13579)

/*
 * Startup callback, to set up private data for the validator.
 */
static void
validator_startup(ValidatorModuleState *state)
{
 /*
  * Make sure the server is correctly setting sversion. (Real modules
  * should not do this; it would defeat upgrade compatibility.)
 */
 if (state->sversion != PG_VERSION_NUM)
  elog(ERROR, "oauth_validator: sversion set to %d", state->sversion);

 state->private_data = PRIVATE_COOKIE;
}

/*
 * Shutdown callback, to tear down the validator.
 */
static void
validator_shutdown(ValidatorModuleState *state)
{
 /* Check to make sure our private state still exists. */
 if (state->private_data != PRIVATE_COOKIE)
  elog(PANIC, "oauth_validator: private state cookie changed to %p in shutdown",
    state->private_data);
}

/*
 * Validator implementation. Logs the incoming data and authorizes the token by
 * default; the behavior can be modified via the module's GUC settings.
 */
static bool
validate_token(const ValidatorModuleState *state,
      const char *token, const char *role,
      ValidatorModuleResult *res)
{
 /* Check to make sure our private state still exists. */
 if (state->private_data != PRIVATE_COOKIE)
  elog(ERROR, "oauth_validator: private state cookie changed to %p in validate",
    state->private_data);

 elog(LOG, "oauth_validator: token=\"%s\", role=\"%s\"", token, role);
 elog(LOG, "oauth_validator: issuer=\"%s\", scope=\"%s\"",
   MyProcPort->hba->oauth_issuer,
   MyProcPort->hba->oauth_scope);

 res->authorized = authorize_tokens;
 if (authn_id)
  res->authn_id = pstrdup(authn_id);
 else
  res->authn_id = pstrdup(role);

 return true;
}