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..343428b 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 @@ -98,6 +98,17 @@ describe("/api/auth/sa-config", () => { expect(res.status).toBe(400); }); + it("returns 400 when customerId format is invalid", async () => { + const req = new NextRequest("http://localhost:3000/api/auth/sa-config", { + method: "POST", + body: JSON.stringify({ customerId: "invalid-id" }), + }); + const res = await POST(req); + expect(res.status).toBe(400); + const body = await res.json(); + expect(body.error).toContain("Invalid customerId format"); + }); + it("returns 400 and blocks cookie saving when DWD token minting fails", async () => { mockMintToken.mockRejectedValue(new Error("unauthorized_client: client not authorized")); const req = new NextRequest("http://localhost:3000/api/auth/sa-config", { diff --git a/mcp-examples/pocket-cep/src/__tests__/integration/api/users-route.test.ts b/mcp-examples/pocket-cep/src/__tests__/integration/api/users-route.test.ts index 0f5154c..4528984 100644 --- a/mcp-examples/pocket-cep/src/__tests__/integration/api/users-route.test.ts +++ b/mcp-examples/pocket-cep/src/__tests__/integration/api/users-route.test.ts @@ -29,6 +29,8 @@ vi.mock("next/headers", () => ({ cookies: async () => ({ getAll: () => [], delete: vi.fn(), + has: vi.fn().mockReturnValue(false), + get: vi.fn().mockReturnValue(undefined), }), })); 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/cache-key.test.ts b/mcp-examples/pocket-cep/src/__tests__/unit/cache-key.test.ts new file mode 100644 index 0000000..079b493 --- /dev/null +++ b/mcp-examples/pocket-cep/src/__tests__/unit/cache-key.test.ts @@ -0,0 +1,55 @@ +/** + * @file Unit tests for buildCallerCacheKey. + * + * Verifies that cache keys are correctly generated and partitioned by + * both caller identity (token hash or "sa") and tenant (customerId) + * to prevent cross-tenant data leakage. + */ + +import { describe, it, expect } from "vitest"; +import { createHash } from "node:crypto"; +import { buildCallerCacheKey } from "@/lib/cache-key"; + +describe("buildCallerCacheKey", () => { + const serverUrl = "http://localhost:4000/mcp"; + const token = "mock-oauth-token-123"; + const tokenHash = createHash("sha256").update(token).digest("hex").slice(0, 16); + + it("generates correct key for service_account mode without customerId", () => { + const key = buildCallerCacheKey(serverUrl, undefined); + expect(key).toBe(`${serverUrl}|sa`); + }); + + it("generates correct key for service_account mode with customerId", () => { + const key = buildCallerCacheKey(serverUrl, undefined, "C0111111"); + expect(key).toBe(`${serverUrl}|sa|c:C0111111`); + }); + + it("generates correct key for user_oauth mode without customerId", () => { + const key = buildCallerCacheKey(serverUrl, token); + expect(key).toBe(`${serverUrl}|u:${tokenHash}`); + }); + + it("generates correct key for user_oauth mode with customerId", () => { + const key = buildCallerCacheKey(serverUrl, token, "C0222222"); + expect(key).toBe(`${serverUrl}|u:${tokenHash}|c:C0222222`); + }); + + it("varies key when token changes", () => { + const token2 = "another-token"; + const tokenHash2 = createHash("sha256").update(token2).digest("hex").slice(0, 16); + + const key1 = buildCallerCacheKey(serverUrl, token); + const key2 = buildCallerCacheKey(serverUrl, token2); + + expect(key1).not.toBe(key2); + expect(key2).toBe(`${serverUrl}|u:${tokenHash2}`); + }); + + it("varies key when customerId changes", () => { + const key1 = buildCallerCacheKey(serverUrl, undefined, "C0111111"); + const key2 = buildCallerCacheKey(serverUrl, undefined, "C0222222"); + + expect(key1).not.toBe(key2); + }); +}); diff --git a/mcp-examples/pocket-cep/src/__tests__/unit/mcp-tools.test.ts b/mcp-examples/pocket-cep/src/__tests__/unit/mcp-tools.test.ts index 4de5638..2006655 100644 --- a/mcp-examples/pocket-cep/src/__tests__/unit/mcp-tools.test.ts +++ b/mcp-examples/pocket-cep/src/__tests__/unit/mcp-tools.test.ts @@ -20,6 +20,7 @@ vi.mock("@/lib/mcp-client", () => ({ vi.mock("@/lib/sa-session", () => ({ getServiceAccountConfig: vi.fn().mockResolvedValue(null), + getActiveCustomerId: vi.fn().mockResolvedValue(undefined), })); import { getMcpToolsForAiSdk, invalidateToolCatalog } from "@/lib/mcp-tools"; 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..58580d1 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 @@ -18,6 +18,7 @@ import { DwdScopeVerificationError, mintServiceAccountTokenOrThrow, } from "@/lib/access-token"; +import { CUSTOMER_ID_REGEX } from "@/lib/constants"; /** * Returns the currently configured Service Account tenant credentials. @@ -61,6 +62,13 @@ export async function POST(request: NextRequest) { return NextResponse.json({ error: "customerId is required" }, { status: 400 }); } + if (!CUSTOMER_ID_REGEX.test(customerId)) { + return NextResponse.json( + { error: "Invalid customerId format. Must start with 'C' (e.g. C01234567)." }, + { status: 400 }, + ); + } + const impersonatedUser = body.impersonatedUser?.trim() || ""; clearServiceAccountTokenCache(); diff --git a/mcp-examples/pocket-cep/src/app/api/insights/risky-activity/route.ts b/mcp-examples/pocket-cep/src/app/api/insights/risky-activity/route.ts index 77a49f2..b8c9586 100644 --- a/mcp-examples/pocket-cep/src/app/api/insights/risky-activity/route.ts +++ b/mcp-examples/pocket-cep/src/app/api/insights/risky-activity/route.ts @@ -14,6 +14,8 @@ import { getGoogleAccessToken } from "@/lib/access-token"; import { getEnv } from "@/lib/env"; import { callMcpTool } from "@/lib/mcp-client"; import { requireSession } from "@/lib/session"; +import { getActiveCustomerId } from "@/lib/sa-session"; +import { buildCallerCacheKey } from "@/lib/cache-key"; import { isAuthError, toAuthError } from "@/lib/auth-errors"; import { CACHE_TAGS, getOrFetch } from "@/lib/server-cache"; import { summarizeChromeActivity } from "@/lib/activity-summarizer"; @@ -38,9 +40,10 @@ export async function POST(request: Request) { const selectedUser = body.selectedUser ?? ""; const config = getEnv(); const accessToken = await getGoogleAccessToken(); - try { - const cacheKey = `insights:risky-activity:${selectedUser}`; + const customerId = await getActiveCustomerId(); + const callerKey = buildCallerCacheKey(config.MCP_SERVER_URL, accessToken, customerId); + const cacheKey = `insights:risky-activity:${callerKey}:${selectedUser}`; const summary = await getOrFetch({ key: cacheKey, ttlMs: INSIGHT_TTL_MS, diff --git a/mcp-examples/pocket-cep/src/app/api/prompts/route.ts b/mcp-examples/pocket-cep/src/app/api/prompts/route.ts index cc1fbb8..f03deac 100644 --- a/mcp-examples/pocket-cep/src/app/api/prompts/route.ts +++ b/mcp-examples/pocket-cep/src/app/api/prompts/route.ts @@ -15,6 +15,7 @@ import { getEnv } from "@/lib/env"; import { getMcpPrompt, listMcpPrompts } from "@/lib/mcp-client"; import { buildCallerCacheKey } from "@/lib/cache-key"; import { requireSession } from "@/lib/session"; +import { getActiveCustomerId } from "@/lib/sa-session"; import { LOG_TAGS } from "@/lib/constants"; import { getErrorMessage } from "@/lib/errors"; import { getOrFetch, CACHE_TAGS } from "@/lib/server-cache"; @@ -27,9 +28,9 @@ export async function GET() { const config = getEnv(); const accessToken = await getGoogleAccessToken(); - const callerKey = buildCallerCacheKey(config.MCP_SERVER_URL, accessToken); - try { + const customerId = await getActiveCustomerId(); + const callerKey = buildCallerCacheKey(config.MCP_SERVER_URL, accessToken, customerId); const prompts = await getOrFetch({ key: `prompts:${callerKey}`, ttlMs: CATALOG_TTL_MS, diff --git a/mcp-examples/pocket-cep/src/app/api/tools/route.ts b/mcp-examples/pocket-cep/src/app/api/tools/route.ts index e939554..ab45410 100644 --- a/mcp-examples/pocket-cep/src/app/api/tools/route.ts +++ b/mcp-examples/pocket-cep/src/app/api/tools/route.ts @@ -24,6 +24,7 @@ import { getGoogleAccessToken } from "@/lib/access-token"; import { getEnv } from "@/lib/env"; import { buildCallerCacheKey } from "@/lib/cache-key"; import { requireSession } from "@/lib/session"; +import { getActiveCustomerId } from "@/lib/sa-session"; import { LOG_TAGS } from "@/lib/constants"; import { getErrorMessage } from "@/lib/errors"; import { getOrFetch, CACHE_TAGS } from "@/lib/server-cache"; @@ -45,9 +46,9 @@ export async function GET() { const config = getEnv(); const accessToken = await getGoogleAccessToken(); - const callerKey = buildCallerCacheKey(config.MCP_SERVER_URL, accessToken); - try { + const customerId = await getActiveCustomerId(); + const callerKey = buildCallerCacheKey(config.MCP_SERVER_URL, accessToken, customerId); const tools = await getOrFetch({ key: `tools:${callerKey}`, ttlMs: TOOLS_TTL_MS, diff --git a/mcp-examples/pocket-cep/src/app/api/users/route.ts b/mcp-examples/pocket-cep/src/app/api/users/route.ts index c0df4c2..e47f09d 100644 --- a/mcp-examples/pocket-cep/src/app/api/users/route.ts +++ b/mcp-examples/pocket-cep/src/app/api/users/route.ts @@ -19,6 +19,7 @@ import { getGoogleAccessToken } from "@/lib/access-token"; import { searchUsers, buildAdminQuery, type DirectoryUser } from "@/lib/admin-sdk"; import { buildCallerCacheKey } from "@/lib/cache-key"; import { requireSession } from "@/lib/session"; +import { getActiveCustomerId } from "@/lib/sa-session"; import { conditionalJson } from "@/lib/http-cache"; import { getOrFetch, CACHE_TAGS } from "@/lib/server-cache"; import { respondWithApiError, unauthenticatedResponse } from "@/lib/api-response"; @@ -38,9 +39,9 @@ export async function GET(request: NextRequest) { const config = getEnv(); const accessToken = await getGoogleAccessToken(); const adminQuery = buildAdminQuery(query); - const callerKey = buildCallerCacheKey(config.MCP_SERVER_URL, accessToken); - try { + const customerId = await getActiveCustomerId(); + const callerKey = buildCallerCacheKey(config.MCP_SERVER_URL, accessToken, customerId); const users = await getOrFetch({ key: `users:${callerKey}:${query.trim().toLowerCase()}`, ttlMs: USERS_TTL_MS, diff --git a/mcp-examples/pocket-cep/src/lib/activity-data.ts b/mcp-examples/pocket-cep/src/lib/activity-data.ts index 564be1d..2cdd4f1 100644 --- a/mcp-examples/pocket-cep/src/lib/activity-data.ts +++ b/mcp-examples/pocket-cep/src/lib/activity-data.ts @@ -67,7 +67,8 @@ export async function getCachedActivity( } const saConfig = config.AUTH_MODE === "service_account" ? await getServiceAccountConfig() : null; const impersonatedUser = saConfig?.impersonatedUser; - const callerKey = buildCallerCacheKey(config.MCP_SERVER_URL, accessToken); + const customerId = saConfig?.customerId; + const callerKey = buildCallerCacheKey(config.MCP_SERVER_URL, accessToken, customerId); return getOrFetch({ key: `activity:${callerKey}:${days}`, @@ -100,7 +101,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/cache-key.ts b/mcp-examples/pocket-cep/src/lib/cache-key.ts index df128a5..ced5e1f 100644 --- a/mcp-examples/pocket-cep/src/lib/cache-key.ts +++ b/mcp-examples/pocket-cep/src/lib/cache-key.ts @@ -1,26 +1,31 @@ /** - * @file Shared cache-key builder for per-caller in-process caches. + * @file Shared cache-key builder for per-caller and per-tenant in-process caches. * * Several server-side caches (MCP tool catalog, MCP prompt catalog, - * Admin Reports activity) isolate entries by caller identity so that - * user_oauth sessions don't share cached data with one another or with - * service_account (ADC) flows. + * Admin Reports activity) isolate entries by caller identity and tenant so that + * user_oauth sessions and service_account sessions don't share cached data + * across different Google Workspace customers (domains) or users. * - * service_account callers share a single `|sa` entry (no token). Each + * service_account callers partition by `|sa|c:`. Each * user_oauth caller gets a per-token entry keyed by a truncated - * SHA-256 — the raw access token never lands in the Map. + * SHA-256, optionally partitioned by `customerId` if available. */ import { createHash } from "node:crypto"; /** - * Builds a cache key of the form `${serverUrl}|sa` (service_account) or - * `${serverUrl}|u:<16-char-hash>` (user_oauth). The hash length is + * Builds a cache key of the form `${serverUrl}|sa|c:${customerId}` (service_account) or + * `${serverUrl}|u:${hash}|c:${customerId}` (user_oauth). The hash length is * enough to make collisions astronomically unlikely while keeping keys * short in logs and heap dumps. */ -export function buildCallerCacheKey(serverUrl: string, accessToken: string | undefined): string { - if (!accessToken) return `${serverUrl}|sa`; +export function buildCallerCacheKey( + serverUrl: string, + accessToken: string | undefined, + customerId?: string, +): string { + const customerPart = customerId ? `|c:${customerId}` : ""; + if (!accessToken) return `${serverUrl}|sa${customerPart}`; const hash = createHash("sha256").update(accessToken).digest("hex").slice(0, 16); - return `${serverUrl}|u:${hash}`; + return `${serverUrl}|u:${hash}${customerPart}`; } diff --git a/mcp-examples/pocket-cep/src/lib/constants.ts b/mcp-examples/pocket-cep/src/lib/constants.ts index cef0927..fd0d2be 100644 --- a/mcp-examples/pocket-cep/src/lib/constants.ts +++ b/mcp-examples/pocket-cep/src/lib/constants.ts @@ -14,6 +14,12 @@ */ export const SA_EMAIL_DOMAIN = "service-account.local"; +/** + * Regex pattern for validating Google Workspace Customer IDs. + * Must start with 'C' followed by alphanumeric characters (e.g. C01234567). + */ +export const CUSTOMER_ID_REGEX = /^C[a-zA-Z0-9]+$/; + /** * DOM id for the header's user-search input. Referenced from the app * bar (focus button), dashboard (`/` keyboard shortcut), and the diff --git a/mcp-examples/pocket-cep/src/lib/env.ts b/mcp-examples/pocket-cep/src/lib/env.ts index c3f1654..a70ac12 100644 --- a/mcp-examples/pocket-cep/src/lib/env.ts +++ b/mcp-examples/pocket-cep/src/lib/env.ts @@ -10,7 +10,7 @@ */ import { z } from "zod"; -import { DEFAULT_MCP_URL } from "./constants"; +import { DEFAULT_MCP_URL, CUSTOMER_ID_REGEX } from "./constants"; /** * Regex-validated Google OAuth client ID. The strict format check catches @@ -47,7 +47,7 @@ const baseFields = { .or(z.literal("")), CEP_CUSTOMER_ID: z .string() - .regex(/^C[a-zA-Z0-9]+$/, "CEP_CUSTOMER_ID must start with 'C' (e.g. C01234567).") + .regex(CUSTOMER_ID_REGEX, "CEP_CUSTOMER_ID must start with 'C' (e.g. C01234567).") .optional() .or(z.literal("")), }; diff --git a/mcp-examples/pocket-cep/src/lib/mcp-tools.ts b/mcp-examples/pocket-cep/src/lib/mcp-tools.ts index aa2d607..98cf5f1 100644 --- a/mcp-examples/pocket-cep/src/lib/mcp-tools.ts +++ b/mcp-examples/pocket-cep/src/lib/mcp-tools.ts @@ -17,7 +17,7 @@ import { callMcpTool, listMcpTools, type McpToolDefinition } from "./mcp-client" import { toAuthError } from "./auth-errors"; import { buildCallerCacheKey } from "./cache-key"; import { LOG_TAGS } from "./constants"; -import { getServiceAccountConfig } from "./sa-session"; +import { getServiceAccountConfig, getActiveCustomerId } from "./sa-session"; const TOOL_CATALOG_TTL_MS = 5 * 60 * 1000; @@ -30,8 +30,9 @@ const toolCatalogCache = new Map { - const key = buildCallerCacheKey(serverUrl, accessToken); + const key = buildCallerCacheKey(serverUrl, accessToken, customerId); const now = Date.now(); const cached = toolCatalogCache.get(key); if (cached && cached.expiresAt > now) return cached.tools; @@ -66,7 +67,7 @@ export async function getMcpToolsForAiSdk( accessToken?: string, customerId?: string, ): Promise { - const mcpTools = await getCachedToolCatalog(serverUrl, accessToken); + const mcpTools = await getCachedToolCatalog(serverUrl, accessToken, customerId); const tools: ToolSet = {}; for (const t of mcpTools) { @@ -144,7 +145,8 @@ function extractErrorText(content: unknown): string { * for injection into reference prompts (e.g. follow-up suggestion brainstorming). */ export async function getMcpToolsSummary(serverUrl: string, accessToken?: string): Promise { - const mcpTools = await getCachedToolCatalog(serverUrl, accessToken); + const customerId = await getActiveCustomerId(); + const mcpTools = await getCachedToolCatalog(serverUrl, accessToken, customerId); return mcpTools .map((t) => { const schema = t.inputSchema as