819 lines · 30.0 KB · TypeScript
RawDownload
// ---------------------------------------------------------------------------// WorkOS AuthKit — Effect-native sealed session management// --------------------------------------------------------------------------- import { env } from "cloudflare:workers";import { Context, Data, Effect, Layer, Option, Predicate, Schema } from "effect";import { GeneratePortalLinkIntent, WorkOS } from "@workos-inc/node/worker";import { defaults as ironDefaults, unseal as unsealIron } from "iron-webcrypto";import { decodeJwt, jwtVerify } from "jose";import { JWKSInvalid, JWKSNoMatchingKey, JWKSTimeout } from "jose/errors";import { parseCookie } from "./cookies";import { createCachedRemoteJWKSet, type CachedRemoteJWKSet } from "./jwks-cache";import {  ServiceAdapterError,  tryPromiseService,  withServiceLogging,  workosErrorFromFailure,} from "./errors"; const COOKIE_NAME = "wos-session";const INVALID_COOKIE_PASSWORD_MESSAGE = "WORKOS_COOKIE_PASSWORD must be at least 32 characters";const WORKOS_SEAL_VERSION_DELIMITER = "~"; type RawWorkOS = WorkOS & {  readonly get: (    path: string,    options?: { readonly query?: Record<string, unknown> },  ) => Promise<{    readonly data: unknown;  }>;  readonly post: (    path: string,    entity: unknown,    options?: { readonly idempotencyKey?: string },  ) => Promise<{ readonly data: unknown }>;}; type WorkOSListMetadata = {  readonly before?: string | null;  readonly after?: string | null;}; type WorkOSAutoPaginatable<Resource> = {  readonly object: "list";  readonly data: Resource[];  readonly listMetadata: WorkOSListMetadata;  readonly autoPagination: () => Promise<Resource[]>;}; export type WorkOSCollectedList<Resource> = {  readonly object: "list";  readonly data: Resource[];  readonly listMetadata: {    readonly before: string | null;    readonly after: string | null;  };}; const RawWorkOSListMetadata = Schema.Struct({  before: Schema.optional(Schema.NullOr(Schema.String)),  after: Schema.optional(Schema.NullOr(Schema.String)),}); const RawWorkOSListResponse = Schema.Struct({  data: Schema.Array(Schema.Unknown),  listMetadata: Schema.optional(RawWorkOSListMetadata),  list_metadata: Schema.optional(RawWorkOSListMetadata),}); const decodeRawWorkOSListResponse = Schema.decodeUnknownOption(RawWorkOSListResponse); const SealedSessionUser = Schema.Struct({  id: Schema.String,  email: Schema.String,  firstName: Schema.NullOr(Schema.String),  lastName: Schema.NullOr(Schema.String),  profilePictureUrl: Schema.NullOr(Schema.String),}); const SealedSessionPayload = Schema.Struct({  accessToken: Schema.String,  refreshToken: Schema.String,  user: SealedSessionUser,  authenticationMethod: Schema.optional(Schema.Unknown),  impersonator: Schema.optional(Schema.Unknown),}); const decodeSealedSessionPayload = Schema.decodeUnknownOption(SealedSessionPayload); const JwtClaims = Schema.Struct({  sid: Schema.String,  org_id: Schema.optional(Schema.String),}); const decodeJwtClaims = Schema.decodeUnknownOption(JwtClaims); const LegacySealedSessionPayload = Schema.Struct({  persistent: Schema.Unknown,}); const decodeLegacySealedSessionPayload = Schema.decodeUnknownOption(LegacySealedSessionPayload); type SealedSessionPayload = typeof SealedSessionPayload.Type; type LocalSessionVerification =  | {      readonly _tag: "Valid";      readonly session: SealedSessionPayload;      readonly organizationId: string | undefined;      readonly sessionId: string;    }  | { readonly _tag: "Refresh" }  | { readonly _tag: "InvalidCookie" }; class LocalSessionCookieError extends Data.TaggedError("LocalSessionCookieError")<{  readonly cause: unknown;}> {} const isLocalSessionValid = Predicate.isTagged("Valid") as (  value: LocalSessionVerification,) => value is Extract<LocalSessionVerification, { readonly _tag: "Valid" }>; const isLocalSessionInvalidCookie = Predicate.isTagged("InvalidCookie") as (  value: LocalSessionVerification,) => value is Extract<LocalSessionVerification, { readonly _tag: "InvalidCookie" }>; const ErrorWithCode = Schema.Struct({ code: Schema.String });const isErrorWithCode = Schema.is(ErrorWithCode); const isJwtVerificationFailure = (cause: unknown): boolean =>  isErrorWithCode(cause) &&  (cause.code.startsWith("ERR_JWT_") ||    cause.code.startsWith("ERR_JWS_") ||    cause.code === JWKSNoMatchingKey.code); const isJwksSystemFailure = (cause: unknown): boolean =>  isErrorWithCode(cause) && (cause.code === JWKSTimeout.code || cause.code === JWKSInvalid.code); const parseWorkOSSeal = (  seal: string,): { readonly sealWithoutVersion: string; readonly tokenVersion: number | null } => {  const [sealWithoutVersion = "", tokenVersionAsString] = seal.split(WORKOS_SEAL_VERSION_DELIMITER);  return {    sealWithoutVersion,    tokenVersion:      tokenVersionAsString === undefined ? null : Number.parseInt(tokenVersionAsString, 10),  };}; const unsealWorkOSSession = async (  sessionData: string,  cookiePassword: string,): Promise<unknown> => {  const { sealWithoutVersion, tokenVersion } = parseWorkOSSeal(sessionData);  const data =    (await unsealIron(sealWithoutVersion, { 1: cookiePassword }, { ...ironDefaults, ttl: 0 })) ??    {};  // Mirrors the SDK's unsealData version handling: current (v2) seals hold the  // payload directly, OTHER versioned seals nest it under `persistent`, and an  // unversioned seal is the payload itself.  if (tokenVersion === 2 || tokenVersion === null) return data;  return Option.match(decodeLegacySealedSessionPayload(data), {    onNone: () => data,    onSome: (legacy) => legacy.persistent,  });}; const getWorkOSSessionJwks = (() => {  const resolvers = new Map<string, CachedRemoteJWKSet>();  return (url: URL): CachedRemoteJWKSet => {    const key = url.toString();    const existing = resolvers.get(key);    if (existing) return existing;    const created = createCachedRemoteJWKSet(url);    resolvers.set(key, created);    return created;  };})(); const verifyJwtOnce = (accessToken: string, jwks: CachedRemoteJWKSet) =>  Effect.tryPromise({    try: () => jwtVerify(accessToken, jwks),    catch: (cause) => new ServiceAdapterError({ cause }),  }); const verifyJwtWithRefreshRetry = (  accessToken: string,  jwks: CachedRemoteJWKSet,): Effect.Effect<boolean, ServiceAdapterError> =>  verifyJwtOnce(accessToken, jwks).pipe(    Effect.as(true),    Effect.catchTag("ServiceAdapterError", (error) => {      if (isErrorWithCode(error.cause) && error.cause.code === JWKSNoMatchingKey.code) {        return Effect.sync(() => jwks.forceRefresh()).pipe(          Effect.flatMap(() => verifyJwtOnce(accessToken, jwks)),          Effect.as(true),          Effect.catchTag("ServiceAdapterError", (retryError) =>            isJwtVerificationFailure(retryError.cause)              ? Effect.succeed(false)              : Effect.fail(retryError),          ),        );      }      if (isJwtVerificationFailure(error.cause)) return Effect.succeed(false);      if (isJwksSystemFailure(error.cause)) return Effect.fail(error);      return Effect.fail(error);    }),  ); const verifySealedSessionLocally = (  sessionData: string,  cookiePassword: string,  jwks: CachedRemoteJWKSet,): Effect.Effect<LocalSessionVerification, ServiceAdapterError> =>  Effect.gen(function* () {    // Phase timings, not just child spans. `local_verify` is a leaf in    // production traces, so a ~3.3s verify has nothing under it to blame —    // and it stayed 3.3s after the JWKS fetch was eliminated entirely    // (jwks.fetch_count == 0), so the cost is one of the phases below. Under    // workerd `Date.now()` only advances at I/O boundaries, which is exactly    // what makes a raw span duration misleading here: recording each phase    // explicitly says which await the wall-clock actually crossed.    const verifyStartedAt = Date.now();     const unsealStartedAt = Date.now();    const unsealed = yield* Effect.tryPromise({      try: () => unsealWorkOSSession(sessionData, cookiePassword),      catch: (cause) => new LocalSessionCookieError({ cause }),    }).pipe(      Effect.catchTag("LocalSessionCookieError", () => Effect.succeed(null as unknown | null)),      Effect.withSpan("workos.session.unseal"),    );    const unsealMs = Date.now() - unsealStartedAt;    yield* Effect.annotateCurrentSpan({ "verify.unseal_ms": unsealMs });    if (!unsealed) return { _tag: "InvalidCookie" };     const decodeStartedAt = Date.now();    const session = Option.match(decodeSealedSessionPayload(unsealed), {      onNone: (): SealedSessionPayload | null => null,      onSome: (payload) => payload,    });    yield* Effect.annotateCurrentSpan({ "verify.decode_ms": Date.now() - decodeStartedAt });    if (!session) return { _tag: "InvalidCookie" };     // Snapshot the JWKS cache around the verify so the `local_verify` span    // says whether THIS verify was a warm-cache signature check or paid for a    // live upstream JWKS fetch. The Aug 2026 latency regression was the cache    // silently missing on most verifies, and no span attribute distinguished    // the two paths.    const jwksBefore = jwks.inspect();    // Entry-state annotation goes on BEFORE the verify so a failing verify    // (the case worth debugging) still records whether the cache was warm.    yield* Effect.annotateCurrentSpan({      "jwks.cache_populated_at_start": jwksBefore.hasJwks,      ...(jwksBefore.fetchedAt === null        ? {}        : { "jwks.cache_age_ms": Date.now() - jwksBefore.fetchedAt }),    });    const jwtStartedAt = Date.now();    const verified = yield* verifyJwtWithRefreshRetry(session.accessToken, jwks).pipe(      Effect.withSpan("workos.session.jwt_verify"),      Effect.onExit(() => {        const jwksAfter = jwks.inspect();        const finishedAt = Date.now();        return Effect.annotateCurrentSpan({          "verify.jwt_ms": finishedAt - jwtStartedAt,          "verify.total_ms": finishedAt - verifyStartedAt,          // Blocking, not total: under stale-while-revalidate a background          // refresh moves `fetchCount` without costing this verify anything.          // Attribute latency to what the caller actually waited on.          "jwks.fetched_during_verify":            jwksAfter.blockingFetchCount > jwksBefore.blockingFetchCount,          "jwks.served_from_store": jwksAfter.storeHitCount > jwksBefore.storeHitCount,          "jwks.fetch_count": jwksAfter.fetchCount,          "jwks.blocking_fetch_count": jwksAfter.blockingFetchCount,          "jwks.fetch_failure_count": jwksAfter.fetchFailureCount,          ...(jwksAfter.lastFetchDurationMs === null            ? {}            : { "jwks.last_fetch_ms": jwksAfter.lastFetchDurationMs }),          // Splits the ~3.4s that sits inside jwt_verify with zero upstream          // fetches: the cross-isolate store read (I/O) vs WebCrypto key          // import vs the signature check itself.          ...(jwksAfter.lastStoreReadMs === null            ? {}            : { "jwks.store_read_ms": jwksAfter.lastStoreReadMs }),          ...(jwksAfter.lastResolveMs === null            ? {}            : { "jwks.key_resolve_ms": jwksAfter.lastResolveMs }),        });      }),    );    if (!verified) return { _tag: "Refresh" };     const claims = Option.getOrNull(decodeJwtClaims(decodeJwt(session.accessToken)));    if (!claims) return { _tag: "Refresh" };     return {      _tag: "Valid",      session,      organizationId: claims.org_id,      sessionId: claims.sid,    };  }); const completedListMetadata = {  before: null,  after: null,} as const; const nextCursorFromRawList = (response: typeof RawWorkOSListResponse.Type): string | null =>  response.listMetadata?.after ?? response.list_metadata?.after ?? null; export const collectWorkOSList = async <Resource>(  response: WorkOSAutoPaginatable<Resource>,): Promise<WorkOSCollectedList<Resource>> => {  const data = response.listMetadata.after ? await response.autoPagination() : response.data;  return {    object: "list",    data,    listMetadata: completedListMetadata,  };}; export const collectRawWorkOSList = async (  loadPage: (after?: string) => Promise<unknown>,): Promise<WorkOSCollectedList<unknown>> => {  const first = Option.getOrNull(decodeRawWorkOSListResponse(await loadPage()));  if (!first) {    return {      object: "list",      data: [],      listMetadata: completedListMetadata,    };  }   const data = [...first.data];  let after = nextCursorFromRawList(first);   while (after) {    const next = Option.getOrNull(decodeRawWorkOSListResponse(await loadPage(after)));    if (!next) break;    data.push(...next.data);    after = nextCursorFromRawList(next);  }   return {    object: "list",    data,    listMetadata: completedListMetadata,  };}; class WorkOSConfigurationError extends Data.TaggedError("WorkOSConfigurationError")<{  readonly message: string;}> {} /** * Optional base-URL override for the WorkOS API (`WORKOS_API_URL`) — points * the REAL SDK at a WorkOS emulator in tests/dev. Unset in production, where * the SDK uses api.workos.com. Sealed-session crypto, JWKS verification, and * every endpoint follow this host, so the whole auth stack runs against the * emulator with zero code substitution. */export const workosApiUrlOptions = (  url: string | undefined,): { apiHostname?: string; port?: number; https?: boolean } => {  if (!url) return {};  const parsed = new URL(url);  return {    apiHostname: parsed.hostname,    ...(parsed.port ? { port: Number(parsed.port) } : {}),    https: parsed.protocol === "https:",  };}; // ---------------------------------------------------------------------------// Service// --------------------------------------------------------------------------- const make = Effect.gen(function* () {  const apiKey = env.WORKOS_API_KEY;  const clientId = env.WORKOS_CLIENT_ID;  const cookiePassword = env.WORKOS_COOKIE_PASSWORD;   if (!cookiePassword || cookiePassword.length < 32) {    return yield* new WorkOSConfigurationError({      message: INVALID_COOKIE_PASSWORD_MESSAGE,    });  }   const workos = new WorkOS({    apiKey,    clientId,    ...workosApiUrlOptions(env.WORKOS_API_URL),  });  const sessionJwks = getWorkOSSessionJwks(new URL(workos.userManagement.getJwksUrl(clientId)));   // The public `WorkOSError` carries the upstream HTTP status when the SDK  // exception had one (all its typed exceptions do), so consumers can tell a  // definitive WorkOS denial (401/403/404 — fail closed) from a transient  // failure (429/5xx/network — retryable).  // `op` names the SDK call (mirroring its `namespace.method` path) so every  // span reads `workos.<operation>` instead of one undifferentiated "workos"  // bucket, and failures log which call actually failed.  const use = <A>(op: string, fn: (wos: WorkOS) => Promise<A>) =>    withServiceLogging(      `workos.${op}`,      workosErrorFromFailure,      tryPromiseService(() => fn(workos)),    );   const authenticateSealedSession = (sessionData: string) =>    Effect.gen(function* () {      if (!sessionData) return null;       const session = workos.userManagement.loadSealedSession({        sessionData,        cookiePassword,      });       const local = yield* withServiceLogging(        "workos.session.local_verify",        workosErrorFromFailure,        verifySealedSessionLocally(sessionData, cookiePassword, sessionJwks),      );       if (isLocalSessionValid(local)) {        return {          userId: local.session.user.id,          email: local.session.user.email,          firstName: local.session.user.firstName,          lastName: local.session.user.lastName,          avatarUrl: local.session.user.profilePictureUrl,          organizationId: local.organizationId,          sessionId: local.sessionId,          refreshedSession: undefined as string | undefined,        };      }       if (isLocalSessionInvalidCookie(local)) return null;       // Try refreshing      const refreshed = yield* use("session.refresh", () => session.refresh()).pipe(        Effect.orElseSucceed(() => ({ authenticated: false as const })),      );       if (!refreshed.authenticated || !("sealedSession" in refreshed) || !refreshed.sealedSession)        return null;       return {        userId: refreshed.user.id,        email: refreshed.user.email,        firstName: refreshed.user.firstName,        lastName: refreshed.user.lastName,        avatarUrl: refreshed.user.profilePictureUrl,        organizationId: refreshed.organizationId,        sessionId: refreshed.sessionId,        refreshedSession: refreshed.sealedSession,      };    });   return {    getAuthorizationUrl: (redirectUri: string, state?: string) =>      workos.userManagement.getAuthorizationUrl({        provider: "authkit",        redirectUri,        clientId,        ...(state ? { state } : {}),      }),     authenticateWithCode: (code: string) =>      use("userManagement.authenticateWithCode", (wos) =>        wos.userManagement.authenticateWithCode({          code,          clientId,          session: { sealSession: true, cookiePassword },        }),      ),     /** Create a new organization in WorkOS. */    createOrganization: (name: string) =>      use("organizations.createOrganization", (wos) =>        wos.organizations.createOrganization({ name }),      ),     /** Add a user to an organization. */    createMembership: (organizationId: string, userId: string, roleSlug?: string) =>      use("userManagement.createOrganizationMembership", (wos) =>        wos.userManagement.createOrganizationMembership({          organizationId,          userId,          ...(roleSlug ? { roleSlug } : {}),        }),      ),     /** List organization memberships for a user. */    listUserMemberships: (userId: string) =>      use("userManagement.listOrganizationMemberships", async (wos) =>        collectWorkOSList(          await wos.userManagement.listOrganizationMemberships({

Lines 1–500 of 819.