import { Injectable, NestMiddleware } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import {
  localhostAllowedHostnames,
  validateHostHeader,
  validateOriginHeader,
} from '@modelcontextprotocol/server';
import type { NextFunction, Request, Response } from 'express';

/**
 * Protection against DNS rebinding on /mcp (what the MCP spec asks of HTTP
 * servers): the Host header must be one of MCP_ALLOWED_HOSTS (default: only
 * localhost), and a browser's Origin, when sent, too. AI clients calling from
 * a server send no Origin, so they only need the right domain.
 */
@Injectable()
export class McpHostsMiddleware implements NestMiddleware {
  private readonly hosts: string[];

  constructor(config: ConfigService) {
    const configured = (config.get<string>('MCP_ALLOWED_HOSTS') ?? '')
      .split(',')
      .map((host) => host.trim().toLowerCase())
      .filter(Boolean);
    this.hosts = configured.length ? configured : localhostAllowedHostnames();
  }

  use(req: Request, res: Response, next: NextFunction): void {
    const error = this.rejection(req);
    if (!error) return next();
    res.status(403).json({
      jsonrpc: '2.0',
      id: null,
      error: { code: -32000, message: error },
    });
  }

  /** Why the request is refused, or null when its Host (and Origin, if any) are allowed. */
  private rejection(req: Request): string | null {
    const host = validateHostHeader(req.headers.host, this.hosts);
    if ('message' in host) return host.message;
    if (!req.headers.origin) return null;
    const origin = validateOriginHeader(req.headers.origin, this.hosts);
    return 'message' in origin ? origin.message : null;
  }
}