cloudly/ts/manager.baseos/classes.baseosmanager.ts

import * as plugins from '../plugins.js';
import * as crypto from 'node:crypto';
import * as nodeStream from 'node:stream';
import type { Cloudly } from '../classes.cloudly.js';
import { logger } from '../logger.js';
import {
  BaseOsNode,
  type IBaseOsDesiredState,
  type IBaseOsNodePublic,
  type IBaseOsRuntimeInfo,
} from './classes.baseosnode.js';
import {
  BaseOsImageBuild,
  type IBaseOsImageArtifact,
  type IBaseOsImageBuildPublic,
  type TBaseOsImageArchitecture,
} from './classes.baseosimagebuild.js';

interface IBaseOsRegisterRequest {
  joinToken?: string;
  nodeToken?: string;
  status?: IBaseOsRuntimeInfo;
}

interface IBaseOsRegisterResponse {
  nodeId?: string;
  nodeToken?: string;
  accepted: boolean;
  message?: string;
}

interface IBaseOsHeartbeatRequest {
  nodeToken?: string;
  status?: IBaseOsRuntimeInfo;
}

interface IBaseOsHeartbeatResponse {
  accepted: boolean;
  message?: string;
  desiredState?: IBaseOsDesiredState;
}

interface IRequestGetBaseOsNodes {
  method: 'getBaseOsNodes';
  request: {
    identity: plugins.servezoneInterfaces.data.IIdentity;
  };
  response: {
    nodes: IBaseOsNodePublic[];
  };
}

interface IBaseOsImageBuildRequest {
  architecture: TBaseOsImageArchitecture;
  cloudlyUrl?: string;
  sourceImageUrl?: string;
  ubuntuVersion?: string;
  hostname?: string;
  wifi?: {
    ssid: string;
    password?: string;
  };
  sshPublicKey?: string;
  artifactRetentionMs?: number;
}

interface IRequestCreateBaseOsImageBuild {
  method: 'createBaseOsImageBuild';
  request: {
    identity: plugins.servezoneInterfaces.data.IIdentity;
    build: IBaseOsImageBuildRequest;
  };
  response: {
    build: IBaseOsImageBuildPublic;
  };
}

interface IRequestGetBaseOsImageBuilds {
  method: 'getBaseOsImageBuilds';
  request: {
    identity: plugins.servezoneInterfaces.data.IIdentity;
  };
  response: {
    builds: IBaseOsImageBuildPublic[];
  };
}

interface IRequestGetBaseOsImageBuildById {
  method: 'getBaseOsImageBuildById';
  request: {
    identity: plugins.servezoneInterfaces.data.IIdentity;
    buildId: string;
  };
  response: {
    build: IBaseOsImageBuildPublic;
  };
}

interface IRequestCreateBaseOsImageDownloadUrl {
  method: 'createBaseOsImageDownloadUrl';
  request: {
    identity: plugins.servezoneInterfaces.data.IIdentity;
    buildId: string;
  };
  response: {
    url: string;
    expiresAt: number;
  };
}

interface ICoreBuildBaseOsImageResponse {
  success: boolean;
  artifact?: IBaseOsImageArtifact;
  logs: string[];
  errorText?: string;
}

export class CloudlyBaseOsManager {
  public cloudlyRef: Cloudly;
  public typedRouter = new plugins.typedrequest.TypedRouter();

  public get db() {
    return this.cloudlyRef.mongodbConnector.smartdataDb;
  }

  public CBaseOsNode = plugins.smartdata.setDefaultManagerForDoc(this, BaseOsNode);
  public CBaseOsImageBuild = plugins.smartdata.setDefaultManagerForDoc(this, BaseOsImageBuild);

  constructor(cloudlyRefArg: Cloudly) {
    this.cloudlyRef = cloudlyRefArg;
    this.cloudlyRef.typedrouter.addTypedRouter(this.typedRouter);

    this.typedRouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<IRequestGetBaseOsNodes>(
        'getBaseOsNodes',
        async (requestDataArg) => {
          await plugins.smartguard.passGuardsOrReject(
            { identity: requestDataArg.identity },
            [this.cloudlyRef.authManager.adminIdentityGuard],
          );
          return {
            nodes: await this.getPublicNodes(),
          };
        },
      ),
    );

    this.typedRouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<IRequestCreateBaseOsImageBuild>(
        'createBaseOsImageBuild',
        async (requestDataArg) => {
          await plugins.smartguard.passGuardsOrReject(
            { identity: requestDataArg.identity },
            [this.cloudlyRef.authManager.adminIdentityGuard],
          );
          return {
            build: await this.createImageBuild(requestDataArg.build),
          };
        },
      ),
    );

    this.typedRouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<IRequestGetBaseOsImageBuilds>(
        'getBaseOsImageBuilds',
        async (requestDataArg) => {
          await plugins.smartguard.passGuardsOrReject(
            { identity: requestDataArg.identity },
            [this.cloudlyRef.authManager.adminIdentityGuard],
          );
          return {
            builds: await this.getPublicImageBuilds(),
          };
        },
      ),
    );

    this.typedRouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<IRequestGetBaseOsImageBuildById>(
        'getBaseOsImageBuildById',
        async (requestDataArg) => {
          await plugins.smartguard.passGuardsOrReject(
            { identity: requestDataArg.identity },
            [this.cloudlyRef.authManager.adminIdentityGuard],
          );
          return {
            build: (await this.getImageBuildById(requestDataArg.buildId)).toPublicBuild(),
          };
        },
      ),
    );

    this.typedRouter.addTypedHandler(
      new plugins.typedrequest.TypedHandler<IRequestCreateBaseOsImageDownloadUrl>(
        'createBaseOsImageDownloadUrl',
        async (requestDataArg) => {
          await plugins.smartguard.passGuardsOrReject(
            { identity: requestDataArg.identity },
            [this.cloudlyRef.authManager.adminIdentityGuard],
          );
          return await this.createImageDownloadUrl(requestDataArg.buildId);
        },
      ),
    );
  }

  public async start() {
    logger.log('info', 'BaseOS manager started');
  }

  public async stop() {
    logger.log('info', 'BaseOS manager stopped');
  }

  public async handleRegisterHttpRequest(
    reqArg: plugins.typedserver.Request,
    resArg: plugins.typedserver.Response,
  ) {
    try {
      const requestData = await this.readJsonBody<IBaseOsRegisterRequest>(reqArg);
      const response = await this.registerNode(requestData);
      this.sendJson(resArg, 200, response);
    } catch (error) {
      this.sendJson(resArg, 400, {
        accepted: false,
        message: `BaseOS registration failed: ${(error as Error).message}`,
      } satisfies IBaseOsRegisterResponse);
    }
  }

  public async handleHeartbeatHttpRequest(
    reqArg: plugins.typedserver.Request,
    resArg: plugins.typedserver.Response,
  ) {
    try {
      const requestData = await this.readJsonBody<IBaseOsHeartbeatRequest>(reqArg);
      const response = await this.acceptHeartbeat(requestData);
      this.sendJson(resArg, 200, response);
    } catch (error) {
      this.sendJson(resArg, 400, {
        accepted: false,
        message: `BaseOS heartbeat failed: ${(error as Error).message}`,
      } satisfies IBaseOsHeartbeatResponse);
    }
  }

  public async handleImageDownloadHttpRequest(
    reqArg: plugins.typedserver.Request,
    resArg: plugins.typedserver.Response,
  ) {
    try {
      const requestUrl = new URL((reqArg as any).originalUrl || reqArg.url || '/', 'http://localhost');
      const buildId = requestUrl.pathname.split('/').at(-2);
      const token = requestUrl.searchParams.get('token');
      if (!buildId || !token) {
        this.sendJson(resArg, 400, { errorText: 'build id or download token missing' });
        return;
      }
      const build = await this.getImageBuildById(buildId);
      if (build.downloadTokenHash !== this.hashSecret(token) || (build.downloadTokenExpiresAt || 0) < Date.now()) {
        this.sendJson(resArg, 403, { errorText: 'download token is invalid or expired' });
        return;
      }
      if (build.data.status !== 'ready' || !build.data.artifact) {
        this.sendJson(resArg, 409, { errorText: 'image build is not ready' });
        return;
      }

      const artifact = build.data.artifact;
      const smartbucket = new plugins.smartbucket.SmartBucket({
        ...this.cloudlyRef.config.data.s3Descriptor,
        port: Number((this.cloudlyRef.config.data.s3Descriptor as any).port || 443),
      } as any);
      const bucket = await smartbucket.getBucketByName(artifact.bucketName);
      const artifactStream = await bucket.fastGetStream({ path: artifact.key }, 'nodestream');
      resArg.status(200);
      resArg.setHeader('Content-Type', artifact.contentType || 'application/octet-stream');
      resArg.setHeader('Content-Length', String(artifact.size));
      resArg.setHeader('Content-Disposition', `attachment; filename="${artifact.filename}"`);
      (artifactStream as nodeStream.Readable).pipe(resArg as any);
    } catch (error) {
      this.sendJson(resArg, 500, {
        errorText: `BaseOS image download failed: ${(error as Error).message}`,
      });
    }
  }

  public async registerNode(
    requestDataArg: IBaseOsRegisterRequest,
  ): Promise<IBaseOsRegisterResponse> {
    if (!this.isRuntimeInfo(requestDataArg.status)) {
      return {
        accepted: false,
        message: 'BaseOS runtime status is missing or invalid',
      };
    }

    if (requestDataArg.nodeToken) {
      const existingNode = await this.getNodeByToken(requestDataArg.nodeToken);
      if (existingNode) {
        await this.updateNodeRuntimeInfo(existingNode, requestDataArg.status, true);
        return {
          accepted: true,
          nodeId: existingNode.id,
        };
      }
    }

    const acceptedBuild = await this.consumeProvisioningToken(requestDataArg.joinToken);
    if (acceptedBuild) {
      const nodeToken = await this.cloudlyRef.authManager.createNewSecureToken();
      const node = await this.upsertNode(requestDataArg.status, nodeToken);
      return {
        accepted: true,
        nodeId: node.id,
        nodeToken,
      };
    }

    const configuredJoinToken = await this.cloudlyRef.settingsManager.getSetting('baseosJoinToken');
    if (!configuredJoinToken) {
      return {
        accepted: false,
        message: 'BaseOS join token is not configured in Cloudly settings',
      };
    }

    if (!requestDataArg.joinToken || requestDataArg.joinToken !== configuredJoinToken) {
      return {
        accepted: false,
        message: 'BaseOS join token is invalid',
      };
    }

    const nodeToken = await this.cloudlyRef.authManager.createNewSecureToken();
    const node = await this.upsertNode(requestDataArg.status, nodeToken);
    return {
      accepted: true,
      nodeId: node.id,
      nodeToken,
    };
  }

  public async acceptHeartbeat(
    requestDataArg: IBaseOsHeartbeatRequest,
  ): Promise<IBaseOsHeartbeatResponse> {
    if (!requestDataArg.nodeToken) {
      return {
        accepted: false,
        message: 'BaseOS node token is missing',
      };
    }
    if (!this.isRuntimeInfo(requestDataArg.status)) {
      return {
        accepted: false,
        message: 'BaseOS runtime status is missing or invalid',
      };
    }

    const node = await this.getNodeByToken(requestDataArg.nodeToken);
    if (!node) {
      return {
        accepted: false,
        message: 'BaseOS node token is invalid',
      };
    }

    await this.updateNodeRuntimeInfo(node, requestDataArg.status, true);
    return {
      accepted: true,
      desiredState: node.data.desiredState || {},
    };
  }

  public async getPublicNodes(): Promise<IBaseOsNodePublic[]> {
    const nodes = await this.CBaseOsNode.getInstances({});
    return nodes.map((nodeArg) => nodeArg.toPublicNode());
  }

  public async createImageBuild(buildRequestArg: IBaseOsImageBuildRequest) {
    const s3Descriptor = this.cloudlyRef.config.data.s3Descriptor;
    if (!s3Descriptor?.bucketName) {
      throw new plugins.typedrequest.TypedResponseError('Cloudly S3 storage is required for BaseOS image builds');
    }

    const workerUrl = await this.cloudlyRef.settingsManager.getSetting('corebuildWorkerUrl');
    const workerToken = await this.cloudlyRef.settingsManager.getSetting('corebuildWorkerToken');
    if (!workerUrl) {
      throw new plugins.typedrequest.TypedResponseError('corebuildWorkerUrl is not configured in Cloudly settings');
    }

    const now = Date.now();
    const buildId = await this.CBaseOsImageBuild.getNewId();
    const provisioningToken = await this.cloudlyRef.authManager.createNewSecureToken();
    const artifactRetentionMs = buildRequestArg.artifactRetentionMs || 1000 * 60 * 60 * 24 * 7;
    const build = new this.CBaseOsImageBuild({
      id: buildId,
      provisioningTokenHash: this.hashSecret(provisioningToken),
      data: {
        status: 'queued',
        architecture: buildRequestArg.architecture,
        cloudlyUrl: buildRequestArg.cloudlyUrl || this.getPublicCloudlyUrl(),
        sourceImageUrl: buildRequestArg.sourceImageUrl,
        ubuntuVersion: buildRequestArg.ubuntuVersion || '24.04',
        hostname: buildRequestArg.hostname,
        wifiSsid: buildRequestArg.wifi?.ssid,
        sshPublicKey: buildRequestArg.sshPublicKey,
        logs: [],
        createdAt: now,
        updatedAt: now,
        expiresAt: now + artifactRetentionMs,
      },
    });
    await build.save();

    this.executeImageBuild(build, provisioningToken, buildRequestArg, workerUrl, workerToken).catch(async (error) => {
      build.data.status = 'failed';
      build.data.errorText = (error as Error).message;
      build.data.updatedAt = Date.now();
      build.data.completedAt = Date.now();
      await build.save();
    });

    return build.toPublicBuild();
  }

  public async getPublicImageBuilds() {
    const builds = await this.CBaseOsImageBuild.getInstances({});
    return builds
      .sort((a, b) => b.data.createdAt - a.data.createdAt)
      .map((buildArg) => buildArg.toPublicBuild());
  }

  public async createImageDownloadUrl(buildIdArg: string) {
    const build = await this.getImageBuildById(buildIdArg);
    if (build.data.status !== 'ready' || !build.data.artifact) {
      throw new plugins.typedrequest.TypedResponseError('BaseOS image build is not ready');
    }
    const token = await this.cloudlyRef.authManager.createNewSecureToken();
    const expiresAt = Date.now() + 1000 * 60 * 15;
    build.downloadTokenHash = this.hashSecret(token);
    build.downloadTokenExpiresAt = expiresAt;
    build.data.updatedAt = Date.now();
    await build.save();
    return {
      url: `/baseos/v1/images/${build.id}/download?token=${encodeURIComponent(token)}`,
      expiresAt,
    };
  }

  private async upsertNode(statusArg: IBaseOsRuntimeInfo, nodeTokenArg: string) {
    const now = Date.now();
    let node = await this.CBaseOsNode.getInstance({
      id: statusArg.nodeId,
    }).catch(() => null);

    if (!node) {
      node = new this.CBaseOsNode({
        id: statusArg.nodeId,
        nodeToken: nodeTokenArg,
        data: {
          runtimeInfo: statusArg,
          createdAt: now,
          updatedAt: now,
          lastHeartbeatAt: now,
        },
      });
    } else {
      node.nodeToken = nodeTokenArg;
      node.data = {
        ...node.data,
        runtimeInfo: statusArg,
        updatedAt: now,
        lastHeartbeatAt: now,
      };
    }

    await node.save();
    return node;
  }

  private async executeImageBuild(
    buildArg: BaseOsImageBuild,
    provisioningTokenArg: string,
    buildRequestArg: IBaseOsImageBuildRequest,
    workerUrlArg: string,
    workerTokenArg?: string,
  ) {
    buildArg.data.status = 'building';
    buildArg.data.startedAt = Date.now();
    buildArg.data.updatedAt = Date.now();
    await buildArg.save();

    const artifactFilename = buildArg.data.architecture === 'amd64'
      ? 'baseos.iso'
      : buildArg.data.architecture === 'arm64'
      ? 'baseos-arm64.iso'
      : 'baseos-rpi.img';
    const artifactKey = `corebuild/baseos/${buildArg.id}/${artifactFilename}`;
    const response = await fetch(
      new URL('/corebuild/v1/jobs/baseos-image', workerUrlArg.endsWith('/') ? workerUrlArg : `${workerUrlArg}/`),
      {
        method: 'POST',
        headers: {
          'content-type': 'application/json',
          ...(workerTokenArg ? { authorization: `Bearer ${workerTokenArg}` } : {}),
        },
        body: JSON.stringify({
          apiToken: workerTokenArg,
          job: {
            id: buildArg.id,
            architecture: buildArg.data.architecture,
            cloudlyUrl: buildArg.data.cloudlyUrl,
            provisioningToken: provisioningTokenArg,
            sourceImageUrl: buildArg.data.sourceImageUrl,
            ubuntuVersion: buildArg.data.ubuntuVersion,
            hostname: buildArg.data.hostname,
            wifi: buildArg.data.wifiSsid
              ? {
                  ssid: buildArg.data.wifiSsid,
                  password: buildRequestArg.wifi?.password,
                }
              : undefined,
            sshPublicKey: buildArg.data.sshPublicKey,
            s3Descriptor: this.cloudlyRef.config.data.s3Descriptor,
            artifactKey,
          },
        }),
      },
    );
    const responseBody = await response.json() as ICoreBuildBaseOsImageResponse;
    buildArg.data.logs = responseBody.logs || [];
    if (!response.ok || !responseBody.success || !responseBody.artifact) {
      buildArg.data.status = 'failed';
      buildArg.data.errorText = responseBody.errorText || `CoreBuild failed with HTTP ${response.status}`;
      buildArg.data.updatedAt = Date.now();
      buildArg.data.completedAt = Date.now();
      await buildArg.save();
      return;
    }

    buildArg.data.status = 'ready';
    buildArg.data.artifact = responseBody.artifact;
    buildArg.data.updatedAt = Date.now();
    buildArg.data.completedAt = Date.now();
    await buildArg.save();
  }

  private async getImageBuildById(buildIdArg: string) {
    const build = await this.CBaseOsImageBuild.getInstance({ id: buildIdArg });
    if (!build) {
      throw new plugins.typedrequest.TypedResponseError(`BaseOS image build ${buildIdArg} not found`);
    }
    return build;
  }

  private async consumeProvisioningToken(joinTokenArg?: string) {
    if (!joinTokenArg) {
      return null;
    }
    const tokenHash = this.hashSecret(joinTokenArg);
    const builds = await this.CBaseOsImageBuild.getInstances({
      provisioningTokenHash: tokenHash,
    });
    const build = builds.find((buildArg) => {
      return !buildArg.provisioningTokenConsumedAt && (!buildArg.data.expiresAt || buildArg.data.expiresAt > Date.now());
    });
    if (!build) {
      return null;
    }
    build.provisioningTokenConsumedAt = Date.now();
    build.data.updatedAt = Date.now();
    await build.save();
    return build;
  }

  private hashSecret(secretArg: string) {
    return crypto.createHash('sha256').update(secretArg).digest('hex');
  }

  private getPublicCloudlyUrl() {
    const sslMode = this.cloudlyRef.config.data.sslMode;
    const protocol = sslMode === 'none' ? 'http' : 'https';
    const port = this.cloudlyRef.config.data.publicPort;
    const includePort = sslMode === 'none' && port && !['80', '443'].includes(port);
    return `${protocol}://${this.cloudlyRef.config.data.publicUrl}${includePort ? `:${port}` : ''}`;
  }

  private async updateNodeRuntimeInfo(
    nodeArg: BaseOsNode,
    statusArg: IBaseOsRuntimeInfo,
    heartbeatArg = false,
  ) {
    nodeArg.data = {
      ...nodeArg.data,
      runtimeInfo: statusArg,
      updatedAt: Date.now(),
      ...(heartbeatArg ? { lastHeartbeatAt: Date.now() } : {}),
    };
    await nodeArg.save();
  }

  private async getNodeByToken(nodeTokenArg: string) {
    const nodes = await this.CBaseOsNode.getInstances({
      nodeToken: nodeTokenArg,
    });
    return nodes[0] || null;
  }

  private isRuntimeInfo(valueArg: unknown): valueArg is IBaseOsRuntimeInfo {
    if (!valueArg || typeof valueArg !== 'object') {
      return false;
    }
    const runtimeInfo = valueArg as Partial<IBaseOsRuntimeInfo>;
    return runtimeInfo.runtime === 'baseos'
      && typeof runtimeInfo.nodeId === 'string'
      && runtimeInfo.nodeId.length > 0
      && typeof runtimeInfo.checkedAt === 'number';
  }

  private async readJsonBody<T>(reqArg: plugins.typedserver.Request): Promise<T> {
    const chunks: Buffer[] = [];
    for await (const chunk of reqArg as any) {
      chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
    }
    const bodyString = Buffer.concat(chunks).toString('utf8').trim();
    return bodyString ? JSON.parse(bodyString) as T : {} as T;
  }

  private sendJson(
    resArg: plugins.typedserver.Response,
    statusCodeArg: number,
    bodyArg: object,
  ) {
    resArg.status(statusCodeArg);
    resArg.setHeader('Content-Type', 'application/json');
    resArg.end(JSON.stringify(bodyArg));
  }
}