/*
 * Qingtian Init Process for Enclaves
 *
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 *
 * Lots of code copied from AWS nitro enclave bootstrap, and modified to
 * be useful in qingtian_enclaves:
 * https://github.com/aws/aws-nitro-enclaves-sdk-bootstrap/blob/main/init/init.c
 * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Code from OpenGCS (https://github.com/microsoft/opengcs) is used in this file.
 * Copyright (c) Microsoft Corporation. All rights reserved.
 */
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <errno.h>
#include <sys/reboot.h>
#include <sys/wait.h>
#include <sys/socket.h>
#include <sys/mount.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/syscall.h>
#include <linux/vm_sockets.h>
#include <unistd.h>
#include <fcntl.h>

struct mount_param {
	const char *source;
	const char *target;
	const char *filesystemtype;
	unsigned long mountflags;
	const void *data;
};

struct symlink_param {
	const char *linkpath;
	const char *target;
};

struct mkdir_param {
	const char *path;
	mode_t mode;
};

struct fs_op {
	void (*fn)(struct fs_op *param);
	int ignore_errval;
	union {
		struct mount_param mount;
		struct symlink_param symlink;
		struct mkdir_param mkdir;
	} u;
};

#define ARRAY_SIZE(x)	(sizeof(x) / sizeof((x)[0]))

#define HEART_BEAT	0xB7
#define VSOCK_PORT	9000
#define VSOCK_CID	3

#define warn(fmt, ...)							\
fprintf(stderr, "qtinit: warn: %s:%d " fmt,				\
	__func__, __LINE__, ##__VA_ARGS__)

#define info(fmt, ...)							\
fprintf(stdout, "qtinit: info: %s:%d " fmt,				\
	__func__, __LINE__, ##__VA_ARGS__)

#define term_cond(cond, fmt, ...)					\
do {									\
	if ((cond))	{						\
		fprintf(stderr, "qtinit: err: %s:%d " fmt,		\
			__func__, __LINE__, ##__VA_ARGS__);		\
		exit(errno);						\
	}								\
} while (0)

#define QTSM_PATH	"/qtsm.ko"
#define DEFAULT_ENV_PATH "PATH=/sbin:/usr/sbin:/bin:/usr/bin"
#define MNT_DEFFLAGS	(MS_NOSUID | MS_NOEXEC | MS_NODEV)
#define DEV_FLAGS		(MS_NOSUID | MS_NOEXEC)
#define FILE_RIGHTS 0755

static void qtinit_block_signals(bool block)
{
	sigset_t set;
	sigfillset(&set);
	sigprocmask(block ? SIG_BLOCK : SIG_UNBLOCK, &set, 0);
}

static inline int finit_module(int fd, const char *param_values, int flags)
{
	return (int)syscall(__NR_finit_module, fd, param_values, flags);
}

static void qtinit_mount(struct fs_op *p)
{
	struct mount_param *mp = &p->u.mount;

	errno = 0;
	if (mount(mp->source, mp->target, mp->filesystemtype, mp->mountflags, mp->data) < 0)
		term_cond(errno != p->ignore_errval,
			  "mount %s failed, %s\n", mp->target, strerror(errno));

	info("mount %s success\n", mp->target);
}

static void qtinit_symlink(struct fs_op *p)
{
	struct symlink_param *sp = &p->u.symlink;

	errno = 0;
	if (symlink(sp->target, sp->linkpath) < 0)
		term_cond(errno != p->ignore_errval,
			  "symlink %s to %s failed, %s\n",
			  sp->linkpath, sp->target, strerror(errno));

	info("symlink %s to %s success\n", sp->linkpath, sp->target);
}

static void qtinit_mkdir(struct fs_op *p)
{
	struct mkdir_param *mp = &p->u.mkdir;

	errno = 0;
	if (mkdir(mp->path, mp->mode) < 0)
		term_cond(errno != p->ignore_errval,
			  "mkdir %s(%3o) failed, %s\n",
			  mp->path, mp->mode & 0777, strerror(errno));

	info("mkdir %s(%3o) success\n", mp->path, mp->mode & 0777);
}

static struct fs_op qtinit_fs[] = {
	/* dev must be the first entry */
	{qtinit_mount, EBUSY, .u.mount = {"dev", "/dev", "devtmpfs", DEV_FLAGS, NULL}},
	{qtinit_mount, EEXIST, .u.mount = {"proc", "/proc", "proc", MNT_DEFFLAGS, NULL}},
	{qtinit_symlink, EEXIST, .u.symlink = {"/dev/fd", "/proc/self/fd"}},
	{qtinit_symlink, EEXIST, .u.symlink = {"/dev/stdin", "/proc/self/fd/0"}},
	{qtinit_symlink, EEXIST, .u.symlink = {"/dev/stdout", "/proc/self/fd/1"}},
	{qtinit_symlink, EEXIST, .u.symlink = {"/dev/stderr", "/proc/self/fd/2"}},
	{qtinit_mount, EEXIST, .u.mount = {"tmpfs", "/tmp", "tmpfs", MNT_DEFFLAGS, NULL}},
	{qtinit_mount, EEXIST, .u.mount = {"tmpfs", "/run", "tmpfs", MNT_DEFFLAGS, "mode=0755"}},
	{qtinit_mkdir, EEXIST, .u.mkdir = {"/dev/shm", FILE_RIGHTS}},
	{qtinit_mount, EEXIST, .u.mount = {"shm", "/dev/shm", "tmpfs", MNT_DEFFLAGS, NULL}},
	{qtinit_mkdir, EEXIST, .u.mkdir = {"/dev/pts", FILE_RIGHTS}},
	{qtinit_mount, EEXIST, .u.mount = {"devpts", "/dev/pts", "devpts", MNT_DEFFLAGS, NULL}},
	{qtinit_mount, EEXIST, .u.mount = {"sysfs", "/sys", "sysfs", MNT_DEFFLAGS, NULL}},
	{qtinit_mount, EEXIST, .u.mount = {"cgroup_root", "/sys/fs/cgroup", "tmpfs", MNT_DEFFLAGS, "mode=0755"}},
};

static void qtinit_prepare_fs(void)
{
	int i;

	for (i = 0; i < ARRAY_SIZE(qtinit_fs); i++)
		qtinit_fs[i].fn(&qtinit_fs[i]);

	info("Prepare fs success\n");
}

static void qtinit_prepare_cgroups(void)
{
	/* read contents from /proc/cgroups */
	char target[128];
	char line[128];
	char name[64];
	int hierarchy, groups, enabled;
	FILE *fp;

	fp = fopen("/proc/cgroups", "r");
	term_cond(!fp, "open /proc/cgroups failed, %s\n", strerror(errno));

	/* skip first line */
	fgets(line, sizeof(line), fp);
	term_cond(ferror(fp), "read /proc/cgroups failed, %s\n", strerror(errno));

	while (fgets(line, sizeof(line), fp) != NULL) {
		memset(name, 0, sizeof(name));
		term_cond(sscanf(line, "%63s %d %d %d",
				 name, &hierarchy, &groups, &enabled) != 4,
			  "parse '%s' failed, %s\n", line, strerror(errno));

		if (!enabled)
			continue;

		sprintf(target, "/sys/fs/cgroup/%s", name);
		term_cond(mkdir(target, FILE_RIGHTS) < 0,
			  "mkdir %s failed, %s\n", target, strerror(errno));
		term_cond(mount(name, target, "cgroup", MNT_DEFFLAGS, name) < 0,
			  "mount %s failed, %s\n", target, strerror(errno));
	}
	term_cond(ferror(fp), "read /proc/cgroups failed, %s\n", strerror(errno));
	fclose(fp);

	info("Prepare cgroups success\n");
}

static void qtinit_prepare_console(void)
{
	const char *console = "/dev/console";

	term_cond(!freopen(console, "r", stdin),
		  "reopen %s failed, %s\n", console, strerror(errno));
	term_cond(!freopen(console, "w", stdout),
		  "reopen %s failed, %s\n", console, strerror(errno));
	term_cond(!freopen(console, "w", stderr),
		  "reopen %s failed, %s\n", console, strerror(errno));

	info("Prepare console success\n");
}

static void qtinit_install_qtsm(void)
{
	int fd;

	fd = open(QTSM_PATH, O_RDONLY | O_CLOEXEC);
	term_cond(fd < 0, "open %s failed, %s\n", QTSM_PATH, strerror(errno));
	term_cond(finit_module(fd, "", 0) < 0,
		  "insmod %s failed, %s\n", QTSM_PATH, strerror(errno));
	close(fd);

	info("Install %s success\n", QTSM_PATH);
}

static void qtinit_enclave_ready(void)
{
	int socket_fd;
	struct sockaddr_vm sa;
	unsigned char buf[1];

	memset(&sa, 0, sizeof(sa));
	sa.svm_family = AF_VSOCK;
	sa.svm_cid = VSOCK_CID;
	sa.svm_port = VSOCK_PORT;

	buf[0] = HEART_BEAT;

	socket_fd = socket(AF_VSOCK, SOCK_STREAM, 0);
	term_cond(socket_fd < 0, "socket failed, %s\n", strerror(errno));
	term_cond(connect(socket_fd, (struct sockaddr*)&sa, sizeof(sa)) < 0,
		  "connect failed, %s\n", strerror(errno));
	term_cond(write(socket_fd, buf, 1) != 1,
		  "write failed, %s\n", strerror(errno));
	term_cond(read(socket_fd, buf, 1) != 1,
		  "read failed, %s\n", strerror(errno));
	term_cond(buf[0] != HEART_BEAT,
		  "received wrong message, %s\n", strerror(errno));
	close(socket_fd);

	info("Heartbeat success\n");
}

static char **qtinit_read_config(const char *file)
{
	char line[512];
	char **s = NULL;
	int cnt;
	FILE *fp;

	fp = fopen(file, "r");
	term_cond(!fp, "open %s failed, %s\n", file, strerror(errno));

	cnt = 0;
	while (fgets(line, sizeof(line), fp) != NULL) {
		size_t len = strlen(line);
		if (line[len - 1] == '\n')
			line[len - 1] = 0;

		cnt++;
		s = realloc(s, cnt * sizeof(char *));
		term_cond(!s, "parse config failed\n");
		s[cnt - 1] = strdup(line);
		term_cond(!s[cnt-1], "parse config failed\n");
	}
	term_cond(ferror(fp), "read %s failed, %s\n", file, strerror(errno));
	if (cnt > 0) {
		s = realloc(s, (cnt + 1) * sizeof(char *));
		term_cond(!s, "parse config failed\n");
		s[cnt] = NULL;
	}
	fclose(fp);

	return s;
}

static const char *const default_env[] = {
	DEFAULT_ENV_PATH,
	NULL,
};

static const char *const default_cmd[] = {
	"sh",
	NULL,
};

static pid_t qtinit_start_app(void)
{
	char **cmd, **env;
	pid_t pid;

	cmd = qtinit_read_config("/cmd");
	if (!cmd) {
		warn("cmd is empty, using default cmd\n");
		cmd = (char **)default_cmd;
	}

	env = qtinit_read_config("/env");
	if (!env) {
		warn("env is empty, using default env\n");
		env = (char **)default_env;
	}

	pid = fork();
	if (pid != 0) {
		term_cond(pid < 0, "fork failed, %s\n", strerror(errno));
		/* Parent */
		return pid;
	}

	/* Child */
	qtinit_block_signals(false);
	setsid();
	setpgid(0, 0);

	info("Begin to run '%s'\n", cmd[0]);
	term_cond(putenv(DEFAULT_ENV_PATH),
		  "putenv '%s' failed, %s\n", DEFAULT_ENV_PATH, strerror(errno));
	execvpe(cmd[0], cmd, env);
	term_cond(true, "Should never reach here.\n");
	return pid;
}

static void qtinit_wait_app(pid_t pid)
{
	int status;

	term_cond(waitpid(pid, &status, 0) == -1,
			  "waitpid failed, %s\n", strerror(errno));
	if (WIFEXITED(status))
		info("pid %ld exited with %d\n", pid, WEXITSTATUS(status));
	else if (WIFSIGNALED(status))
		info("pid %ld exited with signal %d\n", pid, WTERMSIG(status));
	else
		info("pid %ld exited with status=0x%x\n", pid, status);
}

int main()
{
	pid_t pid;

	qtinit_block_signals(true);

	/* init /dev first */
	qtinit_fs[0].fn(&qtinit_fs[0]);
	qtinit_prepare_console();
	qtinit_prepare_fs();
	qtinit_prepare_cgroups();
	qtinit_install_qtsm();
	qtinit_enclave_ready();

	pid = qtinit_start_app();
	qtinit_wait_app(pid);
	reboot(RB_POWER_OFF);

	return 0;
}