Accept a Partner Assertion on Cloudflare

create table app_partner_identity (
  issuer text not null,
  subject text not null,
  user_id text not null references auth_user (id) on delete cascade,
  created_at integer not null,
  primary key (issuer, subject)
);

create index app_partner_identity_user_id_idx
  on app_partner_identity (user_id);

create table app_consumed_partner_assertion (
  issuer text not null,
  assertion_id text not null,
  user_id text,
  expires_at integer not null,
  consumed_at integer not null,
  primary key (issuer, assertion_id)
);

create index app_consumed_partner_assertion_expires_at_idx
  on app_consumed_partner_assertion (expires_at);

import {
  CustomEvidencePoliciesLive,
  defineCustomEvidence,
} from "@effect-auth/core/Assurance";
import { AuthFlow } from "@effect-auth/core/AuthFlow";
import type { AuthMethod } from "@effect-auth/core/AuthFlow";
import { UnixMillis, UserId } from "@effect-auth/core/Identifiers";
import type { LoginRequestContext } from "@effect-auth/core/LoginRisk";
import { Clock, Context, Data, Effect, Schema } from "effect";

export const vendorAssertionMethod =
  "vendor-assertion" as const satisfies AuthMethod;

export const VendorAssertionEvidence = defineCustomEvidence({
  policyId: "app.vendor-assertion",
  policyVersion: 1,
  kind: "vendor-assertion",
  properties: Schema.Struct({
    issuer: Schema.String,
    authenticationContext: Schema.String,
  }),
  evaluate: () => ({
    role: "primary",
    level: "aal1",
    amr: "app:vendor_assertion",
  }),
});

export const AppCustomEvidencePoliciesLive = CustomEvidencePoliciesLive([
  VendorAssertionEvidence.policy,
]);

export class VendorAssertionInvalidError extends Data.TaggedError(
  "VendorAssertionInvalidError"
)<{ readonly message: string }> {}

export class VendorAssertionUnavailableError extends Data.TaggedError(
  "VendorAssertionUnavailableError"
)<{ readonly message: string; readonly cause?: unknown }> {}

export interface VendorAssertionVerifierService {
  readonly verify: (input: {
    readonly assertion: string;
    readonly expectedNonce: string;
  }) => Effect.Effect<
    {
      readonly userId: UserId;
      readonly evidence: {
        readonly issuer: string;
        readonly authenticationContext: string;
      };
    },
    VendorAssertionInvalidError | VendorAssertionUnavailableError
  >;
}

export class VendorAssertionVerifier extends Context.Service<
  VendorAssertionVerifier,
  VendorAssertionVerifierService
>()("app/VendorAssertionVerifier") {}

export const signInWithVendorAssertion = Effect.fn(
  "app.auth.vendor_assertion.sign_in"
)(function* (input: {
  readonly assertion: string;
  readonly expectedNonce: string;
  readonly request: LoginRequestContext;
}) {
  const verifier = yield* VendorAssertionVerifier;
  const authFlow = yield* AuthFlow;
  const verified = yield* verifier.verify({
    assertion: input.assertion,
    expectedNonce: input.expectedNonce,
  });
  const verifiedAt = UnixMillis(yield* Clock.currentTimeMillis);

  return yield* authFlow.completePrimaryFactor({
    userId: verified.userId,
    intent: "sign-in",
    method: vendorAssertionMethod,
    evidence: [
      VendorAssertionEvidence.make({
        verifiedAt,
        properties: verified.evidence,
      }),
    ],
    request: input.request,
  });
});

import type { D1Database } from "@cloudflare/workers-types";
import { UserId } from "@effect-auth/core/Identifiers";
import { Clock, Effect, Layer, Schema } from "effect";
import {
  VendorAssertionInvalidError,
  VendorAssertionUnavailableError,
  VendorAssertionVerifier,
} from "./vendor-assertion";

export interface VendorAssertionEnv {
  readonly DB: D1Database;
  readonly PARTNER_INTROSPECTION_TOKEN: string;
}

const BoundedPartnerString = Schema.String.check(
  Schema.isMinLength(1),
  Schema.isMaxLength(256)
);

const PartnerClaims = Schema.Struct({
  active: Schema.Boolean,
  issuer: BoundedPartnerString,
  subject: BoundedPartnerString,
  audience: BoundedPartnerString,
  purpose: BoundedPartnerString,
  issuedAt: Schema.Int.check(Schema.isGreaterThan(0)),
  expiresAt: Schema.Int.check(Schema.isGreaterThan(0)),
  nonce: BoundedPartnerString,
  assertionId: BoundedPartnerString,
  authenticationContext: BoundedPartnerString,
});

const decodePartnerClaims = Schema.decodeUnknownEffect(PartnerClaims, {
  errors: "all",
  onExcessProperty: "error",
});

const invalid = (message: string) =>
  new VendorAssertionInvalidError({ message });

const unavailable = (message: string, cause?: unknown) =>
  new VendorAssertionUnavailableError({ message, cause });

const readBoundedJson = (response: Response, maxBytes: number) =>
  Effect.tryPromise({
    try: async (): Promise<unknown> => {
      const declaredLength = response.headers.get("content-length");
      if (
        declaredLength !== null &&
        Number.parseInt(declaredLength, 10) > maxBytes
      ) {
        throw new TypeError("Partner response exceeded the byte limit");
      }
      if (response.body === null) {
        throw new TypeError("Partner response body was missing");
      }

      const reader = response.body.getReader();
      const chunks: Uint8Array[] = [];
      let length = 0;
      try {
        while (true) {
          const { done, value } = await reader.read();
          if (done) {
            break;
          }
          length += value.byteLength;
          if (length > maxBytes) {
            await reader.cancel();
            throw new TypeError("Partner response exceeded the byte limit");
          }
          chunks.push(value);
        }
      } finally {
        reader.releaseLock();
      }

      const bytes = new Uint8Array(length);
      let offset = 0;
      for (const chunk of chunks) {
        bytes.set(chunk, offset);
        offset += chunk.byteLength;
      }
      return JSON.parse(new TextDecoder().decode(bytes)) as unknown;
    },
    catch: (cause) => unavailable("Partner response was invalid", cause),
  });

export const VendorAssertionVerifierCloudflareLive = (
  env: VendorAssertionEnv
) =>
  Layer.succeed(VendorAssertionVerifier, {
    verify: Effect.fn("app.vendor_assertion.verify")(function* ({
      assertion,
      expectedNonce,
    }) {
      const response = yield* Effect.tryPromise({
        try: () =>
          fetch("https://login.partner.example/oauth/introspect", {
            method: "POST",
            headers: {
              authorization: `Bearer ${env.PARTNER_INTROSPECTION_TOKEN}`,
              "content-type": "application/json",
            },
            body: JSON.stringify({ assertion }),
            signal: AbortSignal.timeout(3000),
          }),
        catch: (cause) => unavailable("Partner introspection failed", cause),
      });

      if (!response.ok) {
        return yield* Effect.fail(
          unavailable(`Partner introspection returned ${response.status}`)
        );
      }

      const body = yield* readBoundedJson(response, 16_384);
      const claims = yield* decodePartnerClaims(body).pipe(
        Effect.mapError((cause) =>
          unavailable("Partner response failed schema validation", cause)
        )
      );
      const now = yield* Clock.currentTimeMillis;
      const nowSeconds = Math.floor(now / 1000);

      if (
        !claims.active ||
        claims.issuer !== "https://login.partner.example" ||
        claims.audience !== "marketplace-auth" ||
        claims.purpose !== "primary-sign-in" ||
        claims.nonce !== expectedNonce ||
        claims.issuedAt > nowSeconds + 30 ||
        claims.expiresAt <= nowSeconds ||
        claims.expiresAt <= claims.issuedAt ||
        claims.expiresAt - claims.issuedAt > 120
      ) {
        return yield* Effect.fail(invalid("Partner assertion was rejected"));
      }

      const consumed = yield* Effect.tryPromise({
        try: () =>
          env.DB.prepare(
            `insert or ignore into app_consumed_partner_assertion
              (issuer, assertion_id, user_id, expires_at, consumed_at)
             values (
               ?,
               ?,
               (select identity.user_id
                from app_partner_identity as identity
                join auth_user as user on user.id = identity.user_id
                where identity.issuer = ?
                  and identity.subject = ?
                  and user.disabled_at is null),
               ?,
               ?
             )
             returning user_id`
          )
            .bind(
              claims.issuer,
              claims.assertionId,
              claims.issuer,
              claims.subject,
              claims.expiresAt * 1000,
              now
            )
            .first<{ readonly user_id: string | null }>(),
        catch: (cause) => unavailable("D1 assertion consume failed", cause),
      });

      if (consumed === null || consumed.user_id === null) {
        return yield* Effect.fail(invalid("Partner assertion was rejected"));
      }

      return {
        userId: UserId(consumed.user_id),
        evidence: {
          issuer: claims.issuer,
          authenticationContext: claims.authenticationContext,
        },
      };
    }),
  });

import {
  AuthFlowLive,
  defaultMfaRequirementMethods,
  MfaRequirementPolicyLive,
} from "@effect-auth/core/AuthFlow";
import { SessionsLive } from "@effect-auth/core/Sessions";
import { Layer } from "effect";
import {
  AppCustomEvidencePoliciesLive,
  vendorAssertionMethod,
} from "./vendor-assertion";

const CustomMethodMfaPolicyLive = MfaRequirementPolicyLive({
  methods: [...defaultMfaRequirementMethods, vendorAssertionMethod],
});

const SessionsWithCustomEvidenceLive = SessionsLive().pipe(
  Layer.provide(AppCustomEvidencePoliciesLive)
);

export const AuthFlowWithCustomEvidenceLive = AuthFlowLive.pipe(
  Layer.provide(CustomMethodMfaPolicyLive),
  Layer.provide(SessionsWithCustomEvidenceLive)
);

import type { AuthFlowPrimaryFactorError } from "@effect-auth/core/AuthFlow";
import {
  AuthHttp,
  AuthInternalError,
  AuthInvalidCredentialsError,
  AuthOriginCheckMiddleware,
  AuthPolicyDeniedError,
  AuthRateLimitedError,
  AuthSchemaErrorMiddleware,
  PrimaryAuthSuccess,
} from "@effect-auth/core/HttpApi";
import type { LoginRequestContext } from "@effect-auth/core/LoginRisk";
import { Effect, Schema } from "effect";
import {
  HttpApi,
  HttpApiEndpoint,
  HttpApiGroup,
} from "effect/unstable/httpapi";
import {
  signInWithVendorAssertion,
  VendorAssertionInvalidError,
  VendorAssertionUnavailableError,
} from "./vendor-assertion";

export const VendorAssertionSignInBody = Schema.Struct({
  assertion: Schema.String.check(
    Schema.isMinLength(1),
    Schema.isMaxLength(16_384)
  ),
});

export const vendorAssertionSignInEndpoint = HttpApiEndpoint.post(
  "signIn",
  "/sign-in",
  {
    payload: VendorAssertionSignInBody,
    success: PrimaryAuthSuccess,
    error: [
      AuthInvalidCredentialsError,
      AuthPolicyDeniedError,
      AuthRateLimitedError,
      AuthInternalError,
    ],
  }
);

export class VendorAssertionHttpApiGroup extends HttpApiGroup.make(
  "vendorAssertion"
)
  .add(vendorAssertionSignInEndpoint)
  .prefix("/auth/vendor-assertion")
  .middleware(AuthSchemaErrorMiddleware)
  .middleware(AuthOriginCheckMiddleware) {}

export class VendorAssertionHttpApi extends HttpApi.make(
  "VendorAssertionHttpApi"
).add(VendorAssertionHttpApiGroup) {}

const mapVendorAssertionHttpError = (
  error:
    | VendorAssertionInvalidError
    | VendorAssertionUnavailableError
    | AuthFlowPrimaryFactorError
): AuthInvalidCredentialsError | AuthInternalError =>
  error._tag === "VendorAssertionInvalidError"
    ? new AuthInvalidCredentialsError({
        code: "invalid_credentials",
        message: "Invalid credentials",
      })
    : new AuthInternalError({
        code: "internal_error",
        message: "Failed to sign in with vendor assertion",
      });

export const vendorAssertionSignInOperation = Effect.fn(
  "app.http.vendor_assertion.sign_in"
)(function* (input: {
  readonly assertion: string;
  readonly expectedNonce: string;
  readonly request: LoginRequestContext;
}) {
  const authHttp = yield* AuthHttp;
  const result = yield* signInWithVendorAssertion(input).pipe(
    Effect.mapError(mapVendorAssertionHttpError)
  );

  return yield* authHttp.commitPrimaryFactorResult(result);
});

import {
  createAuthClient,
  defineAuthHttpApiExtension,
} from "@effect-auth/core/Client";
import { Schema } from "effect";
import {
  VendorAssertionHttpApi,
  VendorAssertionSignInBody,
} from "./vendor-assertion-api";

type VendorAssertionSignInRequest = Schema.Schema.Type<
  typeof VendorAssertionSignInBody
>;

const vendorAssertionExtension = defineAuthHttpApiExtension(
  VendorAssertionHttpApi,
  ({ run }) => ({
    signIn: (input: VendorAssertionSignInRequest) =>
      run((client) => client.vendorAssertion.signIn({ payload: input })),
  })
);

export const auth = createAuthClient({
  protocol: {
    extensions: { vendorAssertion: vendorAssertionExtension },
  },
});

const result = await auth.extensions.vendorAssertion.signIn({ assertion });

if (result.type === "requires_mfa") {
  await auth.mfa.totp.verify({ flowId: result.flowId, code });
}

import * as Alchemy from "alchemy";
import * as Cloudflare from "alchemy/Cloudflare";
import * as Config from "effect/Config";
import * as Effect from "effect/Effect";
import type { RATE_LIMITER as RateLimitDurableObject } from "./src/worker";

export const Database = Cloudflare.D1.Database("AuthDatabase", {
  migrationsDir: "./migrations",
});

export const RateLimiter = Cloudflare.DurableObject<RateLimitDurableObject>(
  "RATE_LIMITER",
  {
    className: "RATE_LIMITER",
  }
);

export const AuthWorker = Cloudflare.Worker("AuthWorker", {
  main: "./src/worker.ts",
  compatibility: { flags: ["nodejs_compat"] },
  env: {
    DB: Database,
    RATE_LIMITER: RateLimiter,
    PARTNER_INTROSPECTION_TOKEN: Config.redacted("PARTNER_INTROSPECTION_TOKEN"),
  },
});

export default Alchemy.Stack(
  "MarketplaceAuth",
  {
    providers: Cloudflare.providers(),
    state: Cloudflare.state(),
  },
  Effect.gen(function* () {
    const worker = yield* AuthWorker;
    return { url: worker.url.as<string>() };
  })
);