headplane/app/routes/ssh/page.tsx

import { Loader2, WifiOff } from "lucide-react";
import { useEffect, useState } from "react";
import { data, isRouteErrorResponse, type ShouldRevalidateFunction } from "react-router";

import Button from "~/components/button";
import Card from "~/components/card";
import Code from "~/components/code";
import StatusBanner from "~/components/status-banner";
import {
  agentsContext,
  appConfigContext,
  headscaleContext,
  requestApiContext,
} from "~/server/context";
import { findHeadscaleUserBySubject } from "~/server/web/headscale-identity";

import type { Route } from "./+types/page";
import { isSSHError, SSHErrorBoundary, sshErrors } from "./errors";
import Ghostty from "./ghostty.client";
import UserPrompt from "./user-prompt";
import type { HeadplaneSSH } from "./wasm.client";
import { loadHeadplaneWASM } from "./wasm.client";

const WASM_MODULE_URL = `${__PREFIX__}/hp_ssh.wasm`;
const WASM_HELPER_URL = `${__PREFIX__}/wasm_exec.js`;
const SSH_PREAUTH_KEY_TTL_MS = 10 * 60 * 1000;

export const shouldRevalidate: ShouldRevalidateFunction = () => {
  return false;
};

export async function loader({ request, params, context, url }: Route.LoaderArgs) {
  const agents = context.get(agentsContext);
  const config = context.get(appConfigContext);
  const headscale = context.get(headscaleContext);
  const getRequestApi = context.get(requestApiContext);
  const compatibilityWarning = getBrowserSSHCompatibilityWarning(headscale.version);

  const origin = url.origin;
  const assets = [WASM_HELPER_URL, WASM_MODULE_URL];
  const missing: string[] = [];

  for (const file of assets) {
    const res = await fetch(`${origin}${file}`, { method: "HEAD" });
    if (!res.ok) {
      missing.push(file);
    }
  }

  if (missing.length > 0) {
    throw data(sshErrors.wasm_missing, 405);
  }

  if (agents.state !== "enabled") {
    throw data(sshErrors.agent_required, 400);
  }

  const { principal, api } = await getRequestApi(request);
  if (principal.kind === "api_key") {
    throw data(sshErrors.oidc_required, 403);
  }

  const hostname = params.id;
  const username = url.searchParams.get("user") || undefined;

  const nodes = await api.nodes.list();
  const node = nodes.find((n) => n.givenName === hostname);
  if (!node) {
    throw data(sshErrors.node_not_found(hostname), 404);
  }

  if (!node.online) {
    return { hostname, username, offline: true, node: undefined, compatibilityWarning };
  }

  if (!username) {
    return {
      hostname,
      username: undefined,
      offline: false,
      node: undefined,
      compatibilityWarning,
    };
  }

  // The user must exist within Headscale to generate a pre-auth key
  const users = await api.users.list();
  const hsUser = principal.user.headscaleUserId
    ? users.find((u) => u.id === principal.user.headscaleUserId)
    : findHeadscaleUserBySubject(users, principal.user.subject, principal.profile.email);

  if (!hsUser) {
    throw data(sshErrors.user_not_linked, 404);
  }

  const preAuthKey = await api.preAuthKeys.create({
    user: hsUser.id,
    ephemeral: true,
    reusable: false,
    expiration: new Date(Date.now() + SSH_PREAUTH_KEY_TTL_MS),
    aclTags: null,
  });

  const controlURL = config.headscale.public_url ?? config.headscale.url;
  return {
    hostname,
    username,
    offline: false,
    node: {
      ipAddress: node.ipAddresses[0],
      controlURL,
      preAuthKey: preAuthKey.key,
      ephemeralHostname: generateHostname(username),
    },
    compatibilityWarning,
  };
}

function getBrowserSSHCompatibilityWarning(version: {
  unknown: boolean;
  major: number;
  minor: number;
  patch: number;
  raw: string;
}) {
  if (version.unknown) return null;
  if (version.major === 0 && version.minor === 29 && version.patch < 2) {
    return { version: version.raw };
  }
  return null;
}

function generateHostname(username: string) {
  const hex = crypto.randomUUID().replace(/-/g, "").slice(0, 8);
  return `ssh-${hex}-${username}`;
}

export const links: Route.LinksFunction = () => [
  {
    rel: "preload",
    href: WASM_MODULE_URL,
    as: "fetch",
    type: "application/wasm",
    crossOrigin: "anonymous",
  },
];

export default function Page({ loaderData }: Route.ComponentProps) {
  const { hostname, username, offline, node, compatibilityWarning } = loaderData;

  if (offline) {
    return (
      <>
        <BrowserSSHCompatibilityBanner warning={compatibilityWarning} />
        <div className="flex h-screen w-screen items-center justify-center bg-black">
          <Card className="w-screen" variant="flat">
            <div className="flex items-center justify-between gap-4">
              <Card.Title>Node Offline</Card.Title>
              <WifiOff className="mb-2 h-6 w-6 text-red-500" />
            </div>
            <Card.Text>
              <Code>{hostname}</Code> is not currently connected to the Tailnet.
            </Card.Text>
            <Button className="mt-8 w-full" onClick={() => window.location.reload()}>
              Retry Connection
            </Button>
          </Card>
        </div>
      </>
    );
  }

  if (!username || !node) {
    return (
      <>
        <BrowserSSHCompatibilityBanner warning={compatibilityWarning} />
        <UserPrompt hostname={hostname} />
      </>
    );
  }

  return (
    <>
      <BrowserSSHCompatibilityBanner warning={compatibilityWarning} />
      <SSHConsole hostname={hostname} username={username} node={node} />
    </>
  );
}

function BrowserSSHCompatibilityBanner({
  warning,
}: {
  warning: { version: string } | null | undefined;
}) {
  if (!warning) return null;

  return (
    <div className="fixed inset-x-4 top-4 z-[60] mx-auto max-w-2xl">
      <StatusBanner
        variant="warning"
        title={`Browser SSH is broken on Headscale ${warning.version}`}
      >
        Headscale 0.29 beta releases through 0.29.1 reject Tailscale's browser/WASM{" "}
        <Code>/ts2021</Code> WebSocket request with <Code>405 Method Not Allowed</Code>. Upgrade
        Headscale to 0.29.2 or newer, or use Headscale 0.28.x.
      </StatusBanner>
    </div>
  );
}

function SSHConsole({
  hostname,
  username,
  node,
}: {
  hostname: string;
  username: string;
  node: { ipAddress: string; controlURL: string; preAuthKey: string; ephemeralHostname: string };
}) {
  const [ssh, setSsh] = useState<HeadplaneSSH | null>(null);
  const [connected, setConnected] = useState(false);
  const [status, setStatus] = useState("Starting tunnel…");

  useEffect(() => {
    let cancelled = false;

    console.log("[ssh] Loading WASM factory");
    loadHeadplaneWASM().then((create) => {
      console.log("[ssh] Factory loaded, creating IPN", create);

      if (cancelled) {
        return;
      }

      setStatus("Joining Tailnet…");
      const instance = create({
        controlURL: node.controlURL,
        preAuthKey: node.preAuthKey,
        hostname: node.ephemeralHostname,
        onReady: () => {
          console.log("[ssh] IPN ready (Running)");
          if (!cancelled) {
            setStatus(`Connecting to ${hostname}…`);
            setSsh(instance);
          }
        },
        onError: (msg) => {
          console.error("[ssh] IPN error:", msg);
          if (!cancelled) {
            setStatus(`Failed to join Tailnet: ${msg}`);
          }
        },
      });

      console.log("[ssh] IPN instance created", instance);
    });

    return () => {
      cancelled = true;
    };
  }, [node]);

  return (
    <div className="fixed inset-0 flex flex-col bg-black">
      {!connected && (
        <div className="absolute inset-0 z-50 flex items-center justify-center">
          <div className="flex flex-col items-center gap-3">
            <Loader2 className="size-8 animate-spin text-mist-200" />
            <p className="text-sm text-mist-400">{status}</p>
          </div>
        </div>
      )}

      {ssh && (
        <Ghostty
          ssh={ssh}
          username={username}
          ipAddress={node.ipAddress}
          onConnected={() => setConnected(true)}
        />
      )}
    </div>
  );
}

export function ErrorBoundary({ error }: Route.ErrorBoundaryProps) {
  const routeError = isRouteErrorResponse(error) ? error.data : null;
  if (routeError == null || !isSSHError(routeError)) {
    // Pass through further down the tree to the global error boundary
    throw error;
  }

  return (
    <div className="flex h-screen w-screen items-center justify-center">
      <SSHErrorBoundary
        title={routeError.title}
        message={routeError.message}
        anchor={routeError.anchor}
      />
    </div>
  );
}