atproto OAuth client in CloudFlare Worker

atprotoOauthClient.ts

import { JoseKey, type Importable } from "@atproto/jwk-jose";
import type { JWK } from "jose";
// should be usable with nodejs compat of 'node:dns/promises'
import { AppViewHandleResolver } from "@atproto-labs/handle-resolver";
import type { Jwk } from "@atproto/jwk";

// Preserve the original Request constructor
const OriginalRequest = globalThis.Request;

// Create a custom Request constructor
class CustomRequest extends OriginalRequest {
  constructor(input: RequestInfo, init?: RequestInit) {
    // Sanitize the `init` object to handle unsupported fields
    const sanitizedInit = CustomRequest.sanitizeRequestInit(init);

    // Call the original constructor with the sanitized init
    super(input, sanitizedInit);
  }

  // Static method to sanitize the RequestInit object
  private static sanitizeRequestInit(
    init: RequestInit | undefined
  ): RequestInit {
    if (!init) {
      return {};
    }

    // Destructure and handle unsupported fields
    const { cache, redirect, ...rest } = init;

    // Log warnings for unsupported or adjusted fields
    if (cache) {
      console.debug(
        `The 'cache' field is not supported in Cloudflare Workers and will be ignored.`
      );
    }

    if (redirect === "error") {
      console.debug(
        `Sanitizing Request constructor: the 'redirect: "error"' option is not supported in Cloudflare Workers. Using 'redirect: "manual"' instead.`
      );
      (rest as RequestInit).redirect = "manual"; // Replace with 'manual'
    } else if (redirect) {
      (rest as RequestInit).redirect = redirect; // Retain other valid redirect options
    }

    return rest; // Return sanitized RequestInit object
  }
}
globalThis.Request = CustomRequest as unknown as typeof Request;

function addAlgorithmToJWK(jwk: JWK): JWK {
  if (!jwk.alg) {
    switch (jwk.kty) {
      case "RSA":
        // Default to RS256 if no alg is specified
        jwk.alg = "RS256";
        break;
      case "EC":
        // Infer algorithm based on curve
        if (jwk.crv === "P-256") {
          jwk.alg = "ES256";
        } else if (jwk.crv === "P-384") {
          jwk.alg = "ES384";
        } else if (jwk.crv === "P-521") {
          jwk.alg = "ES512";
        } else if (jwk.crv === "secp256k1") {
          jwk.alg = "ES256K";
        } else {
          throw new Error("Unsupported EC curve");
        }
        break;
      case "oct":
        // Default to HS256 for symmetric keys
        jwk.alg = "HS256";
        break;
      default:
        throw new Error(`Unsupported key type: ${jwk.kty}`);
    }
  }

  return jwk;
}

const originalFetch = globalThis.fetch;
globalThis.fetch = async function (...args: Parameters<typeof fetch>) {
  if (args[1]) {
    if (args[1].redirect === "error") {
      console.debug(
        `sanitizing fetch: The 'redirect: "error"' option is not supported in Cloudflare Workers. Using 'redirect: "manual"' instead.`
      );
      args[1].redirect = "manual";
    }
  }

  return originalFetch(...args);
};
const resolver = new AppViewHandleResolver("https://bsky.social", {});

import {
  OAuthClient,
  type InternalStateData,
  type Session,
  type DigestAlgorithm,
  type TokenSet,
  type Key,
} from "@atproto/oauth-client";

type SerializedState = Omit<InternalStateData, "dpopKey"> & {
  dpopKey: Record<string, unknown> | undefined;
};

const base_url = import.meta.env.BASE_BSKY_URL ?? "https://domain.invalid";
// change to workers kv - probably needs dependency injection? because this library
// doesn't have access to the cloudflare/request context
// const stateStore = new Map<string, InternalStateData>();
// const sessionStore = new Map<string, NodeSavedSession>();
const createClient = async (
  stateStore: KVNamespace,
  sessionStore: KVNamespace
) => {
  const client = new OAuthClient({
    fetch: globalThis.fetch,
    responseMode: "query",
    handleResolver: resolver,
    runtimeImplementation: {
      // A runtime specific implementation of the crypto operations needed by the
      // OAuth client. See "@atproto/oauth-client-browser" for a browser specific
      // implementation. The following example is suitable for use in NodeJS.

      async createKey(algs: string[]): Promise<Key> {
        // algs is an ordered array of preferred algorithms (e.g. ['RS256', 'ES256'])

        // Note, in browser environments, it is better to use non extractable keys
        // to prevent the private key from being stolen. This can be done using
        // the WebcryptoKey class from the "@atproto/jwk-webcrypto" package. The
        // inconvenient of these keys (which is also what makes them stronger) is
        // that the only way to persist them across browser reloads is to save
        // them in the indexed DB.
        const joseKey = await JoseKey.generate(algs);
        const jwk = joseKey.privateJwk!;
        addAlgorithmToJWK(jwk);

        return joseKey;
      },
      getRandomValues(byteLength: number): Uint8Array {
        return crypto.getRandomValues(new Uint8Array(byteLength));
      },
      async digest(
        data: Uint8Array,
        { name }: DigestAlgorithm
      ): Promise<Uint8Array> {
        switch (name) {
          case "sha256":
          case "sha384":
          case "sha512": {
            const buf = await crypto.subtle.digest(
              `SHA-${name.slice(3)}`,
              data
            );
            return new Uint8Array(buf);
          }
          default:
            throw new Error(`Unsupported digest algorithm: ${name}`);
        }
      },
    },
    // This object will be used to build the payload of the /client-metadata.json
    // endpoint metadata, exposing the client metadata to the OAuth server.
    clientMetadata: {
      // Must be a URL that will be exposing this metadata
      client_id: `${base_url}/client-metadata.json`,
      client_name: "My App",
      client_uri: `${base_url}`,
      logo_uri: `${base_url}/logo.png`,
      tos_uri: `${base_url}/tos`,
      policy_uri: `${base_url}/policy`,
      redirect_uris: [`${base_url}/callback`],
      scope: "transition:generic atproto",
      token_endpoint_auth_signing_alg: "ES256",
      grant_types: ["authorization_code", "refresh_token"],
      response_types: ["code"],
      application_type: "web",
      token_endpoint_auth_method: "private_key_jwt",
      dpop_bound_access_tokens: true,
      jwks_uri: `${base_url}/jwks.json`,
    },

    // Used to authenticate the client to the token endpoint. Will be used to
    // build the jwks object to be exposed on the "jwks_uri" endpoint.
    keyset: await Promise.all([
      JoseKey.fromImportable(import.meta.env.PRIVATE_KEY_1 ?? "", "0"),
      JoseKey.fromImportable(import.meta.env.PRIVATE_KEY_2 ?? "", "1"),
      JoseKey.fromImportable(import.meta.env.PRIVATE_KEY_3 ?? "", "2"),
    ]),

    // Interface to store authorization state data (during authorization flows)
    stateStore: {
      async set(key: string, internalState: InternalStateData): Promise<void> {
        const serializedState: SerializedState = {
          ...internalState,
          dpopKey: internalState.dpopKey.privateJwk,
        };
        await stateStore.put(key, JSON.stringify(serializedState));
      },
      async get(key: string): Promise<InternalStateData | undefined> {
        const serializedResult = await stateStore.get<SerializedState>(
          key,
          "json"
        );
        if (
          serializedResult === null ||
          serializedResult.dpopKey == undefined
        ) {
          return undefined;
        }
        const result = {
          ...serializedResult,
          dpopKey: await JoseKey.fromJWK(serializedResult.dpopKey),
        };
        return result;
      },
      async del(key: string): Promise<void> {
        await stateStore.delete(key);
      },
    },

    // Interface to store authenticated session data
    sessionStore: {
      async set(sub: string, session: Session): Promise<void> {
        const keyToSerialize = session.dpopKey.privateJwk! as Jwk;
        const serializeableSession = {
          ...session,
          dpopKey: keyToSerialize,
        };
        await sessionStore.put(sub, JSON.stringify(serializeableSession));
      },
      async get(sub: string): Promise<Session | undefined> {
        const serializedSession = await sessionStore.get<SerializedSession>(
          sub,
          "json"
        );
        if (!serializedSession) return undefined;
        console.log(
          `restoring session for key ${JSON.stringify(serializedSession)}`
        );
        const session = {
          ...serializedSession,
          dpopKey: await JoseKey.fromJWK(serializedSession.dpopKey),
        } as unknown as Session;
        return session;
      },
      async del(sub: string): Promise<void> {
        await sessionStore.delete(sub);
      },
    },
  });
  return client;
};
export { createClient };

export type SerializedSession = {
  dpopKey: Jwk;
  tokenSet: TokenSet;
};