feat: secure Service Account session config using signed JWT cookies

mcp-examples/pocket-cep/src/__tests__/integration/api/sa-config-route.test.ts

@@ -39,13 +39,21 @@ vi.mock("@/lib/access-token", async (importOriginal) => {
 });
 
 import { GET, POST, DELETE } from "@/app/api/auth/sa-config/route";
-import { COOKIE_SA_CUSTOMER_ID, COOKIE_SA_IMPERSONATED_USER } from "@/lib/sa-session";
+import {
+  COOKIE_SA_SESSION,
+  COOKIE_SA_CUSTOMER_ID,
+  COOKIE_SA_IMPERSONATED_USER,
+} from "@/lib/sa-session";
+import { verifyJwt } from "@/lib/jwt";
 import { DwdScopeVerificationError } from "@/lib/access-token";
 
 describe("/api/auth/sa-config", () => {
   beforeEach(() => {
     vi.clearAllMocks();
-    mockGetEnv.mockReturnValue({ AUTH_MODE: "service_account" });
+    mockGetEnv.mockReturnValue({
+      AUTH_MODE: "service_account",
+      BETTER_AUTH_SECRET: "mock-secret",
+    });
   });
 
   describe("GET", () => {
@@ -159,8 +167,14 @@ describe("/api/auth/sa-config", () => {
         customerId: "C00woaabb",
         impersonatedUser: "admin@example.com",
       });
-      expect(res.cookies.get(COOKIE_SA_CUSTOMER_ID)?.value).toBe("C00woaabb");
-      expect(res.cookies.get(COOKIE_SA_IMPERSONATED_USER)?.value).toBe("admin@example.com");
+      const sessionCookie = res.cookies.get(COOKIE_SA_SESSION)?.value;
+      expect(sessionCookie).toBeDefined();
+      const payload = verifyJwt(sessionCookie!, "mock-secret");
+      expect(payload?.customerId).toBe("C00woaabb");
+      expect(payload?.impersonatedUser).toBe("admin@example.com");
+      expect(payload?.exp).toBeTypeOf("number");
+      expect(res.cookies.get(COOKIE_SA_CUSTOMER_ID)?.value).toBe("");
+      expect(res.cookies.get(COOKIE_SA_IMPERSONATED_USER)?.value).toBe("");
     });
 
     it("returns 400 and blocks cookie saving when Option 2 (Direct Mode) token minting fails", async () => {
@@ -197,7 +211,13 @@ describe("/api/auth/sa-config", () => {
         customerId: "C00woaabb",
         impersonatedUser: "",
       });
-      expect(res.cookies.get(COOKIE_SA_CUSTOMER_ID)?.value).toBe("C00woaabb");
+      const sessionCookie = res.cookies.get(COOKIE_SA_SESSION)?.value;
+      expect(sessionCookie).toBeDefined();
+      const payload = verifyJwt(sessionCookie!, "mock-secret");
+      expect(payload?.customerId).toBe("C00woaabb");
+      expect(payload?.impersonatedUser).toBe("");
+      expect(payload?.exp).toBeTypeOf("number");
+      expect(res.cookies.get(COOKIE_SA_CUSTOMER_ID)?.value).toBe("");
       expect(res.cookies.get(COOKIE_SA_IMPERSONATED_USER)?.value).toBe("");
     });
   });
@@ -208,6 +228,7 @@ describe("/api/auth/sa-config", () => {
       expect(res.status).toBe(200);
       expect(await res.json()).toEqual({ success: true });
       expect(mockClearTokenCache).toHaveBeenCalled();
+      expect(res.cookies.get(COOKIE_SA_SESSION)?.value).toBe("");
       expect(res.cookies.get(COOKIE_SA_CUSTOMER_ID)?.value).toBe("");
       expect(res.cookies.get(COOKIE_SA_IMPERSONATED_USER)?.value).toBe("");
     });

mcp-examples/pocket-cep/src/__tests__/integration/proxy.test.ts

@@ -8,9 +8,10 @@
 
 import { describe, it, expect, vi, beforeEach } from "vitest";
 
-const { mockGetEnv, mockGetSessionCookie } = vi.hoisted(() => ({
+const { mockGetEnv, mockGetSessionCookie, mockProbeMcpServer } = vi.hoisted(() => ({
   mockGetEnv: vi.fn(),
   mockGetSessionCookie: vi.fn(),
+  mockProbeMcpServer: vi.fn().mockResolvedValue({ ok: true, message: "ok" }),
 }));
 
 vi.mock("@/lib/env", async (importOriginal) => {
@@ -25,9 +26,15 @@ vi.mock("better-auth/cookies", () => ({
   getSessionCookie: mockGetSessionCookie,
 }));
 
+vi.mock("@/lib/doctor-checks", () => ({
+  probeMcpServer: mockProbeMcpServer,
+}));
+
 import { NextRequest } from "next/server";
 import { proxy } from "@/proxy";
 import { EnvValidationError } from "@/lib/env";
+import { signJwt } from "@/lib/jwt";
+import { COOKIE_SA_SESSION } from "@/lib/sa-session";
 
 function makeRequest(
   url = "http://localhost:3000/",
@@ -135,6 +142,40 @@ describe("proxy — normal auth routing", () => {
     expect(res.headers.get("location")).toContain("/sa-setup");
   });
 
+  it("service_account + session + configured SA via valid JWT + /dashboard → allows (no redirect)", async () => {
+    mockGetEnv.mockReturnValue({
+      AUTH_MODE: "service_account",
+      BETTER_AUTH_SECRET: "mock-secret",
+      MCP_SERVER_URL: "http://localhost:4000/mcp",
+    });
+    mockGetSessionCookie.mockReturnValue("signed-cookie");
+
+    const validSaSession = signJwt({ customerId: "C01234567" }, "mock-secret");
+    const res = await proxy(
+      makeRequest("http://localhost:3000/dashboard", {
+        headers: { cookie: `${COOKIE_SA_SESSION}=${validSaSession}` },
+      }),
+    );
+    expect(res.headers.get("location")).toBeNull();
+  });
+
+  it("service_account + session + invalid SA session JWT + /dashboard → redirects to /sa-setup", async () => {
+    mockGetEnv.mockReturnValue({
+      AUTH_MODE: "service_account",
+      BETTER_AUTH_SECRET: "mock-secret",
+    });
+    mockGetSessionCookie.mockReturnValue("signed-cookie");
+
+    const invalidSaSession = "invalid.jwt.signature";
+    const res = await proxy(
+      makeRequest("http://localhost:3000/dashboard", {
+        headers: { cookie: `${COOKIE_SA_SESSION}=${invalidSaSession}` },
+      }),
+    );
+    expect(res.status).toBe(307);
+    expect(res.headers.get("location")).toContain("/sa-setup");
+  });
+
   it("user_oauth + session + /sa-setup → redirects to /", async () => {
     mockGetEnv.mockReturnValue({ AUTH_MODE: "user_oauth" });
     mockGetSessionCookie.mockReturnValue("signed-cookie");
@@ -192,14 +233,16 @@ describe("proxy — MCP reachability gate", () => {
     mockGetEnv.mockReturnValue({
       AUTH_MODE: "service_account",
       MCP_SERVER_URL: "http://localhost:4000/mcp",
+      BETTER_AUTH_SECRET: "mock-secret",
     });
     mockGetSessionCookie.mockReturnValue("signed-cookie");
 
+    const validSaSession = signJwt({ customerId: "C01234567" }, "mock-secret");
     const { proxy: proxyImpl } = await import("@/proxy");
     const { NextRequest: Req } = await import("next/server");
     const res = await proxyImpl(
       new Req(new URL("http://localhost:3000/dashboard"), {
-        headers: { cookie: "cep_sa_customer_id=C01234567" },
+        headers: { cookie: `${COOKIE_SA_SESSION}=${validSaSession}` },
       }),
     );
 
@@ -223,14 +266,16 @@ describe("proxy — MCP reachability gate", () => {
     mockGetEnv.mockReturnValue({
       AUTH_MODE: "service_account",
       MCP_SERVER_URL: "http://localhost:4000/mcp",
+      BETTER_AUTH_SECRET: "mock-secret",
     });
     mockGetSessionCookie.mockReturnValue("signed-cookie");
 
+    const validSaSession = signJwt({ customerId: "C01234567" }, "mock-secret");
     const { proxy: proxyImpl } = await import("@/proxy");
     const { NextRequest: Req } = await import("next/server");
     const res = await proxyImpl(
       new Req(new URL("http://localhost:3000/dashboard"), {
-        headers: { cookie: "cep_sa_customer_id=C01234567" },
+        headers: { cookie: `${COOKIE_SA_SESSION}=${validSaSession}` },
       }),
     );
 
@@ -272,19 +317,21 @@ describe("proxy — MCP reachability gate", () => {
     mockGetEnv.mockReturnValue({
       AUTH_MODE: "service_account",
       MCP_SERVER_URL: "http://localhost:4000/mcp",
+      BETTER_AUTH_SECRET: "mock-secret",
     });
     mockGetSessionCookie.mockReturnValue("signed-cookie");
 
+    const validSaSession = signJwt({ customerId: "C01234567" }, "mock-secret");
     const { proxy: proxyImpl } = await import("@/proxy");
     const { NextRequest: Req } = await import("next/server");
     await proxyImpl(
       new Req(new URL("http://localhost:3000/dashboard"), {
-        headers: { cookie: "cep_sa_customer_id=C01234567" },
+        headers: { cookie: `${COOKIE_SA_SESSION}=${validSaSession}` },
       }),
     );
     await proxyImpl(
       new Req(new URL("http://localhost:3000/dashboard/extra"), {
-        headers: { cookie: "cep_sa_customer_id=C01234567" },
+        headers: { cookie: `${COOKIE_SA_SESSION}=${validSaSession}` },
       }),
     );
 

mcp-examples/pocket-cep/src/__tests__/unit/auth-errors.test.ts

@@ -166,7 +166,9 @@ describe("toAuthError", () => {
       const err = new Error("Invalid Customer Id");
       const result = toAuthError(err, "admin-sdk");
       expect(result?.code).toBe("invalid_customer_id");
-      expect(result?.remedy).toContain("Ensure the Customer ID parameter passed to the tool is correct");
+      expect(result?.remedy).toContain(
+        "Ensure the Customer ID parameter passed to the tool is correct",
+      );
       expect(result?.remedy).not.toContain("/sa-setup");
     });
   });

mcp-examples/pocket-cep/src/__tests__/unit/sa-session.test.ts

@@ -1,28 +1,28 @@
-import { describe, it, expect, vi } from "vitest";
-import {
-  getServiceAccountConfig,
-  COOKIE_SA_CUSTOMER_ID,
-  COOKIE_SA_IMPERSONATED_USER,
-} from "@/lib/sa-session";
-
-let mockCookieHas = (name: string): boolean =>
-  name === COOKIE_SA_CUSTOMER_ID || name === COOKIE_SA_IMPERSONATED_USER;
-let mockCookieGet = (name: string): { value?: string } | undefined => {
-  if (name === COOKIE_SA_CUSTOMER_ID) return { value: "C01234567" };
-  if (name === COOKIE_SA_IMPERSONATED_USER) return { value: "admin@example.com" };
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { getServiceAccountConfig, COOKIE_SA_SESSION } from "@/lib/sa-session";
+import { signJwt } from "@/lib/jwt";
+
+const SECRET = "mock-secret";
+
+let mockSessionCookieValue: string | undefined = undefined;
+
+const mockCookieGet = (name: string): { value?: string } | undefined => {
+  if (name === COOKIE_SA_SESSION) {
+    return mockSessionCookieValue ? { value: mockSessionCookieValue } : undefined;
+  }
   return undefined;
 };
 
 vi.mock("next/headers", () => ({
   cookies: vi.fn(async () => ({
-    has: (name: string) => mockCookieHas(name),
+    has: (name: string) => !!mockCookieGet(name),
     get: (name: string) => mockCookieGet(name),
   })),
 }));
 
 const mockGetEnv = vi.fn(() => ({
   AUTH_MODE: "service_account" as const,
-  BETTER_AUTH_SECRET: "mock-secret",
+  BETTER_AUTH_SECRET: SECRET,
   BETTER_AUTH_URL: "http://localhost:3000",
   MCP_SERVER_URL: "http://localhost:4000/mcp",
   LLM_MODEL: "",
@@ -35,14 +35,17 @@ vi.mock("@/lib/env", () => ({
 }));
 
 describe("getServiceAccountConfig", () => {
-  it("resolves customerId and impersonatedUser from cookies", async () => {
-    mockCookieHas = (name: string): boolean =>
-      name === COOKIE_SA_CUSTOMER_ID || name === COOKIE_SA_IMPERSONATED_USER;
-    mockCookieGet = (name) => {
-      if (name === COOKIE_SA_CUSTOMER_ID) return { value: "C01234567" };
-      if (name === COOKIE_SA_IMPERSONATED_USER) return { value: "admin@example.com" };
-      return undefined;
-    };
+  beforeEach(() => {
+    mockSessionCookieValue = undefined;
+    delete process.env.CEP_CUSTOMER_ID;
+    delete process.env.CEP_IMPERSONATE_SUBJECT;
+  });
+
+  it("resolves customerId and impersonatedUser from valid JWT session cookie", async () => {
+    mockSessionCookieValue = signJwt(
+      { customerId: "C01234567", impersonatedUser: "admin@example.com" },
+      SECRET,
+    );
 
     const config = await getServiceAccountConfig();
     expect(config).toEqual({
@@ -51,46 +54,69 @@ describe("getServiceAccountConfig", () => {
     });
   });
 
-  it("does not fall back to CEP_IMPERSONATE_SUBJECT when COOKIE_SA_CUSTOMER_ID exists (Option 2 Direct Mode)", async () => {
+  it("does not fall back to CEP_IMPERSONATE_SUBJECT when valid session cookie exists but has no impersonation (Direct Mode)", async () => {
     process.env.CEP_IMPERSONATE_SUBJECT = "zombie-admin@example.com";
-    mockCookieHas = (name: string): boolean => name === COOKIE_SA_CUSTOMER_ID;
-    mockCookieGet = (name) => {
-      if (name === COOKIE_SA_CUSTOMER_ID) return { value: "C01234567" };
-      if (name === COOKIE_SA_IMPERSONATED_USER) return undefined;
-      return undefined;
-    };
+    mockSessionCookieValue = signJwt({ customerId: "C01234567" }, SECRET);
 
     const config = await getServiceAccountConfig();
     expect(config).toEqual({
       customerId: "C01234567",
       impersonatedUser: undefined,
     });
-    delete process.env.CEP_IMPERSONATE_SUBJECT;
   });
 
-  it("falls back to CEP_IMPERSONATE_SUBJECT when no customer session cookie has been saved yet but CEP_CUSTOMER_ID is set", async () => {
+  it("falls back to CEP_IMPERSONATE_SUBJECT when session cookie is missing but CEP_CUSTOMER_ID is set", async () => {
     process.env.CEP_CUSTOMER_ID = "C09876543";
     process.env.CEP_IMPERSONATE_SUBJECT = "env-admin@example.com";
-    mockCookieHas = (_name: string): boolean => false;
-    mockCookieGet = () => undefined;
+    mockSessionCookieValue = undefined;
 
     const config = await getServiceAccountConfig();
     expect(config).toEqual({
       customerId: "C09876543",
       impersonatedUser: "env-admin@example.com",
     });
-    delete process.env.CEP_CUSTOMER_ID;
-    delete process.env.CEP_IMPERSONATE_SUBJECT;
   });
 
   it("returns null when customerId is empty even if CEP_IMPERSONATE_SUBJECT is set", async () => {
-    delete process.env.CEP_CUSTOMER_ID;
     process.env.CEP_IMPERSONATE_SUBJECT = "env-admin@example.com";
-    mockCookieHas = (_name: string): boolean => false;
-    mockCookieGet = () => undefined;
+    mockSessionCookieValue = undefined;
 
     const config = await getServiceAccountConfig();
     expect(config).toBeNull();
-    delete process.env.CEP_IMPERSONATE_SUBJECT;
+  });
+
+  it("falls back to env variables when the session cookie signature is invalid (tampered)", async () => {
+    process.env.CEP_CUSTOMER_ID = "C09876543";
+    process.env.CEP_IMPERSONATE_SUBJECT = "env-admin@example.com";
+
+    // Sign with a different secret
+    mockSessionCookieValue = signJwt(
+      { customerId: "C01234567", impersonatedUser: "attacker@example.com" },
+      "attacker-secret-key",
+    );
+
+    const config = await getServiceAccountConfig();
+    // Signature validation should fail and fall back to the env configurations
+    expect(config).toEqual({
+      customerId: "C09876543",
+      impersonatedUser: "env-admin@example.com",
+    });
+  });
+
+  it("falls back to env variables when the session cookie is expired", async () => {
+    process.env.CEP_CUSTOMER_ID = "C09876543";
+    process.env.CEP_IMPERSONATE_SUBJECT = "env-admin@example.com";
+
+    const expiredExp = Math.floor(Date.now() / 1000) - 10;
+    mockSessionCookieValue = signJwt(
+      { customerId: "C01234567", impersonatedUser: "admin@example.com", exp: expiredExp },
+      SECRET,
+    );
+
+    const config = await getServiceAccountConfig();
+    expect(config).toEqual({
+      customerId: "C09876543",
+      impersonatedUser: "env-admin@example.com",
+    });
   });
 });

mcp-examples/pocket-cep/src/app/api/auth/sa-config/route.ts

@@ -8,10 +8,12 @@
 
 import { NextRequest, NextResponse } from "next/server";
 import {
+  COOKIE_SA_SESSION,
   COOKIE_SA_CUSTOMER_ID,
   COOKIE_SA_IMPERSONATED_USER,
   getServiceAccountConfig,
 } from "@/lib/sa-session";
+import { signJwt } from "@/lib/jwt";
 import { getEnv } from "@/lib/env";
 import {
   clearServiceAccountTokenCache,
@@ -109,13 +111,12 @@ export async function POST(request: NextRequest) {
     secure: isProduction,
   };
 
-  response.cookies.set(COOKIE_SA_CUSTOMER_ID, customerId, cookieOptions);
+  const exp = Math.floor(Date.now() / 1000) + 86400 * 30; // 30 days
+  const token = signJwt({ customerId, impersonatedUser, exp }, env.BETTER_AUTH_SECRET);
+  response.cookies.set(COOKIE_SA_SESSION, token, cookieOptions);
 
-  if (impersonatedUser) {
-    response.cookies.set(COOKIE_SA_IMPERSONATED_USER, impersonatedUser, cookieOptions);
-  } else {
-    response.cookies.delete(COOKIE_SA_IMPERSONATED_USER);
-  }
+  response.cookies.delete(COOKIE_SA_CUSTOMER_ID);
+  response.cookies.delete(COOKIE_SA_IMPERSONATED_USER);
 
   return response;
 }
@@ -126,6 +127,7 @@ export async function POST(request: NextRequest) {
 export async function DELETE() {
   clearServiceAccountTokenCache();
   const response = NextResponse.json({ success: true });
+  response.cookies.delete(COOKIE_SA_SESSION);
   response.cookies.delete(COOKIE_SA_CUSTOMER_ID);
   response.cookies.delete(COOKIE_SA_IMPERSONATED_USER);
   return response;