/*
 * tinyrd — the rawnix initramfs.
 *
 * One file, one static binary, two roles:
 *
 *   tinyrd -o FILE    write a newc cpio whose /init is this binary
 *   /init  (pid 1)    optionally unlock a LUKS root, mount root, switch_root
 *
 * The archive carries no kernel modules, so one image serves every kernel
 * and nothing has to be regenerated per version.  The price is that the
 * unlock path must be built in:
 *
 *   BLK_DEV_INITRD DEVTMPFS BLK_DEV_DM DM_CRYPT CRYPTO_XTS
 *   CRYPTO_AES_NI_INTEL CRYPTO_SHA256 CRYPTO_HMAC
 *   CRYPTO_USER_API_SKCIPHER CRYPTO_USER_API_HASH, the root filesystem,
 *   and whatever the keyboard needs (i8042/atkbd, or xhci + usbhid).
 *
 * libcryptsetup here uses the kernel crypto backend: every hash goes
 * through AF_ALG.  Without CRYPTO_USER_API_HASH the header will not load
 * ("Cannot initialize crypto backend").
 *
 * Kernel command line:
 *
 *   root=/dev/... | UUID=<fs uuid>
 *                               required.  A device path is what the
 *                               kernel would use; UUID= survives the
 *                               disk being renamed, which /dev/sda2
 *                               does not — adding a drive, moving a
 *                               cable or booting the same image in a
 *                               VM is enough
 *   cryptroot=UUID=<luks-uuid>  optional; root= must then be
 *                               /dev/mapper/<name>, and <name> is the
 *                               mapping created
 *   rootfstype=   default ext4
 *   rootflags=    passed to mount(2) as data
 *   ro | rw       default ro, last one wins
 *   init=         default /sbin/init
 *
 * Without cryptroot= it mounts root= directly, so the same image boots an
 * unencrypted system.
 *
 * After the switch the new root sees exactly what the kernel leaves it
 * when there is no initramfs: devtmpfs on /dev, root mounted ro, nothing
 * else.  Stage 1 needs no initramfs awareness.
 *
 * Build: see the port.  It links a private static libcryptsetup and
 * libdevmapper; static is mandatory and the builder refuses otherwise.
 */

#define _GNU_SOURCE
#include <dirent.h>
#include <elf.h>
#include <errno.h>
#include <fcntl.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <termios.h>
#include <time.h>
#include <unistd.h>
#include <sys/mount.h>
#include <sys/reboot.h>
#include <sys/stat.h>

#include <libcryptsetup.h>

#define NEWROOT   "/newroot"
#define WAIT_SECS 30

/* ── builder ─────────────────────────────────────────────────────── */

/* Parents before children: the kernel's unpacker does not create them. */
static const struct {
	const char *name;
	unsigned mode, rmaj, rmin;
} skel[] = {
	{ "dev",         S_IFDIR | 0755, 0, 0 },
	{ "dev/console", S_IFCHR | 0600, 5, 1 },  /* kernel opens it before /init */
	{ "newroot",     S_IFDIR | 0755, 0, 0 },
	{ "proc",        S_IFDIR | 0555, 0, 0 },
	{ "run",         S_IFDIR | 0755, 0, 0 },
	{ "sys",         S_IFDIR | 0555, 0, 0 },
};

static void pad4(FILE *f, size_t n)
{
	static const char z[3];
	fwrite(z, 1, (4 - n % 4) % 4, f);
}

/* uid/gid 0, mtime 0: the image is reproducible and needs no root to build */
static void cpio_entry(FILE *f, unsigned ino, unsigned mode, unsigned rmaj,
                       unsigned rmin, const char *name, const void *data,
                       unsigned size)
{
	unsigned nsz = strlen(name) + 1;

	fprintf(f, "070701%08X%08X%08X%08X%08X%08X%08X%08X%08X%08X%08X%08X%08X",
	        ino, mode, 0u, 0u, S_ISDIR(mode) ? 2u : 1u, 0u, size,
	        0u, 0u, rmaj, rmin, nsz, 0u);
	fwrite(name, 1, nsz, f);
	pad4(f, 110 + nsz);
	if (size) {
		fwrite(data, 1, size, f);
		pad4(f, size);
	}
}

static unsigned char *slurp(const char *path, unsigned *size)
{
	unsigned char *buf = NULL;
	struct stat st;
	size_t off = 0;
	ssize_t r;
	int fd;

	if ((fd = open(path, O_RDONLY | O_CLOEXEC)) < 0)
		return NULL;
	if (fstat(fd, &st) == 0 && st.st_size > 0 && (buf = malloc(st.st_size)))
		while (off < (size_t)st.st_size) {
			if ((r = read(fd, buf + off, st.st_size - off)) <= 0) {
				free(buf);
				buf = NULL;
				break;
			}
			off += r;
		}
	close(fd);
	*size = off;
	return buf;
}

/*
 * A dynamic /init fails with ENOENT (its interpreter is missing), and the
 * kernel reports that as a panic nobody connects to the link flags.
 */
static int is_static(const unsigned char *b, unsigned n)
{
	Elf64_Ehdr eh;
	Elf64_Phdr ph;
	unsigned i;

	if (n < sizeof eh || memcmp(b, ELFMAG, SELFMAG) || b[EI_CLASS] != ELFCLASS64)
		return 0;
	memcpy(&eh, b, sizeof eh);
	for (i = 0; i < eh.e_phnum; i++) {
		size_t off = eh.e_phoff + (size_t)i * eh.e_phentsize;
		if (off + sizeof ph > n)
			return 0;
		memcpy(&ph, b + off, sizeof ph);
		if (ph.p_type == PT_INTERP)
			return 0;
	}
	return 1;
}

static int build(const char *out, const char *argv0)
{
	unsigned char *self;
	unsigned size, i, ino = 1;
	char tmp[4096];
	FILE *f;
	int bad;

	/* argv[0] covers a build sandbox that hides /proc */
	if (!(self = slurp("/proc/self/exe", &size)) && !(self = slurp(argv0, &size))) {
		perror("tinyrd: cannot read own executable");
		return 1;
	}
	if (!is_static(self, size)) {
		fputs("tinyrd: not a static binary; the initramfs has no libc\n", stderr);
		return 1;
	}
	snprintf(tmp, sizeof tmp, "%s.tmp", out);
	if (!(f = fopen(tmp, "wb"))) {
		perror(tmp);
		return 1;
	}
	for (i = 0; i < sizeof skel / sizeof *skel; i++)
		cpio_entry(f, ino++, skel[i].mode, skel[i].rmaj, skel[i].rmin,
		           skel[i].name, NULL, 0);
	cpio_entry(f, ino++, S_IFREG | 0755, 0, 0, "init", self, size);
	cpio_entry(f, 0, 0, 0, 0, "TRAILER!!!", NULL, 0);

	bad = ferror(f);
	bad |= fclose(f) != 0;
	if (bad || rename(tmp, out) < 0) {
		perror(out);
		unlink(tmp);
		return 1;
	}
	return 0;
}

/* ── init ────────────────────────────────────────────────────────── */

static char cmdbuf[4096];
static char *tok[256];
static int ntok;
static char founddev[288];

static void say(const char *fmt, ...)
{
	va_list ap;

	fputs("tinyrd: ", stderr);
	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);
	fputc('\n', stderr);
}

/* pid 1 must never return: that is a panic with the message scrolled away */
static void fatal(const char *fmt, ...)
{
	va_list ap;
	char c;

	fputs("tinyrd: ", stderr);
	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);
	fputs("\ntinyrd: press Enter to reboot\n", stderr);
	if (read(0, &c, 1) == 1) {
		sync();
		reboot(RB_AUTOBOOT);
	}
	for (;;)
		pause();
}

static void read_cmdline(void)
{
	ssize_t n = -1;
	char *s;
	int fd;

	if ((fd = open("/proc/cmdline", O_RDONLY | O_CLOEXEC)) >= 0) {
		n = read(fd, cmdbuf, sizeof cmdbuf - 1);
		close(fd);
	}
	if (n < 0)
		fatal("cannot read /proc/cmdline");
	cmdbuf[n] = '\0';
	for (s = strtok(cmdbuf, " \t\n"); s && ntok < 256; s = strtok(NULL, " \t\n"))
		tok[ntok++] = s;
}

/* last key=value wins, as in the kernel */
static const char *param(const char *key)
{
	size_t k = strlen(key);
	const char *v = NULL;
	int i;

	for (i = 0; i < ntok; i++)
		if (!strncmp(tok[i], key, k) && tok[i][k] == '=')
			v = tok[i] + k + 1;
	return v;
}

static int want_rw(void)
{
	int i, rw = 0;

	for (i = 0; i < ntok; i++)
		if (!strcmp(tok[i], "rw"))
			rw = 1;
		else if (!strcmp(tok[i], "ro"))
			rw = 0;
	return rw;
}

/*
 * Filesystem UUIDs, so root= need not name a device.  Each entry is a
 * magic to confirm the type and the offset of the 16 raw bytes; there
 * is no blkid here and no udev to have populated /dev/disk/by-uuid.
 */
static const struct {
	off_t       moff;
	const char *magic;
	size_t      mlen;
	off_t       uoff;
} fs_sigs[] = {
	{ 1080,    "\x53\xEF",   2, 1128    },  /* ext2/3/4 */
	{ 0,       "XFSB",       4, 32      },  /* xfs       */
	{ 0x10040, "_BHRfS_M",   8, 0x10020 },  /* btrfs     */
};

static int fs_uuid_is(const char *dev, const char *uuid)
{
	unsigned char m[8], u[16];
	char got[37];
	int fd, i, ok = 0;

	if ((fd = open(dev, O_RDONLY | O_NONBLOCK | O_CLOEXEC)) < 0)
		return 0;
	for (i = 0; !ok && i < (int)(sizeof fs_sigs / sizeof *fs_sigs); i++) {
		if (pread(fd, m, fs_sigs[i].mlen, fs_sigs[i].moff) != (ssize_t)fs_sigs[i].mlen ||
		    memcmp(m, fs_sigs[i].magic, fs_sigs[i].mlen) ||
		    pread(fd, u, sizeof u, fs_sigs[i].uoff) != (ssize_t)sizeof u)
			continue;
		snprintf(got, sizeof got,
		         "%02x%02x%02x%02x-%02x%02x-%02x%02x-%02x%02x-"
		         "%02x%02x%02x%02x%02x%02x",
		         u[0], u[1], u[2],  u[3],  u[4],  u[5],  u[6],  u[7],
		         u[8], u[9], u[10], u[11], u[12], u[13], u[14], u[15]);
		ok = !strcasecmp(got, uuid);
	}
	close(fd);
	return ok;
}

/* LUKS1 and LUKS2 both keep a 40-byte NUL-padded uuid at offset 168 */
static int luks_uuid_is(const char *dev, const char *uuid)
{
	unsigned char h[208];
	int fd, ok = 0;

	if ((fd = open(dev, O_RDONLY | O_NONBLOCK | O_CLOEXEC)) < 0)
		return 0;
	if (pread(fd, h, sizeof h, 0) == (ssize_t)sizeof h &&
	    !memcmp(h, "LUKS\xba\xbe", 6)) {
		h[207] = '\0';
		ok = !strcasecmp((char *)h + 168, uuid);
	}
	close(fd);
	return ok;
}

/* one scan for both, since the only difference is what is compared */
static int probe_uuid(const char *uuid, int luks)
{
	DIR *d = opendir("/sys/class/block");
	struct dirent *e;
	int found = 0;

	if (!d)
		return 0;
	while (!found && (e = readdir(d))) {
		if (e->d_name[0] == '.')
			continue;
		snprintf(founddev, sizeof founddev, "/dev/%s", e->d_name);
		found = luks ? luks_uuid_is(founddev, uuid)
		             : fs_uuid_is(founddev, uuid);
	}
	closedir(d);
	return found;
}

static int probe_luks(const char *uuid) { return probe_uuid(uuid, 1); }
static int probe_fs(const char *uuid)   { return probe_uuid(uuid, 0); }

static int probe_node(const char *path)
{
	struct stat st;

	return stat(path, &st) == 0 && S_ISBLK(st.st_mode);
}

/* nvme and usb probe asynchronously; /init can start before the disk exists */
static void wait_for(int (*probe)(const char *), const char *arg, const char *what)
{
	struct timespec t = { 0, 100 * 1000 * 1000 };
	int i;

	for (i = 0; !probe(arg); i++) {
		if (i == 10)
			say("waiting for %s", what);
		if (i == WAIT_SECS * 10)
			fatal("%s did not appear within %ds", what, WAIT_SECS);
		nanosleep(&t, NULL);
	}
}

static ssize_t ask(char *buf, size_t n)
{
	struct termios old, t;
	int tty = tcgetattr(0, &old) == 0;
	ssize_t r;

	if (tty) {
		t = old;
		t.c_lflag &= ~ECHO;
		t.c_lflag |= ECHONL;
		tcsetattr(0, TCSAFLUSH, &t);
	}
	fputs("Passphrase: ", stderr);
	r = read(0, buf, n);
	if (tty)
		tcsetattr(0, TCSAFLUSH, &old);
	if (r <= 0)
		return -1;
	if (buf[r - 1] == '\n')
		r--;
	return r;
}

static void cs_log(int level, const char *msg, void *usrptr)
{
	(void)usrptr;
	if (level == CRYPT_LOG_ERROR)
		fprintf(stderr, "tinyrd: %s", msg);
}

static void unlock(const char *dev, const char *name)
{
	struct crypt_device *cd = NULL;
	char pass[1024];
	ssize_t n;
	int r;

	/* no udev here: libdevmapper must create /dev/mapper/<name> itself */
	setenv("DM_DISABLE_UDEV", "1", 1);
	crypt_set_log_callback(NULL, cs_log, NULL);

	if ((r = crypt_init(&cd, dev)) < 0 || (r = crypt_load(cd, CRYPT_LUKS, NULL)) < 0)
		fatal("%s: cannot load LUKS header: %s", dev, strerror(-r));

	/* LUKS2 persistent flags (allow-discards, no_*_workqueue) apply here */
	for (;;) {
		if ((n = ask(pass, sizeof pass)) < 0)
			fatal("cannot read passphrase from console");
		r = crypt_activate_by_passphrase(cd, name, CRYPT_ANY_SLOT, pass, n, 0);
		explicit_bzero(pass, sizeof pass);
		if (r >= 0)
			break;
		if (r != -EPERM)
			fatal("activating %s on %s: %s", name, dev, strerror(-r));
		say("no key available with this passphrase");
	}
	crypt_free(cd);
	unsetenv("DM_DISABLE_UDEV");  /* would leak into stage 1's environment */
}

static int init_main(char **argv)
{
	const char *root, *cr, *fstype, *init;
	char path[4096];
	int fd;

	if (mount("devtmpfs", "/dev", "devtmpfs", 0, NULL) < 0)
		fatal("mount /dev: %s", strerror(errno));
	if (fcntl(0, F_GETFD) < 0 && (fd = open("/dev/console", O_RDWR)) >= 0) {
		dup2(fd, 0), dup2(fd, 1), dup2(fd, 2);
		if (fd > 2)
			close(fd);
	}
	if (mount("proc", "/proc", "proc", MS_NOSUID | MS_NODEV | MS_NOEXEC, NULL) < 0 ||
	    mount("sysfs", "/sys", "sysfs", MS_NOSUID | MS_NODEV | MS_NOEXEC, NULL) < 0 ||
	    mount("tmpfs", "/run", "tmpfs", MS_NOSUID | MS_NODEV, "mode=0755") < 0)
		fatal("mount proc/sys/run: %s", strerror(errno));

	read_cmdline();

	root = param("root");
	if (!root || (strncmp(root, "/dev/", 5) && strncasecmp(root, "UUID=", 5)))
		fatal("root= must be a /dev path or UUID=<fs uuid>, got %s",
		      root ? root : "nothing");

	if ((cr = param("cryptroot"))) {
		if (strncasecmp(cr, "UUID=", 5) || !cr[5])
			fatal("cryptroot= takes UUID=<luks uuid>");
		if (strncmp(root, "/dev/mapper/", 12) || !root[12] || strchr(root + 12, '/'))
			fatal("cryptroot= needs root=/dev/mapper/<name>");
		wait_for(probe_luks, cr + 5, cr);
		unlock(founddev, root + 12);
	} else if (!strncasecmp(root, "UUID=", 5)) {
		if (!root[5])
			fatal("root=UUID= needs a uuid");
		wait_for(probe_fs, root + 5, root);
		root = founddev;
	} else {
		wait_for(probe_node, root, root);
	}

	fstype = param("rootfstype");
	if (!fstype)
		fstype = "ext4";
	if (mount(root, NEWROOT, fstype, want_rw() ? 0 : MS_RDONLY, param("rootflags")) < 0)
		fatal("mount %s (%s): %s", root, fstype, strerror(errno));

	init = param("init");
	if (!init)
		init = "/sbin/init";
	snprintf(path, sizeof path, NEWROOT "%s", init);
	if (access(path, X_OK) < 0)
		fatal("%s: %s on %s", init, strerror(errno), root);

	/* hand over /dev, drop the rest; the new root mounts its own */
	if (mount("/dev", NEWROOT "/dev", NULL, MS_MOVE, NULL) < 0)
		fatal("move /dev: %s", strerror(errno));
	umount2("/run", MNT_DETACH);
	umount2("/sys", MNT_DETACH);
	umount2("/proc", MNT_DETACH);

	/* free the rootfs; we wrote the archive, so we know what is in it */
	unlink("/init");
	unlink("/dev/console");
	rmdir("/dev");
	rmdir("/proc");
	rmdir("/sys");
	rmdir("/run");

	if (chdir(NEWROOT) < 0 || mount(".", "/", NULL, MS_MOVE, NULL) < 0 ||
	    chroot(".") < 0 || chdir("/") < 0)
		fatal("switch_root: %s", strerror(errno));

	/* same argv the kernel would have given it: bare words like "single" */
	argv[0] = (char *)init;
	execv(init, argv);
	fatal("exec %s: %s", init, strerror(errno));
	return 1;
}

int main(int argc, char **argv)
{
	if (getpid() == 1)
		return init_main(argv);
	if (argc == 3 && !strcmp(argv[1], "-o"))
		return build(argv[2], argv[0]);
	fputs("usage: tinyrd -o FILE\n", stderr);
	return 2;
}