Secure TypeScript MCP server

import { createServer } from "node:http";
import { createHash, randomUUID, timingSafeEqual } from "node:crypto";

import { context, trace, SpanStatusCode } from "@opentelemetry/api";
import { createMcpHandler, McpServer } from "@modelcontextprotocol/server";
import { toNodeHandler } from "@modelcontextprotocol/node";
import { jwtVerify, createRemoteJWKSet, type JWTPayload } from "jose";
import * as z from "zod/v4";

type Claims = JWTPayload & {
  sub: string;
  workspace_id: string;
  scope?: string;
};

type CodeHit = {
  path: string;
  startLine: number;
  endLine: number;
  excerpt: string;
  score: number;
};

const config = {
  issuer: mustEnv("OAUTH_ISSUER"),
  audience: mustEnv("MCP_AUDIENCE"),
  jwksUrl: new URL(mustEnv("OAUTH_JWKS_URL")),
  port: Number(process.env.PORT ?? 8787),
};

const jwks = createRemoteJWKSet(config.jwksUrl);
const tracer = trace.getTracer("secure-code-search-mcp");
const limiter = new SlidingWindowLimiter(60_000, 30);

function mustEnv(name: string): string {
  const value = process.env[name]?.trim();
  if (!value) throw new Error(`${name} is required`);
  return value;
}

function scopes(claims: Claims): Set<string> {
  return new Set((claims.scope ?? "").split(/\s+/).filter(Boolean));
}

async function authenticate(request: Request): Promise<Claims> {
  const header = request.headers.get("authorization") ?? "";
  const match = /^Bearer\s+(.+)$/i.exec(header);
  if (!match) throw new HttpError(401, "Bearer token required");

  const { payload } = await jwtVerify(match[1], jwks, {
    issuer: config.issuer,
    audience: config.audience,
    algorithms: ["RS256", "ES256"],
    clockTolerance: 5,
  });
  if (!payload.sub || typeof payload.workspace_id !== "string") {
    throw new HttpError(403, "Token lacks tenant identity");
  }
  return payload as Claims;
}

function requireScope(claims: Claims, required: string): void {
  if (!scopes(claims).has(required)) {
    throw new HttpError(403, `Missing scope: ${required}`);
  }
}

class HttpError extends Error {
  constructor(readonly status: number, message: string) {
    super(message);
  }
}

class SlidingWindowLimiter {
  private readonly requests = new Map<string, number[]>();

  constructor(
    private readonly windowMs: number,
    private readonly maximum: number,
  ) {}

  accept(key: string, now = Date.now()): boolean {
    const floor = now - this.windowMs;
    const recent = (this.requests.get(key) ?? []).filter((time) => time > floor);
    if (recent.length >= this.maximum) return false;
    recent.push(now);
    this.requests.set(key, recent);
    return true;
  }
}

class CodeIndex {
  async search(input: {
    workspaceId: string;
    query: string;
    limit: number;
    signal: AbortSignal;
  }): Promise<CodeHit[]> {
    input.signal.throwIfAborted();
    // Production implementation executes a tenant-scoped hybrid query.
    return [
      {
        path: "src/jobs/claim.ts",
        startLine: 41,
        endLine: 67,
        excerpt: "const job = await claimNextJob(workspaceId);",
        score: 0.91,
      },
    ].filter(() => input.workspaceId.length > 0).slice(0, input.limit);
  }
}

const codeIndex = new CodeIndex();

function stableRequestKey(claims: Claims, query: string): string {
  return createHash("sha256")
    .update(`${claims.workspace_id}\0${claims.sub}\0${query}`)
    .digest("hex");
}

function equalKey(left: string, right: string): boolean {
  const a = Buffer.from(left);
  const b = Buffer.from(right);
  return a.length === b.length && timingSafeEqual(a, b);
}

function buildMcpServer(claims: Claims): McpServer {
  const server = new McpServer({
    name: "secure-code-search",
    version: "2.0.0",
  });

  server.registerTool(
    "search_code",
    {
      title: "Search repository code",
      description:
        "Search code visible to the authenticated workspace. Never accepts a tenant ID from tool input.",
      inputSchema: z.object({
        query: z.string().trim().min(3).max(500),
        limit: z.number().int().min(1).max(20).default(8),
        requestKey: z.string().regex(/^[a-f0-9]{64}$/).optional(),
      }),
    },
    async ({ query, limit, requestKey }, extra) => {
      requireScope(claims, "code:read");
      const expectedKey = stableRequestKey(claims, query);
      if (requestKey && !equalKey(requestKey, expectedKey)) {
        throw new Error("Request key does not match authenticated input");
      }

      const rateKey = `${claims.workspace_id}:${claims.sub}`;
      if (!limiter.accept(rateKey)) throw new Error("Rate limit exceeded");

      return tracer.startActiveSpan("mcp.search_code", async (span) => {
        span.setAttributes({
          "enduser.id": claims.sub,
          "hattrick.workspace_id": claims.workspace_id,
          "mcp.tool.name": "search_code",
          "mcp.request_id": randomUUID(),
        });
        try {
          const deadline = AbortSignal.timeout(8_000);
          const signal = AbortSignal.any([deadline, extra.signal]);
          const hits = await codeIndex.search({
            workspaceId: claims.workspace_id,
            query,
            limit,
            signal,
          });
          span.setAttribute("mcp.result_count", hits.length);
          return {
            structuredContent: { hits },
            content: [{ type: "text", text: JSON.stringify({ hits }) }],
          };
        } catch (error) {
          span.recordException(error as Error);
          span.setStatus({ code: SpanStatusCode.ERROR });
          throw error;
        } finally {
          span.end();
        }
      });
    },
  );

  server.registerTool(
    "explain_access",
    {
      description: "Explain the caller's effective MCP permissions.",
      inputSchema: z.object({}),
    },
    async () => ({
      structuredContent: {
        subject: claims.sub,
        workspaceId: claims.workspace_id,
        scopes: [...scopes(claims)].sort(),
      },
      content: [{ type: "text", text: "Access is derived from the bearer token." }],
    }),
  );
  return server;
}

const mcp = createMcpHandler(({ authInfo }) => {
  if (!authInfo) throw new Error("Authenticated request context missing");
  return buildMcpServer(authInfo.extra?.claims as Claims);
}, {
  responseMode: "sse",
  legacy: "stateless",
});

async function route(request: Request): Promise<Response> {
  const requestId = request.headers.get("x-request-id") ?? randomUUID();
  if (new URL(request.url).pathname === "/healthz") {
    return Response.json({ ok: true, requestId });
  }
  if (new URL(request.url).pathname !== "/mcp") {
    return new Response("Not found", { status: 404 });
  }

  try {
    const claims = await authenticate(request);
    const headers = new Headers(request.headers);
    headers.set("x-auth-subject", claims.sub);
    const authenticated = new Request(request, { headers });
    return await context.with(trace.setSpan(context.active(), undefined), () =>
      mcp.fetch(authenticated, {
        authInfo: {
          token: "verified",
          clientId: claims.sub,
          scopes: [...scopes(claims)],
          extra: { claims },
        },
      })
    );
  } catch (error) {
    const status = error instanceof HttpError ? error.status : 500;
    const message = status === 500 ? "Internal server error" : error.message;
    return Response.json(
      { error: message, requestId },
      {
        status,
        headers: status === 401
          ? { "WWW-Authenticate": `Bearer resource_metadata="${config.issuer}/.well-known/oauth-protected-resource"` }
          : undefined,
      },
    );
  }
}

const nodeHandler = toNodeHandler(route);
const httpServer = createServer(nodeHandler);
httpServer.requestTimeout = 15_000;
httpServer.headersTimeout = 5_000;
httpServer.listen(config.port, "127.0.0.1");

async function shutdown(signal: string): Promise<void> {
  console.info(JSON.stringify({ event: "shutdown", signal }));
  httpServer.closeIdleConnections();
  await Promise.race([
    new Promise<void>((resolve, reject) =>
      httpServer.close((error) => error ? reject(error) : resolve())
    ),
    new Promise<never>((_, reject) =>
      setTimeout(() => reject(new Error("Shutdown deadline exceeded")), 10_000)
    ),
  ]);
  await mcp.close();
}

for (const signal of ["SIGINT", "SIGTERM"] as const) {
  process.once(signal, () => {
    shutdown(signal).then(() => process.exit(0)).catch((error) => {
      console.error(error);
      process.exit(1);
    });
  });
}