yaak/plugins/auth-oauth2/src/grants/clientCredentials.ts

import { createPrivateKey, randomUUID } from "node:crypto";
import type { Context } from "@yaakapp/api";
import jwt, { type Algorithm } from "jsonwebtoken";
import { fetchAccessToken } from "../fetchAccessToken";
import type { TokenStoreArgs } from "../store";
import { getToken, storeToken } from "../store";
import { isTokenExpired } from "../util";

export const jwtAlgorithms = [
  "HS256",
  "HS384",
  "HS512",
  "RS256",
  "RS384",
  "RS512",
  "PS256",
  "PS384",
  "PS512",
  "ES256",
  "ES384",
  "ES512",
  "none",
] as const;

export const defaultJwtAlgorithm = jwtAlgorithms[0];

/**
 * Build a signed JWT for the client_assertion parameter (RFC 7523).
 *
 * The `secret` value is auto-detected as one of:
 *   - **JWK** – a JSON string containing a private-key object (has a `kty` field).
 *   - **PEM** – a string whose trimmed form starts with `-----`.
 *   - **HMAC secret** – anything else, used as-is for HS* algorithms.
 */
function buildClientAssertionJwt(params: {
  clientId: string;
  accessTokenUrl: string;
  secret: string;
  algorithm: Algorithm;
}): string {
  const { clientId, accessTokenUrl, secret, algorithm } = params;

  const isHmac = algorithm.startsWith("HS") || algorithm === "none";

  // Resolve the signing key depending on format
  let signingKey: jwt.Secret;
  let kid: string | undefined;

  const trimmed = secret.trim();

  if (isHmac) {
    // HMAC algorithms use the raw secret (string or Buffer)
    signingKey = secret;
  } else if (trimmed.startsWith("{")) {
    // Looks like JSON - treat as JWK. There is surely a better way to detect JWK vs a raw secret, but this should work in most cases.
    // oxlint-disable-next-line no-explicit-any
    let jwk: any;
    try {
      jwk = JSON.parse(trimmed);
    } catch {
      throw new Error("Client Assertion secret looks like JSON but is not valid");
    }

    kid = jwk?.kid;
    signingKey = createPrivateKey({ key: jwk, format: "jwk" });
  } else if (trimmed.startsWith("-----")) {
    // PEM-encoded key
    signingKey = createPrivateKey({ key: trimmed, format: "pem" });
  } else {
    throw new Error(
      "Client Assertion secret must be a JWK JSON object, a PEM-encoded key " +
        "(starting with -----), or a raw secret for HMAC algorithms.",
    );
  }

  const now = Math.floor(Date.now() / 1000);
  const payload = {
    iss: clientId,
    sub: clientId,
    aud: accessTokenUrl,
    iat: now,
    exp: now + 300, // 5 minutes
    jti: randomUUID(),
  };

  // Build the JWT header; include "kid" when available
  const header: jwt.JwtHeader = { alg: algorithm, typ: "JWT" };
  if (kid) {
    header.kid = kid;
  }

  return jwt.sign(JSON.stringify(payload), signingKey, {
    algorithm: algorithm as jwt.Algorithm,
    header,
  });
}

export async function getClientCredentials(
  ctx: Context,
  contextId: string,
  {
    accessTokenUrl,
    clientId,
    clientSecret,
    scope,
    audience,
    credentialsInBody,
    clientAssertionSecret,
    clientAssertionSecretBase64,
    clientCredentialsMethod,
    clientAssertionAlgorithm,
  }: {
    accessTokenUrl: string;
    clientId: string;
    clientSecret: string;
    scope: string | null;
    audience: string | null;
    credentialsInBody: boolean;
    clientAssertionSecret: string;
    clientAssertionSecretBase64: boolean;
    clientCredentialsMethod: string;
    clientAssertionAlgorithm: string;
  },
) {
  const tokenArgs: TokenStoreArgs = {
    contextId,
    clientId,
    accessTokenUrl,
    authorizationUrl: null,
  };
  const token = await getToken(ctx, tokenArgs);
  if (token && !isTokenExpired(token)) {
    return token;
  }

  const common: Omit<
    Parameters<typeof fetchAccessToken>[1],
    "clientAssertion" | "clientSecret" | "credentialsInBody"
  > = {
    grantType: "client_credentials",
    accessTokenUrl,
    audience,
    clientId,
    scope,
    params: [],
  };

  const fetchParams: Parameters<typeof fetchAccessToken>[1] =
    clientCredentialsMethod === "client_assertion"
      ? {
          ...common,
          clientAssertion: buildClientAssertionJwt({
            clientId,
            algorithm: clientAssertionAlgorithm as Algorithm,
            accessTokenUrl,
            secret: clientAssertionSecretBase64
              ? Buffer.from(clientAssertionSecret, "base64").toString("utf-8")
              : clientAssertionSecret,
          }),
        }
      : {
          ...common,
          clientSecret,
          credentialsInBody,
        };

  const response = await fetchAccessToken(ctx, fetchParams);

  return storeToken(ctx, tokenArgs, response);
}