socketcomputer/qemu/src/QemuVM.ts

import { execa, ExecaChildProcess } from 'execa';
import { EventEmitter } from 'events';
import QmpClient from './QmpClient.js';
import { QemuDisplay } from './QemuDisplay.js';
import { unlink } from 'node:fs/promises';

export enum VMState {
	Stopped,
	Starting,
	Started,
	Stopping
}

export type QemuVmDefinition = {
	id: string;
	command: string[];
};

/// Temporary path base (for UNIX sockets/etc.)
const kVmTmpPathBase = `/tmp`;

/// The max amount of times QMP connection is allowed to fail before
/// the VM is forcefully stopped.
const kMaxFailCount = 5;

let gVMShouldSnapshot = false;

async function Sleep(ms: number) {
	return new Promise((resolve) => setTimeout(resolve, ms));
}

export function setSnapshot(val: boolean) {
	gVMShouldSnapshot = val;
}

export class QemuVM extends EventEmitter {
	private state = VMState.Stopped;

	private qmpInstance: QmpClient | null = null;
	private qmpConnected = false;
	private qmpFailCount = 0;

	private qemuProcess: ExecaChildProcess | null = null;
	private qemuRunning = false;

	private display: QemuDisplay | null = null;

	private definition: QemuVmDefinition;
	private addedCommandShit = false;

	constructor(def: QemuVmDefinition) {
		super();
		this.definition = def;
	}

	async Start() {
		// Don't start while either trying to start or starting.
		if (this.state == VMState.Started || this.state == VMState.Starting) return;


		var cmd = this.definition.command;

		// build additional command line statements to enable qmp/vnc over unix sockets
		if(!this.addedCommandShit) {
			cmd.push('-no-shutdown');
			if(gVMShouldSnapshot)
				cmd.push('-snapshot');
			cmd.push(`-qmp`);
			cmd.push(`unix:${this.GetQmpPath()},server,nowait`);
			cmd.push(`-vnc`);
			cmd.push(`unix:${this.GetVncPath()}`);
			this.addedCommandShit = true;
		}

		this.VMLog(`Starting QEMU with command \"${cmd.join(' ')}\"`);
		await this.StartQemu(cmd);
	}

	async Stop() {
		// This is called in certain lifecycle places where we can't safely assert state yet
		//this.AssertState(VMState.Started, 'cannot use QemuVM#Stop on a non-started VM');

		// Start indicating we're stopping, so we don't
		// erroneously start trying to restart everything
		// we're going to tear down in this function call.
		this.SetState(VMState.Stopping);

		// Kill the QEMU process and QMP/display connections if they are running.
		await this.DisconnectQmp();
		this.DisconnectDisplay();
		await this.StopQemu();
	}

	async Reset() {
		this.AssertState(VMState.Started, 'cannot use QemuVM#Reset on a non-started VM');

		// let code know the VM is going to reset
		// N.B: In the crusttest world, a reset simply amounts to a
		// mean cold reboot of the qemu process basically
		this.emit('reset');
		await this.Stop();
		await Sleep(500);
		await this.Start();
	}

	async QmpCommand(command: string, args: any | null): Promise<any> {
		return await this.qmpInstance?.Execute(command, args);
	}

	async MonitorCommand(command: string) {
		this.AssertState(VMState.Started, 'cannot use QemuVM#MonitorCommand on a non-started VM');
		return await this.QmpCommand('human-monitor-command', {
			'command-line': command
		});
	}

	async ChangeRemovableMedia(deviceName: string, imagePath: string): Promise<void> {
		this.AssertState(VMState.Started, 'cannot use QemuVM#ChangeRemovableMedia on a non-started VM');
		// N.B: if this throws, the code which called this should handle the error accordingly
		await this.QmpCommand('blockdev-change-medium', {
			device: deviceName, // techinically deprecated, but I don't feel like figuring out QOM path just for a simple function
			filename: imagePath
		});
	}

	async EjectRemovableMedia(deviceName: string) {
		this.AssertState(VMState.Started, 'cannot use QemuVM#EjectRemovableMedia on a non-started VM');
		await this.QmpCommand('eject', {
			device: deviceName
		});
	}

	GetDisplay() {
		return this.display;
	}

	/// Private fun bits :)

	private VMLog(...args: any[]) {
		// TODO: hook this into a logger of some sort
		console.log(`[QemuVM] [${this.definition.id}] ${args.join('')}`);
	}

	private AssertState(stateShouldBe: VMState, message: string) {
		if (this.state !== stateShouldBe) throw new Error(message);
	}

	private SetState(state: VMState) {
		this.state = state;
		this.emit('statechange', this.state);
	}

	private GetQmpPath() {
		return `${kVmTmpPathBase}/socket2-${this.definition.id}-ctrl`;
	}

	private GetVncPath() {
		return `${kVmTmpPathBase}/socket2-${this.definition.id}-vnc`;
	}

	private async StartQemu(split: Array<string>) {
		let self = this;

		this.SetState(VMState.Starting);

		// Start QEMU
		this.qemuProcess = execa(split[0], split.slice(1));

		this.qemuProcess.on('spawn', async () => {
			self.qemuRunning = true;
			await Sleep(500);
			await self.ConnectQmp();
		});

		this.qemuProcess.on('exit', async (code) => {
			self.qemuRunning = false;
			console.log("qemu process go boom")

			// ?
			if (self.qmpConnected) {
				await self.DisconnectQmp();
			}
			
			self.DisconnectDisplay();

			if (self.state != VMState.Stopping) {
				if (code == 0) {
					await Sleep(500);
					await self.StartQemu(split);
				} else {
					self.VMLog('QEMU exited with a non-zero exit code. This usually means an error in the command line. Stopping VM.');
					await self.Stop();
				}
			} else {
				this.SetState(VMState.Stopped);
			}
		});
	}

	private async StopQemu() {
		if (this.qemuRunning == true) this.qemuProcess?.kill('SIGKILL');
	}

	private async ConnectQmp() {
		let self = this;

		if (!this.qmpConnected) {
			self.qmpInstance = new QmpClient();

			self.qmpInstance.on('close', async () => {
				self.qmpConnected = false;

				// If we aren't stopping, then we do actually need to care QMP disconnected
				if (self.state != VMState.Stopping) {
					if (self.qmpFailCount++ < kMaxFailCount) {
						this.VMLog(`Failed to connect to QMP ${self.qmpFailCount} times`);
						await Sleep(500);
						await self.ConnectQmp();
					} else {
						this.VMLog(`Failed to connect to QMP ${self.qmpFailCount} times, giving up`);
						await self.Stop();
					}
				}
			});

			self.qmpInstance.on('event', async (ev) => {
				switch (ev.event) {
					// Handle the STOP event sent when using -no-shutdown
					case 'STOP':
						await self.qmpInstance?.Execute('system_reset');
						break;
					case 'RESET':
						self.VMLog('got a reset event!');
						await self.qmpInstance?.Execute('cont');
						break;
				}
			});

			self.qmpInstance.on('qmp-ready', async (hadError) => {
				self.VMLog('QMP ready');

				self.display = new QemuDisplay(self.GetVncPath());
				self.display.Connect();

				// QMP has been connected so the VM is ready to be considered started
				self.qmpFailCount = 0;
				self.qmpConnected = true;
				self.SetState(VMState.Started);
			});

			try {
				await Sleep(500);
				this.qmpInstance?.ConnectUNIX(this.GetQmpPath());
			} catch (err) {
				// just try again
				await Sleep(500);
				await this.ConnectQmp();
			}
		}
	}

	private async DisconnectDisplay() {
		try {
			this.display?.Disconnect();
			this.display = null; // disassociate with that display object.

			await unlink(this.GetVncPath());
		} catch (err) {
			// oh well lol
		}
	}

	private async DisconnectQmp() {
		if (this.qmpConnected) return;
		if(this.qmpInstance == null)
			return;

		this.qmpConnected = false;
		this.qmpInstance.end();
		this.qmpInstance = null;
		try {
			await unlink(this.GetQmpPath());
		} catch(err) {
			
		}
	}
}