Lab 14 — Add bearer auth, per-tenant isolation, and rate limit to PKB-TS

Lesson 14 · MCP Servers · ~75 min, code changes committed under mcp-servers/sample-server/typescript/

L13 left PKB-TS reachable from the public internet behind Caddy. Anyone who knows the URL can open an SSE session and call echo. That is the right place to stop for a personal demo and the wrong place to stop for anything else. This lab adds the three orthogonal guarantees that make a remote MCP server safe to expose: identity (bearer token → tenant id), isolation (per-tenant note store), and quota (token-bucket rate limit per tenant). By the end, two Inspector tabs with two different tokens will see two independent note stores, and a tight call loop will trip the per-tenant rate limit.

Prereqs

• L13 lab complete. docker compose up brings up Caddy + pkb; Inspector can call echo over HTTPS at localhost (tls internal path) or your real domain. /healthz returns {"ok":true,"service":"pkb"}. The four-file TS layout (server.ts, tools.ts, schemas.ts, logger.ts) is in src/ and npm run build exits clean.
• Inspector with custom-header support. @modelcontextprotocol/inspector recent enough to set the Authorization header on a connection. If your version doesn't have a header field, the troubleshooting section at the bottom shows how to drive the smoke test with curl instead.
• Working directory. All commands run from mcp-servers/sample-server/typescript/ unless noted.

Step 1 — Write src/auth.ts

Create src/auth.ts. The middleware extracts the bearer token from the Authorization header, looks it up in a token → tenant map seeded from PKB_TOKENS, and sets req.tenant on success or returns 401 on failure:

import type { Request, Response, NextFunction } from "express";
import { toolLogger } from "./logger.js";

const log = toolLogger("auth");

// Token -> tenant id. For a real deployment swap for a DB / JWT / OAuth.
// Format: PKB_TOKENS="alice:tok_alice_xxx,bob:tok_bob_yyy"
function parseTokens(): Map<string, string> {
  const raw = process.env.PKB_TOKENS ?? "";
  const out = new Map<string, string>();
  for (const pair of raw.split(",").map((s) => s.trim()).filter(Boolean)) {
    const [tenant, token] = pair.split(":");
    if (tenant && token) out.set(token, tenant);
  }
  return out;
}

const TOKENS = parseTokens();

export interface AuthedRequest extends Request {
  tenant?: string;
}

export function bearerAuth(req: AuthedRequest, res: Response, next: NextFunction): void {
  const hdr = req.header("authorization") ?? "";
  const m = hdr.match(/^Bearer\s+(\S+)$/i);
  if (!m) {
    log.warn({ phase: "denied", reason: "no_bearer" });
    res.status(401).json({ error: "missing bearer token" });
    return;
  }
  const tenant = TOKENS.get(m[1]);
  if (!tenant) {
    log.warn({ phase: "denied", reason: "unknown_token" });
    res.status(401).json({ error: "invalid token" });
    return;
  }
  req.tenant = tenant;
  next();
}

Three things worth pausing on:

• parseTokens() runs once at module load. The env var is read into a Map<token, tenant> at startup; rotating tokens means restarting the process. That is the simplest model that works; production swaps the in-memory map for an OAuth introspection call or a JWT verifier with key caching.
• Every denial is logged with a reason field. no_bearer vs unknown_token are two different attack shapes; counting them separately gives you an abuse signal in L15's audit pipeline. The token value itself is never logged — only the fact of failure.
• The handler returns instead of falling through. Express middleware must either call next() or end the response. Forgetting the return after res.status(401).json(...) means execution continues into the SSE handler with req.tenant === undefined, and you leak an unauthenticated session.

Step 2 — Write src/store.ts

Create src/store.ts. Two levels of Map: tenant id → (note id → body):

// Per-tenant in-memory note store. Production swap: SQLite/Postgres keyed by tenant id.

type NoteId = string;
type Notes = Map<NoteId, string>;

const STORES = new Map<string, Notes>();

export function storeFor(tenant: string): Notes {
  let s = STORES.get(tenant);
  if (!s) {
    s = new Map<NoteId, string>([
      ["hello", `Hello note for ${tenant}`],
    ]);
    STORES.set(tenant, s);
  }
  return s;
}

export function tenantCount(): number {
  return STORES.size;
}

storeFor(tenant) is the only path tool code uses to reach notes. On first access for a new tenant, a placeholder hello note is seeded so the search/fetch tools always have something to return in the smoke test — otherwise an empty bucket is hard to tell apart from a broken store. The tenantCount() helper is purely for diagnostics; tools never read it.

Step 3 — Write src/rateLimit.ts

Create src/rateLimit.ts. A token bucket per tenant: CAPACITY tokens of burst, REFILL_PER_MIN sustained rate, refill computed lazily on each check():

// Token-bucket rate limit per tenant. Default: 60 req/min refill, 10-req burst.

interface Bucket { tokens: number; lastRefill: number; }

const BUCKETS = new Map<string, Bucket>();
const CAPACITY = Number(process.env.PKB_RL_CAPACITY ?? 10);
const REFILL_PER_MIN = Number(process.env.PKB_RL_REFILL_PER_MIN ?? 60);

export interface RateLimitResult {
  allowed: boolean;
  remaining: number;
  retryAfterMs: number;
}

export function check(tenant: string): RateLimitResult {
  const now = Date.now();
  let b = BUCKETS.get(tenant);
  if (!b) { b = { tokens: CAPACITY, lastRefill: now }; BUCKETS.set(tenant, b); }
  // Refill since last check, capped at CAPACITY.
  const elapsedMin = (now - b.lastRefill) / 60_000;
  b.tokens = Math.min(CAPACITY, b.tokens + elapsedMin * REFILL_PER_MIN);
  b.lastRefill = now;
  if (b.tokens >= 1) {
    b.tokens -= 1;
    return { allowed: true, remaining: Math.floor(b.tokens), retryAfterMs: 0 };
  }
  const retryAfterMs = Math.ceil(((1 - b.tokens) / REFILL_PER_MIN) * 60_000);
  return { allowed: false, remaining: 0, retryAfterMs };
}

How the math falls out:

• Defaults: 10-token burst, 60 / minute sustained. A fresh tenant can make 10 calls right now without waiting. Sustained above 1 / second and the bucket runs dry; each subsequent call returns a non-zero retryAfterMs until enough time has passed to refill at least one token.
• Lazy refill, not a background timer. No setInterval, no per-tenant ticker. Refill is computed on each check() from elapsed * rate. This scales to thousands of tenants without thousands of timers.
• Configurable via env. PKB_RL_CAPACITY / PKB_RL_REFILL_PER_MIN let you crank the burst and rate up for production or down for the lab where you want to see the limit trip in a few curls.

Step 4 — Add the PKB tools to src/tools.ts

Append three tenant-scoped tools to src/tools.ts. Each takes (tenant: string, rawArgs: unknown) — the tenant is the function's first argument, not a Zod-validated field on rawArgs:

import { storeFor } from "./store.js";
import { z } from "zod";

const SearchInput = z.object({ query: z.string().min(1).max(1000), limit: z.number().int().min(1).max(50).default(10) });
const FetchInput = z.object({ id: z.string().min(1).max(200) });
const WriteInput = z.object({ id: z.string().min(1).max(200), body: z.string().min(1).max(50_000) });

export async function pkbSearch(tenant: string, rawArgs: unknown): Promise<string[]> {
  const log = toolLogger("pkb.search").child({ tenant });
  const t0 = performance.now();
  log.info({ phase: "start" });
  try {
    const args = SearchInput.parse(rawArgs);
    const store = storeFor(tenant);
    const out: string[] = [];
    for (const [id, body] of store.entries()) {
      if (out.length >= args.limit) break;
      if (id.includes(args.query) || body.includes(args.query)) out.push(id);
    }
    log.info({ phase: "ok", dur_ms: +(performance.now() - t0).toFixed(1), n: out.length });
    return out;
  } catch (err) {
    log.warn({ phase: "err", dur_ms: +(performance.now() - t0).toFixed(1), type: (err as Error).name });
    throw err;
  }
}

export async function pkbFetch(tenant: string, rawArgs: unknown): Promise<string> {
  const log = toolLogger("pkb.fetch").child({ tenant });
  const t0 = performance.now();
  log.info({ phase: "start" });
  try {
    const args = FetchInput.parse(rawArgs);
    const store = storeFor(tenant);
    const body = store.get(args.id);
    if (body === undefined) {
      log.warn({ phase: "err", dur_ms: +(performance.now() - t0).toFixed(1), type: "UnknownId" });
      throw new Error(`unknown note id: ${args.id}`);
    }
    log.info({ phase: "ok", dur_ms: +(performance.now() - t0).toFixed(1), body_len: body.length });
    return body;
  } catch (err) {
    if (!(err as Error).message?.startsWith("unknown note id")) {
      log.warn({ phase: "err", dur_ms: +(performance.now() - t0).toFixed(1), type: (err as Error).name });
    }
    throw err;
  }
}

export async function pkbWrite(tenant: string, rawArgs: unknown): Promise<string> {
  const log = toolLogger("pkb.write").child({ tenant });
  const t0 = performance.now();
  log.info({ phase: "start" });
  try {
    const args = WriteInput.parse(rawArgs);
    const store = storeFor(tenant);
    store.set(args.id, args.body);
    log.info({ phase: "ok", dur_ms: +(performance.now() - t0).toFixed(1), body_len: args.body.length });
    return `saved: ${args.id}`;
  } catch (err) {
    log.warn({ phase: "err", dur_ms: +(performance.now() - t0).toFixed(1), type: (err as Error).name });
    throw err;
  }
}

Step 5 — Wire it all in src/server.ts

Open src/server.ts. At the top, import the new modules and the three new tool handlers:

import { bearerAuth, type AuthedRequest } from "./auth.js";
import { check as rlCheck } from "./rateLimit.js";
import { echoTool, pkbSearch, pkbFetch, pkbWrite } from "./tools.js";

Replace the CallTool handler with one that resolves the tenant from the session, checks the rate limit, and dispatches to the right tool:

server.setRequestHandler(CallToolRequestSchema, async (req, extra) => {
  const tenant = (extra as any)?.tenant ?? "anonymous";
  const rl = rlCheck(tenant);
  if (!rl.allowed) {
    return {
      isError: true,
      content: [{ type: "text", text: `rate limit exceeded; retry after ${rl.retryAfterMs} ms` }],
    };
  }
  const { name, arguments: args } = req.params;
  try {
    if (name === "echo") return { content: [{ type: "text", text: await echoTool(args) }] };
    if (name === "pkb.search") {
      const ids = await pkbSearch(tenant, args);
      return { content: [{ type: "text", text: JSON.stringify(ids) }] };
    }
    if (name === "pkb.fetch") return { content: [{ type: "text", text: await pkbFetch(tenant, args) }] };
    if (name === "pkb.write") return { content: [{ type: "text", text: await pkbWrite(tenant, args) }] };
    return { isError: true, content: [{ type: "text", text: `unknown tool: ${name}` }] };
  } catch (err) {
    return { isError: true, content: [{ type: "text", text: (err as Error).message }] };
  }
});

Then replace the single-session SSE block in the http branch of main() with the per-tenant session map plus middleware:

if (mode === "http") {
  const app = express();
  const port = Number(process.env.PORT ?? 8080);
  const sessions = new Map<string, { transport: SSEServerTransport; tenant: string }>();

  app.get("/healthz", (_req, res) => res.json({ ok: true, service: "pkb" }));

  app.get("/sse", bearerAuth, async (req: AuthedRequest, res) => {
    const tenant = req.tenant!;
    const transport = new SSEServerTransport("/messages", res);
    sessions.set(transport.sessionId, { transport, tenant });
    (server as any).tenant = tenant;   // simple plumbing for the handler closure
    res.on("close", () => sessions.delete(transport.sessionId));
    await server.connect(transport);
    logger.info({ phase: "session_open", tenant, session_id: transport.sessionId });
  });

  app.post("/messages", express.json(), async (req, res) => {
    const sid = String(req.query.sessionId ?? "");
    const sess = sessions.get(sid);
    if (!sess) { res.status(404).send("unknown session"); return; }
    // Inject the tenant into the request handler context.
    (req as any).tenant = sess.tenant;
    await sess.transport.handlePostMessage(req, res, req.body);
  });

  app.listen(port, () => logger.info({ phase: "ready", transport: "sse", port }));
  return;
}

Finally, extend the ListTools handler so the new tools show up in Inspector. Add three entries alongside echo: pkb.search (input: query string + limit int), pkb.fetch (input: id string), and pkb.write (input: id string + body string). The schemas mirror the Zod definitions from Step 4 — no tenant field anywhere.

Step 6 — Rebuild and start with two tenants

Compile and start the server with Alice and Bob in the token map:

npm run build
PKB_TRANSPORT=http \
  PKB_TOKENS="alice:tok_alice,bob:tok_bob" \
  node dist/server.js &
sleep 1

Confirm /healthz still answers — it is unauthenticated on purpose so load balancers can probe without a token:

curl -s localhost:8080/healthz
# {"ok":true,"service":"pkb"}

Step 7 — Verify auth: anonymous is rejected

Hit /sse with no Authorization header:

curl -i localhost:8080/sse

Expect a clean HTTP/1.1 401 with body {"error":"missing bearer token"}. In the server logs you should see one line per attempt:

{"level":"warn","tool":"auth","phase":"denied","reason":"no_bearer"}

Now try a bogus token:

curl -i -H "Authorization: Bearer not-a-real-token" localhost:8080/sse

Still 401, but the log line now reads "reason":"unknown_token". The two reasons being separable is the abuse-signal hook L15's audit pipeline reads.

Step 8 — Verify isolation: Alice and Bob get separate notes

Open Inspector in two tabs (or two browser windows), each pointed at http://localhost:8080/sse with a different bearer header. In recent Inspector versions there is an Authorization field next to the URL; older versions require the curl-based workaround in the troubleshooting section.

• Tab 1 — header: Authorization: Bearer tok_alice.
• Tab 2 — header: Authorization: Bearer tok_bob.

In Alice's tab, call pkb.write with arguments:

{ "id": "secret", "body": "alice-secret" }

Then call pkb.fetch with {"id":"secret"} — you should see "alice-secret" come back. Switch to Bob's tab and call pkb.fetch with the same args. You should get an isError result with text unknown note id: secret — Bob's store is empty for that id. Now call pkb.write in Bob's tab with {"id":"secret","body":"bob-secret"} and pkb.fetch it — you get "bob-secret", never Alice's value. Switch back to Alice and re-fetch — still "alice-secret". Cross-tenant reads cannot happen because the tenant id flows from the authenticated session, not from the tool args.

In the server logs, every tool-call line is tagged with the right tenant:

{"level":"info","tool":"pkb.write","tenant":"alice","phase":"ok","dur_ms":0.3,"body_len":12}
{"level":"info","tool":"pkb.fetch","tenant":"bob","phase":"start"}
{"level":"warn","tool":"pkb.fetch","tenant":"bob","phase":"err","dur_ms":0.2,"type":"UnknownId"}

Step 9 — Verify rate limit

With Alice's session still open in Inspector, fire pkb.search as fast as the UI lets you. After ~10 calls in a second you should see an isError: true content with text rate limit exceeded; retry after <N> ms. To make this reliable from a shell, run a small concurrent burst against /messages:

# Replace SID with the sessionId Inspector negotiated (see browser devtools or the URL)
SID="..."
seq 1 50 | xargs -P 20 -I{} curl -s -o /dev/null -w "%{http_code} " \
  -H "Content-Type: application/json" \
  -d '{"jsonrpc":"2.0","id":{},"method":"tools/call","params":{"name":"pkb.search","arguments":{"query":"hello"}}}' \
  "http://localhost:8080/messages?sessionId=$SID"
echo

The first ten or so return 200; the rest still return 200 at the HTTP layer (the rate-limit denial is an MCP-level isError, not an HTTP 429) but the response body now contains the rate-limit text. In the logs you'll see successful calls flip to denials:

{"level":"info","tool":"pkb.search","tenant":"alice","phase":"ok","dur_ms":0.1,"n":1}
... (10 of these)
{"level":"info","tool":"pkb.search","tenant":"alice","phase":"start"}
# (no matching phase=ok because CallTool short-circuited to isError before pkbSearch ran)

Wait one second — the bucket refills at one token per second — and a single follow-up call succeeds again. If you crank PKB_RL_CAPACITY=2 and restart, the limit trips on the third call instead of the eleventh; useful when you want to see the behaviour in two curls.

Step 10 — Tear down

Stop the local server:

kill %1
# or: pkill -f "node dist/server.js"

If you want to test the full container path next, rebuild the image with the new code and start it with the same env vars:

docker build -t pkb-ts:l14 .
docker run --rm -p 8080:8080 \
  -e PKB_TRANSPORT=http \
  -e PKB_TOKENS="alice:tok_alice,bob:tok_bob" \
  -d --name pkb-l14 pkb-ts:l14
sleep 2
curl -s localhost:8080/healthz
curl -i localhost:8080/sse                                     # 401
curl -i -H "Authorization: Bearer tok_alice" localhost:8080/sse  # 200 SSE stream
docker stop pkb-l14

Same three checks (healthz works, anonymous is 401, Alice's token flips it to 200) reproduced inside the L13 container.

Step 11 — Audit checklist

What you should leave with

• A PKB-TS server that requires a bearer token on /sse, with denials logged by reason. Anonymous and bad tokens both return 401 with distinct log lines.
• Per-tenant note isolation verified across two Inspector tabs with two tokens. Cross-tenant reads are impossible because the tenant is sourced from the authenticated session, not the tool args.
• A token-bucket rate limiter that returns isError with a retryAfterMs hint when a single tenant bursts past 10 calls/second.
• Three new tools advertised to clients (pkb.search, pkb.fetch, pkb.write) whose schemas carry note ids and bodies but never carry identity.
• Four small, focused new files (src/auth.ts, src/store.ts, src/rateLimit.ts, and the bottom half of src/tools.ts) that you can read end-to-end in under five minutes — the production version of each will be 3–10x larger, but the shape stays the same.

Troubleshooting