aniplay-api/src/libs/gcloud/getGoogleAuthToken.ts

import type { GetTokenOptions, TokenData as GoogleTokenData } from "gtoken";
import { SignJWT, importPKCS8 } from "jose";
import { DateTime } from "luxon";

import { lazy } from "../lazy";
import type { AdminSdkCredentials } from "./getAdminSdkCredentials";

export async function getGoogleAuthToken(adminSdkJson: AdminSdkCredentials) {
  const { privateKey, clientEmail } = adminSdkJson;

  const gToken = new GoogleToken(
    {
      key: privateKey,
      email: clientEmail,
      scope: [
        "https://www.googleapis.com/auth/firebase.messaging",
        "https://www.googleapis.com/auth/cloud-tasks",
      ],
    },
    adminSdkJson,
  );
  return gToken.getToken().then((token) => token.access_token);
}

const GOOGLE_TOKEN_URL = "https://www.googleapis.com/oauth2/v4/token";

class GoogleToken {
  #inFlightRequest?: Promise<GoogleTokenData>;
  #tokenData?: TokenData;

  #scope = lazy(() => {
    const scope = this.options.scope;

    if (Array.isArray(scope)) {
      return scope;
    } else if (typeof scope === "string") {
      return [scope];
    } else {
      return [];
    }
  });

  constructor(
    private options: TokenOptions,
    credentialsJson: AdminSdkCredentials,
  ) {
    this.options.key = this.options.key ?? credentialsJson.privateKey;
    this.options.email = this.options.email ?? credentialsJson.clientEmail;
  }

  getToken({
    forceRefresh = false,
    ...options
  }: GetTokenOptions = {}): Promise<GoogleTokenData> {
    return this.#getTokenAsync({ ...options, forceRefresh });
  }

  isTokenExpiring(): boolean {
    const now = DateTime.now();
    const eagerRefreshThresholdMillis =
      this.options.eagerRefreshThresholdMillis ?? 0;
    if (this.#tokenData) {
      return (
        now.plus({ milliseconds: eagerRefreshThresholdMillis }) >=
        this.#tokenData.expiresAt
      );
    }

    return true;
  }

  #getTokenAsync(options: GetTokenOptions): Promise<GoogleTokenData> {
    const { forceRefresh } = options;

    if (this.#inFlightRequest && !forceRefresh) {
      return this.#inFlightRequest;
    }

    try {
      this.#inFlightRequest = this.#getTokenAsyncInternal(options);
      return this.#inFlightRequest!;
    } finally {
      this.#inFlightRequest = undefined;
    }
  }

  #getTokenAsyncInternal(options: GetTokenOptions): Promise<GoogleTokenData> {
    if (!this.isTokenExpiring() && options.forceRefresh) {
      return Promise.resolve(this.#tokenData!.token);
    }

    return this.#requestToken();
  }

  async #requestToken(): Promise<GoogleTokenData> {
    if (!this.options.email) {
      throw new Error("No email provided");
    }
    if (!this.options.key) {
      throw new Error("No private key provided");
    }

    const issuedTokenAt = DateTime.now().toSeconds();
    const additionalClaims = this.options.additionalClaims ?? {};
    const jwtPayload = {
      iss: this.options.email,
      scope: this.#scope.get().join(" "),
      aud: GOOGLE_TOKEN_URL,
      exp: issuedTokenAt + 3600,
      iat: issuedTokenAt,
      additionalClaims,
      sub: this.options.sub,
    };

    const key = await importPKCS8(this.options.key, "RS256");
    const signedJwt = await new SignJWT(jwtPayload)
      .setProtectedHeader({ alg: "RS256" })
      .sign(key);

    try {
      const res = await fetch(GOOGLE_TOKEN_URL, {
        method: "POST",
        headers: {
          "Content-Type": "application/x-www-form-urlencoded",
        },
        body: new URLSearchParams({
          grant_type: "urn:ietf:params:oauth:grant-type:jwt-bearer",
          assertion: signedJwt,
        }),
      });
      this.#tokenData = await res.json<GoogleTokenData>().then((data) => ({
        token: data,
        expiresAt: DateTime.fromSeconds(issuedTokenAt).plus({
          seconds: data.expires_in,
        }),
      }));

      return this.#tokenData!.token;
    } catch (e) {
      console.error(e);
      throw e;
    }
  }
}

export interface TokenOptions {
  keyFile?: string;
  key?: string;
  email?: string;
  iss?: string;
  sub?: string;
  scope?: string | string[];
  additionalClaims?: {};
  /** Eagerly refresh the token if it is within this many milliseconds from expiring. Defaults to 0. */
  eagerRefreshThresholdMillis?: number;
}

interface TokenData {
  token: GoogleTokenData;
  expiresAt: DateTime;
}