kuber/server/registry.ts

import { createHash } from "node:crypto";
import {
  assertSha256Digest,
  type Sha256Digest,
} from "../shared/build-protocol";

const ACCEPT = [
  "application/vnd.oci.image.index.v1+json",
  "application/vnd.oci.image.manifest.v1+json",
  "application/vnd.docker.distribution.manifest.list.v2+json",
  "application/vnd.docker.distribution.manifest.v2+json",
].join(", ");

export type RegistryCredentials = { username: string; password: string };
export type RegistryFetch = (
  input: string | URL | Request,
  init?: RequestInit,
) => Promise<Response>;
export type RegistryResolveOptions = {
  fetch?: RegistryFetch;
  credentials?: RegistryCredentials;
  insecure?: boolean;
  origin?: string;
  cacheTtlMs?: number;
  cacheMaxEntries?: number;
  clock?: () => number;
};

const DEFAULT_DIGEST_CACHE_TTL_MS = 30_000;
const DEFAULT_DIGEST_CACHE_MAX_ENTRIES = 256;
const digestCache = new Map<
  string,
  { digest: Sha256Digest; expiresAt: number }
>();

export type ParsedImageReference = {
  registry: string;
  repository: string;
  reference: string;
  digest?: Sha256Digest;
};

function validateRepository(repository: string, image: string): void {
  if (
    !repository ||
    repository.split("/").some((part) => !part || part === "." || part === "..")
  )
    throw new Error(`Invalid image reference: ${image}`);
}

export function parseImageReference(image: string): ParsedImageReference {
  const slash = image.indexOf("/");
  if (slash <= 0)
    throw new Error("Image reference must include a registry host");
  const registry = image.slice(0, slash);
  let repositoryAndReference = image.slice(slash + 1);
  if (
    !repositoryAndReference ||
    !registry ||
    /[/?#@]/.test(registry) ||
    /\s/.test(image)
  )
    throw new Error(`Invalid image reference: ${image}`);

  const at = repositoryAndReference.lastIndexOf("@");
  if (at !== -1) {
    const value = repositoryAndReference.slice(at + 1);
    assertSha256Digest(value);
    repositoryAndReference = repositoryAndReference.slice(0, at);
    validateRepository(repositoryAndReference, image);
    return {
      registry,
      repository: repositoryAndReference,
      reference: value,
      digest: value,
    };
  }
  const lastSlash = repositoryAndReference.lastIndexOf("/");
  const colon = repositoryAndReference.lastIndexOf(":");
  const reference =
    colon > lastSlash ? repositoryAndReference.slice(colon + 1) : "latest";
  const repository =
    colon > lastSlash
      ? repositoryAndReference.slice(0, colon)
      : repositoryAndReference;
  validateRepository(repository, image);
  if (!reference) throw new Error(`Invalid image reference: ${image}`);
  return { registry, repository, reference };
}

function bearerParameters(
  challenge: string,
): Record<string, string> | undefined {
  const match = /^Bearer\s+(.+)$/i.exec(challenge.trim());
  if (!match?.[1]) return;
  const values: Record<string, string> = {};
  const expression = /([a-z][a-z0-9_-]*)=(?:"((?:\\.|[^"])*)"|([^,\s]+))/gi;
  for (const item of match[1].matchAll(expression))
    values[item[1]!.toLowerCase()] = (item[2] ?? item[3] ?? "").replace(
      /\\"/g,
      '"',
    );
  return values.realm ? values : undefined;
}

async function responseError(response: Response): Promise<Error> {
  const detail = (await response.text()).slice(0, 512).trim();
  return new Error(
    `Registry request failed (${response.status})${detail ? `: ${detail}` : ""}`,
  );
}

export async function resolveRegistryDigest(
  image: string,
  options: RegistryResolveOptions = {},
): Promise<Sha256Digest> {
  const parsed = parseImageReference(image);
  if (parsed.digest) return parsed.digest;
  const now = options.clock ?? Date.now;
  const cacheTtlMs = Number.isFinite(options.cacheTtlMs)
    ? Math.max(0, options.cacheTtlMs!)
    : DEFAULT_DIGEST_CACHE_TTL_MS;
  const cacheMaxEntries = Number.isFinite(options.cacheMaxEntries)
    ? Math.max(0, Math.floor(options.cacheMaxEntries!))
    : DEFAULT_DIGEST_CACHE_MAX_ENTRIES;
  const fetcher: RegistryFetch = options.fetch ?? globalThis.fetch;
  const scheme = options.insecure ? "http" : "https";
  const repository = parsed.repository
    .split("/")
    .map(encodeURIComponent)
    .join("/");
  const origin = (options.origin ?? `${scheme}://${parsed.registry}`).replace(
    /\/+$/,
    "",
  );
  const credentialsKey = options.credentials
    ? createHash("sha256")
        .update(
          `${options.credentials.username}\0${options.credentials.password}`,
        )
        .digest("hex")
    : "anonymous";
  const cacheKey = `${origin}\0${parsed.repository}\0${parsed.reference}\0${credentialsKey}`;
  const cached = digestCache.get(cacheKey);
  if (cached) {
    if (cached.expiresAt > now()) return cached.digest;
    digestCache.delete(cacheKey);
  }
  const manifestUrl = `${origin}/v2/${repository}/manifests/${encodeURIComponent(parsed.reference)}`;
  const headers = new Headers({ accept: ACCEPT });
  if (options.credentials) {
    headers.set(
      "authorization",
      `Basic ${Buffer.from(`${options.credentials.username}:${options.credentials.password}`).toString("base64")}`,
    );
  }

  let response = await fetcher(manifestUrl, { headers });
  if (response.status === 401) {
    const challenge = bearerParameters(
      response.headers.get("www-authenticate") ?? "",
    );
    if (!challenge) throw await responseError(response);
    const tokenUrl = new URL(challenge.realm!);
    if (tokenUrl.protocol !== "https:" && !options.insecure)
      throw new Error("Registry bearer token realm must use HTTPS");
    if (tokenUrl.protocol !== "https:" && tokenUrl.protocol !== "http:")
      throw new Error("Registry bearer token realm must use HTTP or HTTPS");
    if (challenge.service)
      tokenUrl.searchParams.set("service", challenge.service);
    tokenUrl.searchParams.set(
      "scope",
      challenge.scope ?? `repository:${parsed.repository}:pull`,
    );
    const tokenHeaders = new Headers();
    if (options.credentials)
      tokenHeaders.set(
        "authorization",
        `Basic ${Buffer.from(`${options.credentials.username}:${options.credentials.password}`).toString("base64")}`,
      );
    const tokenResponse = await fetcher(tokenUrl, { headers: tokenHeaders });
    if (!tokenResponse.ok) throw await responseError(tokenResponse);
    const payload = (await tokenResponse.json()) as {
      token?: unknown;
      access_token?: unknown;
    };
    const token = payload.token ?? payload.access_token;
    if (typeof token !== "string" || !token)
      throw new Error("Registry token response did not contain a token");
    headers.set("authorization", `Bearer ${token}`);
    response = await fetcher(manifestUrl, { headers });
  }
  if (!response.ok) throw await responseError(response);

  const body = new Uint8Array(await response.arrayBuffer());
  const advertised = response.headers
    .get("docker-content-digest")
    ?.trim()
    ?.toLowerCase();
  let digest: Sha256Digest;
  if (advertised !== undefined) {
    assertSha256Digest(advertised);
    digest = advertised;
  } else {
    digest = `sha256:${createHash("sha256").update(body).digest("hex")}`;
  }
  if (cacheTtlMs && cacheMaxEntries) {
    digestCache.delete(cacheKey);
    while (digestCache.size >= cacheMaxEntries)
      digestCache.delete(digestCache.keys().next().value!);
    digestCache.set(cacheKey, {
      digest,
      expiresAt: now() + cacheTtlMs,
    });
  }
  return digest;
}