diff --git a/mcp-examples/pocket-cep/src/__tests__/integration/api/sa-config-route.test.ts b/mcp-examples/pocket-cep/src/__tests__/integration/api/sa-config-route.test.ts index b047802..f7ba238 100644 --- a/mcp-examples/pocket-cep/src/__tests__/integration/api/sa-config-route.test.ts +++ b/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(""); }); diff --git a/mcp-examples/pocket-cep/src/__tests__/integration/proxy.test.ts b/mcp-examples/pocket-cep/src/__tests__/integration/proxy.test.ts index 974ee9b..7f4f8c9 100644 --- a/mcp-examples/pocket-cep/src/__tests__/integration/proxy.test.ts +++ b/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}` }, }), ); diff --git a/mcp-examples/pocket-cep/src/__tests__/unit/auth-errors.test.ts b/mcp-examples/pocket-cep/src/__tests__/unit/auth-errors.test.ts index d4e1486..8e88e3b 100644 --- a/mcp-examples/pocket-cep/src/__tests__/unit/auth-errors.test.ts +++ b/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"); }); }); diff --git a/mcp-examples/pocket-cep/src/__tests__/unit/sa-session.test.ts b/mcp-examples/pocket-cep/src/__tests__/unit/sa-session.test.ts index cdf8ade..e70eb2f 100644 --- a/mcp-examples/pocket-cep/src/__tests__/unit/sa-session.test.ts +++ b/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", + }); }); }); diff --git a/mcp-examples/pocket-cep/src/app/api/auth/sa-config/route.ts b/mcp-examples/pocket-cep/src/app/api/auth/sa-config/route.ts index 8412081..55e4092 100644 --- a/mcp-examples/pocket-cep/src/app/api/auth/sa-config/route.ts +++ b/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; diff --git a/mcp-examples/pocket-cep/src/lib/activity-data.ts b/mcp-examples/pocket-cep/src/lib/activity-data.ts index 564be1d..83ee38a 100644 --- a/mcp-examples/pocket-cep/src/lib/activity-data.ts +++ b/mcp-examples/pocket-cep/src/lib/activity-data.ts @@ -100,7 +100,11 @@ export async function getActivitySafe(days: number = DEFAULT_ACTIVITY_DAYS): Pro * Pulls and groups Chrome audit events for the given caller, scoped to * `days` of history. Pagination stops at {@link ACTIVITY_MAX_EVENTS}. */ -async function fetchActivity(tokenToUse: string, days: number, impersonatedUser?: string): Promise { +async function fetchActivity( + tokenToUse: string, + days: number, + impersonatedUser?: string, +): Promise { const requestHeaders = await buildGoogleApiHeaders(tokenToUse); const baseUrl = new URL( diff --git a/mcp-examples/pocket-cep/src/lib/jwt.ts b/mcp-examples/pocket-cep/src/lib/jwt.ts new file mode 100644 index 0000000..4e6ca30 --- /dev/null +++ b/mcp-examples/pocket-cep/src/lib/jwt.ts @@ -0,0 +1,49 @@ +/** + * @file Stateless JWT utility module for Pocket CEP using Node's crypto library. + */ + +import { createHmac, timingSafeEqual } from "node:crypto"; + +/** + * Generates a signed HS256 JWT string for the given payload. + */ +export function signJwt(payload: Record, secret: string): string { + const header = Buffer.from(JSON.stringify({ alg: "HS256", typ: "JWT" })).toString("base64url"); + const body = Buffer.from(JSON.stringify(payload)).toString("base64url"); + const signature = createHmac("sha256", secret).update(`${header}.${body}`).digest("base64url"); + return `${header}.${body}.${signature}`; +} + +/** + * Verifies the signature of an HS256 JWT string and parses the payload. + * Returns null if the signature is invalid or the token is malformed. + */ +export function verifyJwt(token: string, secret: string): Record | null { + const parts = token.split("."); + if (parts.length !== 3) return null; + + const [header, body, signature] = parts; + const expectedSignature = createHmac("sha256", secret) + .update(`${header}.${body}`) + .digest("base64url"); + + const sigBuffer = Buffer.from(signature, "base64url"); + const expectedBuffer = Buffer.from(expectedSignature, "base64url"); + + if (sigBuffer.length !== expectedBuffer.length || !timingSafeEqual(sigBuffer, expectedBuffer)) { + return null; + } + + try { + const payload = JSON.parse(Buffer.from(body, "base64url").toString("utf8")) as Record< + string, + unknown + >; + if (typeof payload.exp === "number" && Date.now() / 1000 > payload.exp) { + return null; + } + return payload; + } catch { + return null; + } +} diff --git a/mcp-examples/pocket-cep/src/lib/sa-session.ts b/mcp-examples/pocket-cep/src/lib/sa-session.ts index 60bf305..5931519 100644 --- a/mcp-examples/pocket-cep/src/lib/sa-session.ts +++ b/mcp-examples/pocket-cep/src/lib/sa-session.ts @@ -3,7 +3,8 @@ * * In service_account mode, Pocket CEP stores the administrator's selected * Google Workspace Customer ID (`customerId`) and optional Domain-Wide Delegation - * Impersonated User (`impersonatedUser`) inside HTTP-only cookies. + * Impersonated User (`impersonatedUser`) inside a cryptographically signed + * HTTP-only JWT cookie. * * These helpers read configuration on the server side so the token minter * (`src/lib/access-token.ts`) and MCP tool caller (`src/lib/mcp-tools.ts`) can @@ -12,9 +13,11 @@ import { cookies } from "next/headers"; import { getEnv } from "@/lib/env"; +import { verifyJwt } from "./jwt"; -export const COOKIE_SA_CUSTOMER_ID = "cep_sa_customer_id"; -export const COOKIE_SA_IMPERSONATED_USER = "cep_sa_impersonated_user"; +export const COOKIE_SA_SESSION = "cep_sa_session"; +export const COOKIE_SA_CUSTOMER_ID = "cep_sa_customer_id"; // Deprecated but kept for transition cleanup +export const COOKIE_SA_IMPERSONATED_USER = "cep_sa_impersonated_user"; // Deprecated but kept for transition cleanup export interface ServiceAccountConfig { customerId: string; @@ -22,7 +25,7 @@ export interface ServiceAccountConfig { } /** - * Retrieves the configured Service Account tenant credentials from cookies. + * Retrieves the configured Service Account tenant credentials from the signed session JWT. * Must be called inside a Next.js server request context (Route Handler or Server Action). * * @returns The ServiceAccountConfig object, or null if customerId has not been set. @@ -30,17 +33,32 @@ export interface ServiceAccountConfig { export async function getServiceAccountConfig(): Promise { const cookieStore = await cookies(); const env = getEnv(); - const hasCustomerCookie = cookieStore.has(COOKIE_SA_CUSTOMER_ID); - const customerId = - cookieStore.get(COOKIE_SA_CUSTOMER_ID)?.value?.trim() || env.CEP_CUSTOMER_ID?.trim(); - const rawImpersonated = cookieStore.get(COOKIE_SA_IMPERSONATED_USER)?.value?.trim(); - const impersonatedUser = hasCustomerCookie - ? rawImpersonated && rawImpersonated.length > 0 - ? rawImpersonated - : undefined - : (rawImpersonated && rawImpersonated.length > 0 ? rawImpersonated : undefined) || - env.CEP_IMPERSONATE_SUBJECT?.trim(); + let customerId = env.CEP_CUSTOMER_ID?.trim() || ""; + let impersonatedUser = env.CEP_IMPERSONATE_SUBJECT?.trim() || undefined; + let hasSessionCookie = false; + + const sessionCookie = cookieStore.get(COOKIE_SA_SESSION)?.value; + if (sessionCookie) { + const payload = verifyJwt(sessionCookie, env.BETTER_AUTH_SECRET); + if (payload) { + hasSessionCookie = true; + if (typeof payload.customerId === "string") { + customerId = payload.customerId.trim(); + } + if (typeof payload.impersonatedUser === "string" && payload.impersonatedUser.trim()) { + impersonatedUser = payload.impersonatedUser.trim(); + } else { + impersonatedUser = undefined; + } + } + } + + // Fallback to environment variables if no valid session cookie is present + if (!hasSessionCookie) { + customerId = env.CEP_CUSTOMER_ID?.trim() || ""; + impersonatedUser = env.CEP_IMPERSONATE_SUBJECT?.trim() || undefined; + } if (!customerId) { return null; diff --git a/mcp-examples/pocket-cep/src/lib/session.ts b/mcp-examples/pocket-cep/src/lib/session.ts index 44c8b1f..158c128 100644 --- a/mcp-examples/pocket-cep/src/lib/session.ts +++ b/mcp-examples/pocket-cep/src/lib/session.ts @@ -31,14 +31,14 @@ import { SA_EMAIL_DOMAIN } from "./constants"; export async function requireSession() { const auth = getAuth(); const session = await auth.api.getSession({ headers: await headers() }); - + let isStaleAnonymous = false; try { const config = getEnv(); isStaleAnonymous = Boolean( session && - config.AUTH_MODE === "user_oauth" && - session.user.email?.endsWith(`@${SA_EMAIL_DOMAIN}`), + config.AUTH_MODE === "user_oauth" && + session.user.email?.endsWith(`@${SA_EMAIL_DOMAIN}`), ); } catch { // If env cannot be loaded, fallback to safe false diff --git a/mcp-examples/pocket-cep/src/proxy.ts b/mcp-examples/pocket-cep/src/proxy.ts index 396cda8..7827bc3 100644 --- a/mcp-examples/pocket-cep/src/proxy.ts +++ b/mcp-examples/pocket-cep/src/proxy.ts @@ -30,6 +30,8 @@ import { getSessionCookie } from "better-auth/cookies"; import { getEnv, isEnvValidationError } from "@/lib/env"; import { renderEnvErrorHtml, renderMcpUnreachableHtml } from "@/lib/env-error-page"; import { probeMcpServer } from "@/lib/doctor-checks"; +import { COOKIE_SA_SESSION } from "@/lib/sa-session"; +import { verifyJwt } from "@/lib/jwt"; /** * Cached "ok" results for the MCP reachability check. The dashboard @@ -97,10 +99,14 @@ export async function proxy(request: NextRequest) { if (pathname === "/") { return NextResponse.redirect(new URL("/sa-setup", request.url)); } - const hasSaCustomer = Boolean( - request.cookies.get("cep_sa_customer_id")?.value?.trim() || - process.env.CEP_CUSTOMER_ID?.trim(), - ); + const sessionCookieVal = request.cookies.get(COOKIE_SA_SESSION)?.value; + let hasSaCustomer = Boolean(env.CEP_CUSTOMER_ID?.trim()); + if (sessionCookieVal) { + const payload = verifyJwt(sessionCookieVal, env.BETTER_AUTH_SECRET); + if (payload && typeof payload.customerId === "string" && payload.customerId.trim()) { + hasSaCustomer = true; + } + } if (pathname.startsWith("/dashboard") && !hasSaCustomer) { return NextResponse.redirect(new URL("/sa-setup", request.url)); }