dispatcher.ts

import { mediaTypes } from "@hapi/accept";
import createDebugger from "debug";

import fetch, { Headers } from "node-fetch";

import type { ContextRegistry } from "./context-registry.js";
import type {
  HttpMethods,
  CounterfactResponseObject,
  Registry,
} from "./registry.js";
import { createResponseBuilder } from "./response-builder.js";
import { validateRequest } from "./request-validator.js";
import { validateResponse } from "./response-validator.js";
import { Tools } from "./tools.js";
import type {
  OpenApiOperation,
  OpenApiParameters,
} from "../counterfact-types/index.js";
import type { Config } from "./config.js";

const debug = createDebugger("counterfact:server:dispatcher");

function parseCookies(cookieHeader: string): Record<string, string> {
  const cookies: Record<string, string> = {};

  for (const part of cookieHeader.split(";")) {
    const eqIndex = part.indexOf("=");

    if (eqIndex === -1) {
      continue;
    }

    const key = part.slice(0, eqIndex).trim();
    const value = part.slice(eqIndex + 1).trim();

    if (key && !(key in cookies)) {
      try {
        cookies[key] = decodeURIComponent(value);
      } catch (error) {
        debug("could not decode cookie value for key %s: %o", key, error);
        cookies[key] = value;
      }
    }
  }

  return cookies;
}

interface ParameterTypes {
  body: Map<string, string>;
  cookie: Map<string, string>;
  formData: Map<string, string>;
  header: Map<string, string>;
  path: Map<string, string>;
  query: Map<string, string>;
}

export interface OpenApiDocument {
  basePath?: string;
  paths: {
    [key: string]: {
      [key in Lowercase<HttpMethods>]?: OpenApiOperation;
    };
  };
  produces?: string[];
}

export type DispatcherRequest = {
  auth?: {
    password?: string;
    username?: string;
  };
  body: unknown;
  headers: {
    [key: string]: string;
  };
  method: HttpMethods;
  path: string;
  query: {
    [key: string]: string;
  };
  rawBody?: string;
  req: {
    path?: string;
  };
};

export class Dispatcher {
  public registry: Registry;

  public contextRegistry: ContextRegistry;

  public openApiDocument: OpenApiDocument | undefined;

  public fetch: typeof fetch;

  public config?: Config; // Add config property

  public constructor(
    registry: Registry,
    contextRegistry: ContextRegistry,
    openApiDocument?: OpenApiDocument,
    config?: Config,
  ) {
    this.registry = registry;
    this.contextRegistry = contextRegistry;
    this.openApiDocument = openApiDocument;
    this.fetch = fetch;
    this.config = config;
  }

  private parameterTypes(
    parameters: OpenApiParameters[] | undefined,
  ): ParameterTypes {
    const types: ParameterTypes = {
      body: new Map(),
      cookie: new Map(),
      formData: new Map(),
      header: new Map(),
      path: new Map(),
      query: new Map(),
    };

    if (!parameters) {
      return types;
    }

    for (const parameter of parameters) {
      const type = parameter?.type;

      if (type !== undefined) {
        types[parameter.in].set(
          parameter.name,
          type === "integer" ? "number" : type,
        );
      }
    }

    return types;
  }

  private findOperation(
    path: string,
    method: HttpMethods,
  ): OpenApiOperation | undefined {
    if (this.openApiDocument) {
      for (const key in this.openApiDocument.paths) {
        if (key.toLowerCase() === path.toLowerCase()) {
          return this.openApiDocument.paths[key]?.[
            method.toLowerCase() as Lowercase<HttpMethods>
          ];
        }
      }
    }

    return undefined;
  }

  public operationForPathAndMethod(
    path: string,
    method: HttpMethods,
  ): OpenApiOperation | undefined {
    const operation = this.findOperation(path, method);

    if (operation === undefined) {
      return undefined;
    }

    if (this.openApiDocument?.produces) {
      return {
        produces: this.openApiDocument.produces,
        ...operation,
      };
    }

    return operation;
  }

  private normalizeResponse(
    response: CounterfactResponseObject,
    acceptHeader: string,
  ): CounterfactResponseObject {
    if (response.content !== undefined) {
      const content = this.selectContent(acceptHeader, response.content);

      if (content === undefined) {
        return {
          body: `Not Acceptable: could not produce a response matching any of the following content types: ${acceptHeader}`,
          contentType: "text/plain",
          status: 406,
        };
      }

      const normalizedResponse = {
        ...response,

        body: content.body as Uint8Array | string | undefined,
        contentType: content.type,
      };

      delete normalizedResponse.content;

      return normalizedResponse;
    }

    return {
      ...response,

      contentType:
        response.headers?.["content-type"]?.toString() ??
        response.contentType ??
        "unknown/unknown",
    };
  }

  public selectContent(
    acceptHeader: string,
    content: { body: unknown; type: string }[],
  ) {
    const preferredMediaTypes = mediaTypes(acceptHeader);

    for (const mediaType of preferredMediaTypes) {
      const contentItem = content.find((item) =>
        this.isMediaType(item.type, mediaType),
      );

      if (contentItem) {
        return contentItem;
      }
    }

    return undefined;
  }

  private isMediaType(type: string, pattern: string) {
    if (pattern === "*/*") {
      return true;
    }

    const [baseType, subType] = type.split("/");

    const [patternType, patternSubType] = pattern.split("/");

    if (baseType === patternType) {
      return subType === patternSubType || patternSubType === "*";
    }

    if (subType === patternSubType) {
      return baseType === patternType || patternType === "*";
    }

    return false;
  }

  public async request({
    auth,
    body,
    headers = {},
    method,
    path,
    query,
    rawBody,
    req,
  }: DispatcherRequest): Promise<CounterfactResponseObject> {
    debug(`request: ${method} ${path}`);
    debug(`headers: ${JSON.stringify(headers)}`);
    debug(`body: ${JSON.stringify(body)}`);

    // If the incoming path includes the base path, remove it
    if (
      this.openApiDocument?.basePath !== undefined &&
      path.toLowerCase().startsWith(this.openApiDocument.basePath.toLowerCase())
    ) {
      path = path.slice(this.openApiDocument.basePath.length);
    }

    const { matchedPath } = this.registry.handler(path, method);

    if (
      !this.registry.exists(method, path) &&
      this.registry.pathExistsWithAnyMethod(path, method)
    ) {
      return {
        body: `The ${method} method is not allowed for ${path}\n`,
        contentType: "text/plain",
        headers: { allow: this.registry.allowedMethods(path) },
        status: 405,
      };
    }

    const operation = this.operationForPathAndMethod(matchedPath, method);

    if (this.config?.validateRequests !== false) {
      const validation = validateRequest(operation, { body, headers, query });

      if (!validation.valid) {
        return {
          body: `Request validation failed:\n${validation.errors.join("\n")}`,
          contentType: "text/plain",
          status: 400,
        };
      }
    }

    const continuousDistribution = (min: number, max: number) => {
      return min + Math.random() * (max - min);
    };

    const response = await this.registry.endpoint(
      method,
      path,
      this.parameterTypes(operation?.parameters),
    )({
      auth,
      body,
      context: this.contextRegistry.find(matchedPath),

      async delay(milliseconds = 0, maxMilliseconds = 0) {
        const delayInMs =
          maxMilliseconds - milliseconds <= 0
            ? milliseconds
            : continuousDistribution(milliseconds, maxMilliseconds);

        return new Promise((resolve) => setTimeout(resolve, delayInMs));
      },

      cookie: parseCookies(headers.cookie ?? headers.Cookie ?? ""),

      headers,

      proxy: async (url: string) => {
        delete headers.host;

        const fetchResponse = await this.fetch(`${url}${req.path ?? ""}`, {
          body: body === undefined ? undefined : rawBody,
          headers: new Headers(headers),
          method,
        });

        const responseHeaders = Object.fromEntries(
          fetchResponse.headers.entries(),
        );

        return {
          body: await fetchResponse.text(),
          contentType: responseHeaders["content-type"] ?? "unknown/unknown",
          headers: responseHeaders,
          status: fetchResponse.status,
        };
      },

      query,

      // @ts-expect-error - Might be pushing the limits of what TypeScript can do here
      response: createResponseBuilder(
        operation ?? { responses: {} },
        this.config,
      ), // Pass config

      tools: new Tools({ headers }),
    });

    if (response === undefined) {
      return {
        body: `The ${method} function did not return anything. Did you forget a return statement?`,
        status: 500,
      };
    }

    const normalizedResponse = this.normalizeResponse(
      response,
      headers.accept ?? "*/*",
    );

    if (
      normalizedResponse.body !== undefined &&
      !mediaTypes(headers.accept ?? "*/*").some((type) =>
        this.isMediaType(normalizedResponse.contentType ?? "", type),
      )
    ) {
      return {
        body: JSON.stringify(mediaTypes(headers.accept ?? "*/*")),
        status: 406,
      };
    }

    if (this.config?.validateResponses !== false) {
      const validation = validateResponse(operation, normalizedResponse);

      if (!validation.valid) {
        return {
          ...normalizedResponse,
          headers: {
            ...normalizedResponse.headers,
            "response-type-error": validation.errors,
          },
        };
      }
    }

    return normalizedResponse;
  }
}