photoncloud-monorepo/nix/test-cluster/vm-guest-image.nix

{ modulesPath, lib, pkgs, ... }:

let
  photonVmDemoApi = pkgs.writeText "photon-vm-demo-api.py" ''
    import json
    import os
    import socket
    import sqlite3
    from http import HTTPStatus
    from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

    DATA_MOUNT = "/mnt/photon-vm-data"
    DB_PATH = os.path.join(DATA_MOUNT, "demo.sqlite3")
    ROOT_BOOT_COUNT_PATH = "/var/lib/photon-vm-smoke/boot-count"
    DATA_BOOT_COUNT_PATH = os.path.join(DATA_MOUNT, "boot-count")
    CONSOLE_PATH = "/dev/ttyS0"
    LISTEN_HOST = "0.0.0.0"
    LISTEN_PORT = 8080


    def log_console(message: str) -> None:
        try:
            with open(CONSOLE_PATH, "a", encoding="utf-8") as console:
                console.write(message + "\n")
        except OSError:
            pass


    def read_int(path: str) -> int:
        try:
            with open(path, "r", encoding="utf-8") as handle:
                return int(handle.read().strip() or "0")
        except (FileNotFoundError, ValueError, OSError):
            return 0


    def init_db() -> None:
        os.makedirs(DATA_MOUNT, exist_ok=True)
        conn = sqlite3.connect(DB_PATH)
        try:
            conn.execute(
                "CREATE TABLE IF NOT EXISTS counters (name TEXT PRIMARY KEY, value INTEGER NOT NULL)"
            )
            conn.execute(
                "INSERT INTO counters (name, value) VALUES ('visits', 0) "
                "ON CONFLICT(name) DO NOTHING"
            )
            conn.commit()
        finally:
            conn.close()


    def current_state(increment: bool = False) -> dict:
        conn = sqlite3.connect(DB_PATH, timeout=30)
        try:
            conn.execute(
                "CREATE TABLE IF NOT EXISTS counters (name TEXT PRIMARY KEY, value INTEGER NOT NULL)"
            )
            conn.execute(
                "INSERT INTO counters (name, value) VALUES ('visits', 0) "
                "ON CONFLICT(name) DO NOTHING"
            )
            if increment:
                conn.execute(
                    "UPDATE counters SET value = value + 1 WHERE name = 'visits'"
                )
            visits = conn.execute(
                "SELECT value FROM counters WHERE name = 'visits'"
            ).fetchone()[0]
            conn.commit()
        finally:
            conn.close()
        return {
            "status": "ok",
            "hostname": socket.gethostname(),
            "listen_port": LISTEN_PORT,
            "db_path": DB_PATH,
            "visits": visits,
            "root_boot_count": read_int(ROOT_BOOT_COUNT_PATH),
            "data_boot_count": read_int(DATA_BOOT_COUNT_PATH),
        }


    class Handler(BaseHTTPRequestHandler):
        server_version = "PhotonVMDemo/1.0"

        def log_message(self, format: str, *args) -> None:
            return

        def _send_json(self, payload: dict, status: int = HTTPStatus.OK) -> None:
            body = json.dumps(payload, sort_keys=True).encode("utf-8")
            self.send_response(status)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self) -> None:
            if self.path == "/health":
                self._send_json({"status": "ok"})
                return
            if self.path == "/state":
                self._send_json(current_state())
                return
            self._send_json({"error": "not_found"}, HTTPStatus.NOT_FOUND)

        def do_POST(self) -> None:
            if self.path == "/visit":
                payload = current_state(increment=True)
                log_console("PHOTON_VM_DEMO_VISIT visits=%s" % payload["visits"])
                self._send_json(payload)
                return
            self._send_json({"error": "not_found"}, HTTPStatus.NOT_FOUND)


    def main() -> None:
        init_db()
        server = ThreadingHTTPServer((LISTEN_HOST, LISTEN_PORT), Handler)
        log_console(
            "PHOTON_VM_DEMO_WEB_READY count=%s port=%s db=%s"
            % (read_int(ROOT_BOOT_COUNT_PATH), LISTEN_PORT, DB_PATH)
        )
        server.serve_forever()


    if __name__ == "__main__":
        main()
  '';
in {
  imports = [
    (modulesPath + "/virtualisation/disk-image.nix")
    (modulesPath + "/profiles/qemu-guest.nix")
  ];

  image = {
    baseName = "photon-vm-smoke";
    format = "qcow2";
    efiSupport = false;
  };

  virtualisation.diskSize = 4096;

  boot.kernelParams = [ "console=ttyS0" "console=tty0" ];

  networking.hostName = "photon-vm-smoke";
  networking.useDHCP = lib.mkDefault true;
  networking.firewall.enable = false;

  services.getty.autologinUser = "root";
  users.mutableUsers = false;
  users.users.root.hashedPassword = "$6$photoncloud$aUJCEE5wm/b5O.9KIKGm84qUWdWXwnebsFEiMBF7u9Y7AOWodaMrjbbKGMOf0X59VJyJeMRsgbT7VWeqMHpUe.";

  documentation.enable = false;
  services.openssh.enable = false;
  environment.systemPackages = [ pkgs.e2fsprogs pkgs.util-linux ];

  systemd.services.photon-vm-smoke = {
    description = "PhotonCloud VM smoke marker";
    wantedBy = [ "multi-user.target" ];
    wants = [ "systemd-udev-settle.service" ];
    after = [ "local-fs.target" "systemd-udev-settle.service" ];
    path = with pkgs; [
      bash
      coreutils
      e2fsprogs
      gawk
      gnugrep
      gnused
      util-linux
    ];
    serviceConfig = {
      Type = "simple";
      Restart = "always";
      RestartSec = "1";
    };
    script = ''
      mkdir -p /var/lib/photon-vm-smoke
      count_file=/var/lib/photon-vm-smoke/boot-count
      if [ -f "$count_file" ]; then
        count=$(( $(cat "$count_file") + 1 ))
      else
        count=1
      fi
      echo "$count" > "$count_file"
      echo "PHOTON_VM_SMOKE_READY count=$count" >/dev/ttyS0

      root_source="$(lsblk -nrpo NAME,MOUNTPOINT | awk '$2 == "/" { print $1; exit }' 2>/dev/null || true)"
      root_disk=""
      if [ -n "$root_source" ] && [ -b "$root_source" ]; then
        root_disk="$(lsblk -ndo PKNAME "$root_source" 2>/dev/null || true)"
        if [ -z "$root_disk" ]; then
          root_disk="$(basename "$root_source")"
        else
          root_disk="/dev/$root_disk"
        fi
      fi
      echo "PHOTON_VM_SMOKE_DATA_ROOT count=$count source=''${root_source:-none} root=''${root_disk:-unknown}" >/dev/ttyS0

      data_disk=""
      if [ -b /dev/disk/by-label/photon-vm-data ]; then
        data_disk="$(readlink -f /dev/disk/by-label/photon-vm-data)"
      fi

      pick_data_disk() {
        while IFS= read -r disk; do
          [ -n "$disk" ] || continue
          if [ -n "$root_source" ] && [ "$disk" = "$root_source" ]; then
            continue
          fi
          if [ -n "$root_disk" ] && [ "$disk" = "$root_disk" ]; then
            continue
          fi
          if lsblk -nrpo MOUNTPOINT "$disk" 2>/dev/null | grep -qx '/'; then
            continue
          fi
          printf '%s\n' "$disk"
          return 0
        done < <(lsblk -dnpr -o NAME,TYPE,RO | awk '$2 == "disk" && $3 == "0" { print $1 }')
        return 1
      }

      deadline=$((SECONDS + 60))
      scan_attempt=0
      while [ -z "$data_disk" ] && [ "$SECONDS" -lt "$deadline" ]; do
        scan_attempt=$((scan_attempt + 1))
        data_disk="$(pick_data_disk || true)"
        echo "PHOTON_VM_SMOKE_DATA_SCAN count=$count attempt=$scan_attempt data=''${data_disk:-none}" >/dev/ttyS0
        [ -n "$data_disk" ] && break
        udevadm settle >/dev/null 2>&1 || true
        sleep 1
      done

      if [ -z "$data_disk" ]; then
        echo "PHOTON_VM_SMOKE_DATA_MISSING count=$count" >/dev/ttyS0
        lsblk -dn -o NAME,TYPE,SIZE >/dev/ttyS0 2>&1 || true
        exit 1
      fi

      echo "PHOTON_VM_SMOKE_DATA_PROBE count=$count root=''${root_disk:-unknown} data=$(basename "$data_disk")" >/dev/ttyS0
      mkdir -p /mnt/photon-vm-data
      if ! blkid "$data_disk" >/dev/null 2>&1; then
        mkfs_output="$(mkfs.ext4 -L photon-vm-data -F "$data_disk" 2>&1)" || {
          mkfs_output="$(printf '%s' "$mkfs_output" | tr '\r\n' ' ' | sed 's/  */ /g')"
          echo "PHOTON_VM_SMOKE_DATA_ERROR count=$count step=mkfs device=$(basename "$data_disk") detail=''${mkfs_output}" >/dev/ttyS0
          lsblk -dn -o NAME,TYPE,RO,SIZE >/dev/ttyS0 2>&1 || true
          blockdev --getsize64 "$data_disk" >/dev/ttyS0 2>&1 || true
          exit 1
        }
      fi
      if ! mountpoint -q /mnt/photon-vm-data; then
        if ! mount "$data_disk" /mnt/photon-vm-data; then
          echo "PHOTON_VM_SMOKE_DATA_ERROR count=$count step=mount device=$(basename "$data_disk")" >/dev/ttyS0
          lsblk -f >/dev/ttyS0 2>&1 || true
          exit 1
        fi
      fi
      data_count_file=/mnt/photon-vm-data/boot-count
      if [ -f "$data_count_file" ]; then
        data_count=$(( $(cat "$data_count_file") + 1 ))
      else
        data_count=1
      fi
      echo "$data_count" > "$data_count_file"
      sync
      echo "PHOTON_VM_SMOKE_DATA_READY count=$data_count device=$(basename "$data_disk")" >/dev/ttyS0

      while true; do
        echo "PHOTON_VM_SMOKE_HEARTBEAT count=$count ts=$(date +%s)" >/dev/ttyS0
        sleep 2
      done
    '';
  };

  systemd.services.photon-vm-demo-api = {
    description = "PhotonCloud VM demo web app";
    wantedBy = [ "multi-user.target" ];
    wants = [ "network-online.target" "photon-vm-smoke.service" ];
    after = [ "network-online.target" "photon-vm-smoke.service" ];
    path = with pkgs; [
      bash
      coreutils
      python3
      util-linux
    ];
    serviceConfig = {
      Type = "simple";
      Restart = "always";
      RestartSec = "1";
    };
    script = ''
      deadline=$((SECONDS + 60))
      while ! mountpoint -q /mnt/photon-vm-data; do
        if [ "$SECONDS" -ge "$deadline" ]; then
          echo "PHOTON_VM_DEMO_WEB_ERROR step=mount-timeout" >/dev/ttyS0
          exit 1
        fi
        sleep 1
      done

      exec python3 ${photonVmDemoApi}
    '';
  };

  system.stateVersion = "24.05";
}