feat(auth): add the device grant protocol state machine

Author: chajac

src/core/deviceAuth/pollState.test.ts

@@ -0,0 +1,161 @@
+import { describe, expect, it } from "bun:test";
+
+import { nextPollStep, slowDownIncrementMs } from "./pollState.js";
+import type { DeviceTokens, PollState } from "./types.js";
+
+const state: PollState = { intervalMs: 5_000, deadlineMs: 300_000 };
+
+const tokens: DeviceTokens = {
+  accessToken: "access",
+  refreshToken: "refresh",
+  expiresAt: undefined,
+  email: "person@example.com",
+  organizationId: undefined,
+};
+
+describe("nextPollStep", () => {
+  it("finishes when the response carries tokens", () => {
+    const step = nextPollStep(state, { kind: "tokens", tokens }, 0);
+
+    expect(step).toEqual({ action: "done", tokens });
+  });
+
+  it("finishes on tokens even after the deadline has passed", () => {
+    const step = nextPollStep(state, { kind: "tokens", tokens }, 300_001);
+
+    expect(step).toEqual({ action: "done", tokens });
+  });
+
+  it("polls again at the current interval while authorization is pending", () => {
+    const step = nextPollStep(state, { kind: "pending" }, 0);
+
+    expect(step).toEqual({ action: "poll", delayMs: 5_000, state });
+  });
+
+  it("raises the interval by five seconds on slow-down", () => {
+    const step = nextPollStep(state, { kind: "slow-down" }, 0);
+
+    expect(step).toEqual({
+      action: "poll",
+      delayMs: 5_000 + slowDownIncrementMs,
+      state: { intervalMs: 5_000 + slowDownIncrementMs, deadlineMs: 300_000 },
+    });
+  });
+
+  it("keeps the raised interval for later pending responses", () => {
+    const slowed = nextPollStep(state, { kind: "slow-down" }, 0);
+    if (slowed.action !== "poll") throw Error("expected to keep polling");
+
+    const step = nextPollStep(slowed.state, { kind: "pending" }, 0);
+
+    expect(step).toEqual({
+      action: "poll",
+      delayMs: 10_000,
+      state: slowed.state,
+    });
+  });
+
+  it("fails when the person rejects the request", () => {
+    const step = nextPollStep(state, { kind: "denied" }, 0);
+
+    expect(step).toEqual({
+      action: "fail",
+      reason: "access-denied",
+      detail: undefined,
+    });
+  });
+
+  it("fails when the device code expires", () => {
+    const step = nextPollStep(state, { kind: "expired" }, 0);
+
+    expect(step).toEqual({
+      action: "fail",
+      reason: "expired",
+      detail: undefined,
+    });
+  });
+
+  it("times out once the deadline passes, however long the server stalls", () => {
+    const step = nextPollStep(state, { kind: "pending" }, 300_001);
+
+    expect(step).toEqual({
+      action: "fail",
+      reason: "timeout",
+      detail: undefined,
+    });
+  });
+
+  it("keeps polling at the deadline itself", () => {
+    const step = nextPollStep(state, { kind: "pending" }, 300_000);
+
+    expect(step).toEqual({ action: "poll", delayMs: 5_000, state });
+  });
+
+  it("retries a server it could not reach, at double the interval", () => {
+    const step = nextPollStep(state, { kind: "unreachable", detail: "x" }, 0);
+
+    expect(step).toEqual({
+      action: "poll",
+      delayMs: 10_000,
+      state: { intervalMs: 10_000, deadlineMs: 300_000 },
+    });
+  });
+
+  it("doubles again while it still cannot reach the server", () => {
+    const first = nextPollStep(state, { kind: "unreachable", detail: "x" }, 0);
+    if (first.action !== "poll") throw Error("expected to keep polling");
+
+    const second = nextPollStep(
+      first.state,
+      { kind: "unreachable", detail: "x" },
+      0,
+    );
+
+    expect(second).toEqual({
+      action: "poll",
+      delayMs: 20_000,
+      state: { intervalMs: 20_000, deadlineMs: 300_000 },
+    });
+  });
+
+  it("keeps the backed-off interval once the server answers again", () => {
+    const backedOff = nextPollStep(
+      state,
+      { kind: "unreachable", detail: "x" },
+      0,
+    );
+    if (backedOff.action !== "poll") throw Error("expected to keep polling");
+
+    const step = nextPollStep(backedOff.state, { kind: "pending" }, 0);
+
+    expect(step).toEqual({
+      action: "poll",
+      delayMs: 10_000,
+      state: backedOff.state,
+    });
+  });
+
+  it("stops retrying an unreachable server once the deadline passes", () => {
+    const step = nextPollStep(
+      state,
+      { kind: "unreachable", detail: "x" },
+      300_001,
+    );
+
+    expect(step).toEqual({
+      action: "fail",
+      reason: "timeout",
+      detail: undefined,
+    });
+  });
+
+  it("fails with the detail from an unexpected error", () => {
+    const step = nextPollStep(state, { kind: "error", detail: "boom" }, 0);
+
+    expect(step).toEqual({
+      action: "fail",
+      reason: "network",
+      detail: "boom",
+    });
+  });
+});

src/core/deviceAuth/pollState.ts

@@ -0,0 +1,68 @@
+import type { PollResponse, PollState, PollStep } from "./types.js";
+
+/**
+ * The increase lasts for the rest of the flow. Treating it as a one-off skip
+ * would return to the very rate the server just objected to.
+ */
+export const slowDownIncrementMs = 5_000;
+
+/**
+ * RFC 8628 asks a client meeting a connection error to slow down before
+ * retrying, and recommends doubling. Persists for the rest of the flow: a
+ * network that dropped one request is likelier to drop the next.
+ */
+const unreachableBackoffFactor = 2;
+
+/** Pure, so the protocol is testable without a clock or a socket. */
+export function nextPollStep(
+  state: PollState,
+  response: PollResponse,
+  nowMs: number,
+): PollStep {
+  // Tokens win over the deadline. An approval that lands as the code expires is
+  // still an approval, and rejecting it would strand someone who just finished.
+  if (response.kind === "tokens") {
+    return { action: "done", tokens: response.tokens };
+  }
+
+  if (nowMs > state.deadlineMs) {
+    return { action: "fail", reason: "timeout", detail: undefined };
+  }
+
+  switch (response.kind) {
+    case "pending":
+      return { action: "poll", delayMs: state.intervalMs, state };
+
+    case "slow-down": {
+      const slowed: PollState = {
+        intervalMs: state.intervalMs + slowDownIncrementMs,
+        deadlineMs: state.deadlineMs,
+      };
+      return { action: "poll", delayMs: slowed.intervalMs, state: slowed };
+    }
+
+    case "denied":
+      return { action: "fail", reason: "access-denied", detail: undefined };
+
+    case "expired":
+      return { action: "fail", reason: "expired", detail: undefined };
+
+    // Retryable, unlike `error`. The person has often already approved in the
+    // browser by now, so abandoning the flow over one dropped request would
+    // throw away work they have done and cannot see failing.
+    case "unreachable": {
+      const backedOff: PollState = {
+        intervalMs: state.intervalMs * unreachableBackoffFactor,
+        deadlineMs: state.deadlineMs,
+      };
+      return {
+        action: "poll",
+        delayMs: backedOff.intervalMs,
+        state: backedOff,
+      };
+    }
+
+    case "error":
+      return { action: "fail", reason: "network", detail: response.detail };
+  }
+}

src/core/deviceAuth/tokenExpiry.test.ts

@@ -0,0 +1,52 @@
+import { describe, expect, it } from "bun:test";
+
+import { readAccessTokenExpiry } from "./tokenExpiry.js";
+
+function base64Url(value: string): string {
+  return Buffer.from(value, "utf8").toString("base64url");
+}
+
+function makeJwt(payload: unknown): string {
+  return [
+    base64Url(JSON.stringify({ alg: "RS256", typ: "JWT" })),
+    base64Url(JSON.stringify(payload)),
+    "signature-is-not-checked-here",
+  ].join(".");
+}
+
+describe("readAccessTokenExpiry", () => {
+  it("converts the exp claim from seconds to epoch milliseconds", () => {
+    const token = makeJwt({ exp: 1_700_000_000, sub: "user_1" });
+
+    expect(readAccessTokenExpiry(token)).toBe(1_700_000_000_000);
+  });
+
+  it("returns undefined when the payload carries no exp claim", () => {
+    expect(readAccessTokenExpiry(makeJwt({ sub: "user_1" }))).toBeUndefined();
+  });
+
+  it("returns undefined when exp is not a number", () => {
+    expect(readAccessTokenExpiry(makeJwt({ exp: "soon" }))).toBeUndefined();
+  });
+
+  it("returns undefined for a token that is not three segments", () => {
+    expect(readAccessTokenExpiry("not.ajwt")).toBeUndefined();
+  });
+
+  it("returns undefined when the payload segment is not JSON", () => {
+    const token = ["header", base64Url("not json at all"), "sig"].join(".");
+
+    expect(readAccessTokenExpiry(token)).toBeUndefined();
+  });
+
+  it("returns undefined for an empty token", () => {
+    expect(readAccessTokenExpiry("")).toBeUndefined();
+  });
+
+  it("decodes payloads containing base64url-only characters", () => {
+    // A payload whose base64 encoding needs - and _ rather than + and /.
+    const token = makeJwt({ exp: 1_700_000_001, note: "??~~??>>>" });
+
+    expect(readAccessTokenExpiry(token)).toBe(1_700_000_001_000);
+  });
+});

src/core/deviceAuth/tokenExpiry.ts

@@ -0,0 +1,30 @@
+/**
+ * Epoch milliseconds. The device token response carries no `expires_in`, so the
+ * only expiry available is the `exp` claim inside the token. Decoded without
+ * verifying the signature: the value decides when to refresh, and the API
+ * judges whether a token is genuine.
+ *
+ * Undefined for anything malformed — an unreadable expiry means "refresh it",
+ * not "crash the command".
+ */
+export function readAccessTokenExpiry(accessToken: string): number | undefined {
+  const segments = accessToken.split(".");
+  if (segments.length !== 3) return undefined;
+
+  const [, payload] = segments;
+  if (!payload) return undefined;
+
+  let claims: unknown;
+  try {
+    claims = JSON.parse(Buffer.from(payload, "base64url").toString("utf8"));
+  } catch {
+    return undefined;
+  }
+
+  if (typeof claims !== "object" || claims === null) return undefined;
+
+  const exp = (claims as { exp?: unknown }).exp;
+  if (typeof exp !== "number" || !Number.isFinite(exp)) return undefined;
+
+  return exp * 1_000;
+}

src/core/deviceAuth/types.ts

@@ -0,0 +1,41 @@
+export type DeviceTokens = {
+  accessToken: string;
+  refreshToken: string;
+  /** Epoch ms; absent when the token carried no readable expiry. */
+  expiresAt: number | undefined;
+  email: string;
+  /**
+   * WorkOS organization the session is scoped to. WorkOS picks one when the
+   * person belongs to several and the device flow never asks, so it is kept
+   * both to show which was chosen and to pin refreshes to it.
+   */
+  organizationId: string | undefined;
+};
+
+/** One token-endpoint answer, already narrowed from its OAuth error code. */
+export type PollResponse =
+  | { kind: "tokens"; tokens: DeviceTokens }
+  | { kind: "pending" }
+  | { kind: "slow-down" }
+  | { kind: "denied" }
+  | { kind: "expired" }
+  /** The server answered something unrecognised, or refused outright. */
+  | { kind: "error"; detail: string }
+  /**
+   * The server never gave a usable answer — unreachable, or failing with a 5xx
+   * or 429. Transient either way, so worth retrying.
+   */
+  | { kind: "unreachable"; detail: string };
+
+export type PollState = {
+  intervalMs: number;
+  /** Epoch ms the device code stops being redeemable. */
+  deadlineMs: number;
+};
+
+type PollFailure = "access-denied" | "expired" | "timeout" | "network";
+
+export type PollStep =
+  | { action: "poll"; delayMs: number; state: PollState }
+  | { action: "done"; tokens: DeviceTokens }
+  | { action: "fail"; reason: PollFailure; detail: string | undefined };