corebuild/ts/classes.baseosimagebuilder.ts

import * as crypto from 'node:crypto';
import * as fs from 'node:fs';
import * as fsp from 'node:fs/promises';
import * as os from 'node:os';
import * as path from 'node:path';
import { spawn } from 'node:child_process';

import * as plugins from './plugins.js';
import type {
  IBaseOsImageArtifactResult,
  IBaseOsImageJob,
  IS3Descriptor,
  TBaseOsImageKind,
  TBaseOsImageSourcePreset,
} from './types.js';

interface IBalenaSourcePreset {
  preset: TBaseOsImageSourcePreset;
  architecture: IBaseOsImageJob['architecture'];
  deviceType: string;
}

const balenaSourcePresets: IBalenaSourcePreset[] = [
  {
    preset: 'balena-generic-amd64',
    architecture: 'amd64',
    deviceType: 'generic-amd64',
  },
  {
    preset: 'balena-generic-aarch64',
    architecture: 'arm64',
    deviceType: 'generic-aarch64',
  },
  {
    preset: 'balena-raspberrypi4-64',
    architecture: 'rpi',
    deviceType: 'raspberrypi4-64',
  },
];

export interface IBaseOsImageBuilderOptions {
  workdir: string;
  isoCreatorCommand: string;
}

export class BaseOsImageBuilder {
  constructor(private options: IBaseOsImageBuilderOptions) {}

  public async build(jobArg: IBaseOsImageJob): Promise<{
    artifact: IBaseOsImageArtifactResult;
    logs: string[];
  }> {
    const logs: string[] = [];
    const imageKind = this.getImageKind(jobArg);
    const jobDir = path.join(this.options.workdir, jobArg.id);
    const outputDir = path.join(jobDir, 'output');
    await fsp.rm(jobDir, { recursive: true, force: true });
    await fsp.mkdir(outputDir, { recursive: true });

    const filename = this.getArtifactFilename(jobArg, imageKind);
    const outputPath = path.join(outputDir, filename);
    const configPath = path.join(jobDir, 'isocreator.config.json');
    const isoCreatorConfig = imageKind === 'balena-raw'
      ? await this.createRawImageConfig(jobArg, outputDir, filename, logs)
      : this.createIsoCreatorConfig(jobArg, outputDir, filename);
    await fsp.writeFile(configPath, `${JSON.stringify(isoCreatorConfig, null, 2)}\n`);

    logs.push(`Starting isocreator for ${jobArg.architecture} ${imageKind}`);
    await this.runIsoCreator(configPath, logs);

    const stat = await fsp.stat(outputPath);
    const sha256 = await this.sha256File(outputPath);
    await this.uploadArtifact(
      jobArg,
      outputPath,
      stat.size,
      this.getContentType(filename, imageKind),
      logs,
    );

    await fsp.rm(jobDir, { recursive: true, force: true }).catch(() => undefined);
    return {
      artifact: {
        bucketName: jobArg.s3Descriptor.bucketName,
        key: jobArg.artifactKey,
        filename,
        contentType: this.getContentType(filename, imageKind),
        size: stat.size,
        sha256,
        createdAt: Date.now(),
      },
      logs,
    };
  }

  private createIsoCreatorConfig(jobArg: IBaseOsImageJob, outputDirArg: string, filenameArg: string) {
    const installScript = this.createBaseOsInstallScript();
    const serviceFile = this.createBaseOsServiceFile();
    const envFile = this.createBaseOsEnvFile(jobArg);
    return {
      version: '1.0',
      imageKind: 'iso',
      ...(jobArg.sourceImageUrl
        ? {
            source: {
              type: 'url',
              url: jobArg.sourceImageUrl,
            },
          }
        : {}),
      iso: {
        ubuntu_version: jobArg.ubuntuVersion || '24.04',
        architecture: jobArg.architecture === 'amd64' ? 'amd64' : 'arm64',
        flavor: 'server',
      },
      output: {
        filename: filenameArg,
        path: outputDirArg,
      },
      ...(jobArg.wifi?.ssid
        ? {
            network: {
              wifi: jobArg.wifi,
            },
          }
        : {}),
      cloud_init: {
        hostname: jobArg.hostname || `baseos-${jobArg.id.slice(0, 8)}`,
        users: jobArg.sshPublicKey
          ? [
              {
                name: 'baseos',
                ssh_authorized_keys: [jobArg.sshPublicKey],
                sudo: 'ALL=(ALL) NOPASSWD:ALL',
                shell: '/bin/bash',
                groups: ['sudo'],
              },
            ]
          : undefined,
        package_update: true,
        packages: ['curl', 'git', 'ca-certificates'],
        write_files: [
          {
            path: '/etc/baseos/baserunner.env',
            owner: 'root:root',
            permissions: '0600',
            content: envFile,
          },
          {
            path: '/etc/systemd/system/baseos-baserunner.service',
            owner: 'root:root',
            permissions: '0644',
            content: serviceFile,
          },
          {
            path: '/usr/local/bin/install-baseos.sh',
            owner: 'root:root',
            permissions: '0755',
            content: installScript,
          },
        ],
        runcmd: [
          'mkdir -p /var/lib/baseos /opt/baseos',
          '/usr/local/bin/install-baseos.sh',
          'systemctl daemon-reload',
          'systemctl enable baseos-baserunner.service',
          'systemctl start baseos-baserunner.service',
        ],
      },
    };
  }

  private async createRawImageConfig(
    jobArg: IBaseOsImageJob,
    outputDirArg: string,
    filenameArg: string,
    logsArg: string[],
  ) {
    const source = await this.resolveRawImageSource(jobArg, logsArg);

    return {
      version: '1.0',
      imageKind: 'raw-image',
      source,
      output: {
        filename: filenameArg,
        path: outputDirArg,
      },
      ...(jobArg.wifi?.ssid
        ? {
            network: {
              wifi: jobArg.wifi,
            },
          }
        : {}),
      raw_image: {
        sourceFormat: 'auto',
        bootPartition: '/dev/sda1',
        outputCompression: filenameArg.endsWith('.xz') ? 'xz' : 'none',
      },
      balena_os: {
        hostname: jobArg.hostname || `baseos-${jobArg.id.slice(0, 8)}`,
        sshPublicKeys: jobArg.sshPublicKey ? [jobArg.sshPublicKey] : [],
        configJson: {
          serveZone: {
            baseos: {
              buildId: jobArg.id,
              cloudlyUrl: jobArg.cloudlyUrl,
              provisioningToken: jobArg.provisioningToken,
            },
          },
        },
        baseOsEnv: {
          BASEOS_CLOUDLY_URL: jobArg.cloudlyUrl,
          BASEOS_JOIN_TOKEN: jobArg.provisioningToken,
          BASEOS_STATE_PATH: '/data/baseos/state.json',
          BASEOS_PRELOAD_TARGET_STATE_PATH: '/data/baseos/preload-target-state.json',
          BASEOS_HEARTBEAT_INTERVAL_MS: '60000',
          SERVEZONE_RUNTIME: 'baseos',
        },
      },
    };
  }

  private getImageKind(jobArg: IBaseOsImageJob): TBaseOsImageKind {
    if (jobArg.imageKind) {
      return jobArg.imageKind;
    }
    if (jobArg.architecture === 'rpi') {
      return 'balena-raw';
    }
    if (jobArg.sourceImagePreset || jobArg.balenaOsVersion) {
      return 'balena-raw';
    }
    if (jobArg.sourceImageUrl && this.isRawImageUrl(jobArg.sourceImageUrl)) {
      return 'balena-raw';
    }
    return 'ubuntu-iso';
  }

  private async resolveRawImageSource(jobArg: IBaseOsImageJob, logsArg: string[]) {
    if (jobArg.sourceImageUrl) {
      const filename = this.getSourceFilenameFromUrl(jobArg.sourceImageUrl);
      return {
        type: 'url' as const,
        url: jobArg.sourceImageUrl,
        ...(filename ? { filename } : {}),
      };
    }

    const preset = this.getBalenaSourcePreset(jobArg);
    const version = await this.resolveBalenaOsVersion(preset, jobArg.balenaOsVersion, logsArg);
    const downloadUrl = new URL('https://api.balena-cloud.com/download');
    downloadUrl.searchParams.set('deviceType', preset.deviceType);
    downloadUrl.searchParams.set('version', version);
    downloadUrl.searchParams.set('fileType', '.zip');
    return {
      type: 'url' as const,
      url: downloadUrl.toString(),
      filename: `balenaos-${preset.deviceType}-${version}.img.zip`,
    };
  }

  private getBalenaSourcePreset(jobArg: IBaseOsImageJob) {
    const preset = jobArg.sourceImagePreset
      ? balenaSourcePresets.find((presetArg) => presetArg.preset === jobArg.sourceImagePreset)
      : balenaSourcePresets.find((presetArg) => presetArg.architecture === jobArg.architecture);
    if (!preset) {
      throw new Error(`No balenaOS source preset is available for ${jobArg.architecture}`);
    }
    if (preset.architecture !== jobArg.architecture) {
      throw new Error(`${preset.preset} is only valid for ${preset.architecture} BaseOS images`);
    }
    return preset;
  }

  private async resolveBalenaOsVersion(
    presetArg: IBalenaSourcePreset,
    versionArg: string | undefined,
    logsArg: string[],
  ) {
    const requestedVersion = versionArg?.trim() || 'latest';
    if (requestedVersion !== 'latest') {
      return requestedVersion;
    }

    const releaseUrl = new URL('https://api.balena-cloud.com/v7/release');
    releaseUrl.searchParams.set('$select', 'raw_version');
    releaseUrl.searchParams.set(
      '$filter',
      `(is_final eq true) and (is_invalidated eq false) and (status eq 'success') and (semver_major gt 0) and (belongs_to__application/any(bta:(bta/is_host eq true) and (bta/is_for__device_type/any(dt:dt/slug eq '${presetArg.deviceType}'))))`,
    );
    releaseUrl.searchParams.set('$orderby', 'semver_major desc,semver_minor desc,semver_patch desc,revision desc');
    releaseUrl.searchParams.set('$top', '1');
    const response = await fetch(releaseUrl);
    if (!response.ok) {
      throw new Error(`Failed to resolve latest balenaOS version for ${presetArg.deviceType}: HTTP ${response.status}`);
    }
    const responseBody = await response.json() as { d?: Array<{ raw_version?: string }> };
    const latestVersion = responseBody.d?.[0]?.raw_version;
    if (!latestVersion) {
      throw new Error(`No balenaOS version found for ${presetArg.deviceType}`);
    }
    logsArg.push(`Resolved balenaOS ${presetArg.deviceType} latest version ${latestVersion}`);
    return latestVersion;
  }

  private isRawImageUrl(sourceImageUrlArg: string) {
    if (/\.(img|img\.xz|zip)(\?|$)/i.test(sourceImageUrlArg)) {
      return true;
    }
    try {
      const sourceUrl = new URL(sourceImageUrlArg);
      return sourceUrl.searchParams.get('fileType') === '.zip';
    } catch {
      return false;
    }
  }

  private getSourceFilenameFromUrl(sourceImageUrlArg: string) {
    try {
      const sourceUrl = new URL(sourceImageUrlArg);
      const lowerPath = sourceUrl.pathname.toLowerCase();
      if (lowerPath.endsWith('.img') || lowerPath.endsWith('.img.xz') || lowerPath.endsWith('.zip')) {
        return undefined;
      }
      if (sourceUrl.searchParams.get('fileType') === '.zip') {
        return 'source.img.zip';
      }
    } catch {
      return undefined;
    }
    return undefined;
  }

  private getArtifactFilename(jobArg: IBaseOsImageJob, imageKindArg: TBaseOsImageKind) {
    const architectureSuffix = jobArg.architecture === 'amd64' ? '' : `-${jobArg.architecture}`;
    if (imageKindArg === 'balena-raw') {
      return `baseos${architectureSuffix}.img.xz`;
    }
    return `baseos${architectureSuffix}.iso`;
  }

  private getContentType(filenameArg: string, imageKindArg: TBaseOsImageKind) {
    if (imageKindArg === 'ubuntu-iso') {
      return 'application/x-iso9660-image';
    }
    return filenameArg.endsWith('.xz') ? 'application/x-xz' : 'application/octet-stream';
  }

  private createBaseOsEnvFile(jobArg: IBaseOsImageJob) {
    return [
      `BASEOS_CLOUDLY_URL=${this.escapeEnvValue(jobArg.cloudlyUrl)}`,
      `BASEOS_JOIN_TOKEN=${this.escapeEnvValue(jobArg.provisioningToken)}`,
      'BASEOS_STATE_PATH=/var/lib/baseos/state.json',
      'BASEOS_HEARTBEAT_INTERVAL_MS=60000',
      '',
    ].join('\n');
  }

  private createBaseOsServiceFile() {
    return `[Unit]
Description=BaseOS Runner
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
EnvironmentFile=/etc/baseos/baserunner.env
WorkingDirectory=/opt/baseos
ExecStart=/usr/local/bin/deno run --allow-all /opt/baseos/mod.ts start
Restart=always
RestartSec=5

[Install]
WantedBy=multi-user.target
`;
  }

  private createBaseOsInstallScript() {
    return `#!/bin/sh
set -eu

if ! command -v /usr/local/bin/deno >/dev/null 2>&1; then
  curl -fsSL https://deno.land/install.sh | DENO_INSTALL=/usr/local sh
fi

if [ ! -d /opt/baseos/.git ]; then
  rm -rf /opt/baseos
  git clone https://code.foss.global/serve.zone/baseos.git /opt/baseos
else
  git -C /opt/baseos pull --ff-only || true
fi
`;
  }

  private escapeEnvValue(valueArg: string) {
    return JSON.stringify(valueArg);
  }

  private async runIsoCreator(configPathArg: string, logsArg: string[]) {
    const command = `${this.options.isoCreatorCommand} build --config ${this.shellQuote(configPathArg)}`;
    await this.runShellCommand(command, logsArg);
  }

  private async runShellCommand(commandArg: string, logsArg: string[]) {
    await new Promise<void>((resolve, reject) => {
      const child = spawn(commandArg, {
        shell: true,
        stdio: ['ignore', 'pipe', 'pipe'],
      });
      child.stdout.on('data', (chunk) => this.collectLog(logsArg, chunk));
      child.stderr.on('data', (chunk) => this.collectLog(logsArg, chunk));
      child.on('error', reject);
      child.on('close', (code) => {
        if (code === 0) {
          resolve();
        } else {
          reject(new Error(`Command failed with exit code ${code}: ${commandArg}`));
        }
      });
    });
  }

  private collectLog(logsArg: string[], chunkArg: Buffer) {
    const text = chunkArg.toString('utf8');
    for (const line of text.split('\n')) {
      const trimmed = line.trim();
      if (trimmed) {
        logsArg.push(trimmed);
      }
    }
    if (logsArg.length > 500) {
      logsArg.splice(0, logsArg.length - 500);
    }
  }

  private async uploadArtifact(
    jobArg: IBaseOsImageJob,
    outputPathArg: string,
    contentLengthArg: number,
    contentTypeArg: string,
    logsArg: string[],
  ) {
    logsArg.push(`Uploading artifact to ${jobArg.s3Descriptor.bucketName}/${jobArg.artifactKey}`);
    const s3Client = this.createS3Client(jobArg.s3Descriptor);
    await s3Client.send(new plugins.awsS3.PutObjectCommand({
      Bucket: jobArg.s3Descriptor.bucketName,
      Key: jobArg.artifactKey,
      Body: fs.createReadStream(outputPathArg),
      ContentLength: contentLengthArg,
      ContentType: contentTypeArg,
    }));
  }

  private createS3Client(s3DescriptorArg: IS3Descriptor) {
    return new plugins.awsS3.S3Client({
      endpoint: this.getS3Endpoint(s3DescriptorArg),
      region: s3DescriptorArg.region || 'us-east-1',
      credentials: {
        accessKeyId: s3DescriptorArg.accessKey,
        secretAccessKey: s3DescriptorArg.accessSecret,
      },
      forcePathStyle: true,
      requestChecksumCalculation: 'WHEN_REQUIRED',
    });
  }

  private getS3Endpoint(s3DescriptorArg: IS3Descriptor) {
    const rawEndpoint = s3DescriptorArg.endpoint.trim();
    const useSsl = s3DescriptorArg.useSsl !== false;
    const endpointUrl = /^https?:\/\//i.test(rawEndpoint)
      ? new URL(rawEndpoint)
      : new URL(`${useSsl ? 'https' : 'http'}://${rawEndpoint}`);
    if (s3DescriptorArg.port) {
      endpointUrl.port = String(s3DescriptorArg.port);
    }
    return endpointUrl.origin;
  }

  private async sha256File(filePathArg: string) {
    return await new Promise<string>((resolve, reject) => {
      const hash = crypto.createHash('sha256');
      const stream = fs.createReadStream(filePathArg);
      stream.on('error', reject);
      stream.on('data', (chunk) => hash.update(chunk));
      stream.on('end', () => resolve(hash.digest('hex')));
    });
  }

  private shellQuote(valueArg: string) {
    return `'${valueArg.replace(/'/g, `'\\''`)}'`;
  }

  public static getDefaultWorkdir() {
    return path.join(process.cwd(), '.nogit', 'workdir');
  }

  public static getDefaultWorkerId() {
    return `${os.hostname()}-${process.pid}`;
  }
}