/*
 *  go32-compatible COFF, PE32 and ELF loader stub.
 *  Copyright (C) 2022,  stsp <stsp@users.sourceforge.net>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <https://www.gnu.org/licenses/>.
 *
 */
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include <unistd.h>
#include <stdint.h>
#include <fcntl.h>
#include <assert.h>
#include "stubinfo.h"
#include "dpmiwrp.h"
#include "dos.h"
#include "dpmiwrp.h"
#include "coff.h"
#include "elfp.h"
#include "util.h"
#include "stub_priv.h"
#include "djdev64/dj64init.h"
#include "djdev64/stub.h"

#define DJSTUB_VERSION 8
#define DJSTUB_API_VERSION 5

#define STUB_DEBUG 1
#if STUB_DEBUG
#define stub_debug(f, ...) J_printf(do_printf, DJ64_PRINT_LOG, "stub: " f, ##__VA_ARGS__)
#else
#define stub_debug(...)
#endif

#define FLG1_OFF STFLAGS_OFF
#define FLG2_OFF (STFLAGS_OFF + 1)

#define STFLG1_rsv2    0x20  // compact 32bit VA layout - pre-v7 stubs
#define STFLG1_STATIC  SIFLG_STATIC  // 0x40 - static linking
#define STFLG1_ELF     0x80  // ELF payload (since stub-v8)
#define STFLG2_DJ32    0x20  // dj32 payload
#define STFLG2_C32PL   0x40  // have core 32bit payload
#define STFLG2_EMBOV   0x80  // embedded overlay layout

#define HAS_32PL(buf) (buf[FLG2_OFF] & STFLG2_EMBOV)

static const int STFLAGS_OFF = 0x38;

#define MB (1024 * 1024)
#define VA_SZ (2*MB)

typedef struct
{
    uint32_t offset32;
    unsigned short selector;
} DPMI_FP;

typedef struct
{
    uint16_t rm;
    int pm;
} dpmi_dos_block;

static DPMI_FP clnt_entry;

static const char *_basename(const char *name)
{
    const char *p;
    p = strrchr(name, '\\');
    if (!p)
        p = name;
    else
        p++;
    return p;
}

#if 0
static char *_fname(char *name)
{
    char *p, *p1;
    p = strrchr(name, '\\');
    if (!p)
        p = name;
    else
        p++;
    p1 = strrchr(p, '.');
    if (p1)
        p1[0] = '\0';
    return p;
}
#endif

static unsigned _host_open(const char *pathname, unsigned mode, int *handle)
{
    int fd = open(pathname, mode | O_CLOEXEC);
    if (fd == -1)
        return 1;
    *handle = fd;
    return 0;
}

static unsigned _host_read(int handle, void *buffer, unsigned count,
        unsigned *numread)
{
    int rd = read(handle, buffer, count);
    if (rd == -1)
        return 1;
    *numread = rd;
    return 0;
}

static unsigned long _host_seek(int handle, unsigned long offset, int origin)
{
    return lseek(handle, offset, origin);
}

static int _host_close(int handle)
{
    return close(handle);
}

static struct dos_ops hops = {
    ._dos_open = _host_open,
    ._dos_read = _host_read,
    ._dos_seek = _host_seek,
    ._dos_close = _host_close,
};

__attribute__((format(printf, 3, 4)))
static void J_printf(void (*do_printf)(int prio, const char *fmt, va_list ap),
    int prio, const char *fmt, ...)
{
    va_list val;

    va_start(val, fmt);
    do_printf(prio, fmt, val);
    va_end(val);
}

#define IS_ELF(buf) (buf[0] == 0x7f && buf[1] == 0x45 && \
                buf[2] == 0x4c && buf[3] == 0x46)

#define EI_CLASS 4
#define ELFCLASS32 1
#define ELFCLASS64 2
#define ELF_IS64(buf) (buf[EI_CLASS] == ELFCLASS64)

#define SHM_NOEXEC 1
#define SHM_EXCL   2
#define SHM_NEW_NS 4
#define SHM_NS     8

#define SHM_FLAGS0 (SHM_NOEXEC | SHM_EXCL | SHM_NEW_NS)
#define SHM_FLAGS1 (SHM_NOEXEC | SHM_EXCL | SHM_NS)
#define SHM_FLAGS (SHM_FLAGS0 | (SHM_FLAGS1 << 8))

static int open_dyn(int32_t *cpl_fd, struct dos_ops **ioops,
    struct ldops **ops, int (*uput)(int))
{
    int pfile = open(CRT0, O_RDONLY | O_CLOEXEC);
    if (pfile == -1)
        return -1;
    *cpl_fd = uput(pfile);
    *ioops = &hops;
    *ops = &elf_ops;
    return pfile;
}

#define OPEN_DYN() do { \
    assert(pfile < 0); \
    pfile = open_dyn(&stubinfo.cpl_fd, &ioops, &ops, uput); \
    if (pfile == -1) { \
        error("unable to open %s\n", CRT0); \
        return -1; \
    } \
} while(0)

enum { OT_COFF, OT_DJ32, OT_E32, OT_E64, OT_D64, OT_D32, OT_MAX };

static int find_idx(int type, char *buf)
{
    uint16_t type_map;
    int cnt = 0;
    memcpy(&type_map, &buf[0x36], sizeof(type_map));
    if (type == 0 && type_map == 0)
        return 0;
    while (type_map) {
        uint8_t t = type_map & 0xf;
        if (t == type)
            return cnt;
        type_map >>= 4;
        cnt++;
    }
    return -1;
}

static int find_size(int type, char *buf)
{
    uint32_t size;
    int idx = find_idx(type, buf);
    if (idx == -1)
        return -1;
    memcpy(&size, &buf[0x1c + idx * 4], sizeof(size));
    return size;
}

static int find_offs(int type, char *buf)
{
    uint16_t type_map;
    int offs = 0, cnt = 0;
    memcpy(&type_map, &buf[0x36], sizeof(type_map));
    if (type == 0 && type_map == 0)
        return 0;
    while (type_map) {
        uint32_t off;
        uint8_t t = type_map & 0xf;
        if (t == type)
            return offs;
        memcpy(&off, &buf[0x1c + cnt * 4], sizeof(off));
        offs += off;
        type_map >>= 4;
        cnt++;
    }
    return -1;
}

static int get_type0(char *buf)
{
    if (!(buf[FLG1_OFF] & STFLG1_ELF))
        return OT_COFF;
    if (buf[FLG2_OFF] & STFLG2_DJ32)
        return OT_DJ32;
    return OT_E32;
}

#define exit(x) return -(x)
#define error(...) J_printf(do_printf, DJ64_PRINT_TERMINAL, __VA_ARGS__)
int djstub_main(int argc, char *argv[], char *envp[],
    unsigned psp_sel, int ifile, int ver,
    struct stub_ret_regs *regs, char *(*lin2ptr)(unsigned lin),
    struct dos_ops *dosops, struct dpmi_ops *dpmiops,
    void (*do_printf)(int prio, const char *fmt, va_list ap),
    int (*uput)(int), int (*elf32_open)(int), int api_ver)
{
    int pfile = -1;
    off_t coffset = 0;
    uint32_t coffsize = 0;
    uint32_t noffset = 0;
    uint32_t nsize = 0;
    uint32_t noffset2 = 0;
    uint32_t nsize2 = 0;
    int rc, i;
#define BUF_SIZE 0x40
    char buf[BUF_SIZE];
    int done = 0;
    int dyn = 0;
    int pl32 = 0;
    uint32_t va;
    uint32_t va_size;
    uint32_t mem_lin;
    uint32_t mem_base;
    unsigned short clnt_ds;
    unsigned short stubinfo_fs;
    __dpmi_meminfo info;
    dpmi_dos_block db;
    void *handle;
    _GO32_StubInfo stubinfo = {0};
    _GO32_StubInfo *stubinfo_p;
    struct ldops *ops = NULL;
    int dj32 = 0;
    struct dos_ops *ioops = dosops;
    uint8_t stub_ver = 0;

    if (ver == 0) {
        /* backward-compat code */
        stub_debug("Opening self at %s\n", argv[0]);
        rc = dosops->_dos_open(argv[0], O_RDONLY, &ifile);
        if (rc) {
            error("cannot open %s\n", argv[0]);
            exit(EXIT_FAILURE);
        }
        ver = DJSTUB_API_VERSION;
    }
    if ((ver & 0xff) != DJSTUB_API_VERSION) {
        error("Stub version mismatch: want 0x%x got 0x%x\n", DJSTUB_API_VERSION, ver);
//        exit(1);
    }

    register_dpmiops(dpmiops);
    stubinfo.cpl_fd = -1;
    for (i = 0; envp && envp[i]; i++) {
        const char *s = "ELFLOAD=";
        int l = strlen(s);
        int el = 0, ee = 0;

        if (strncmp(envp[i], s, l) == 0) {
            el++;
        } else {
            s = "ELFEXEC=";
            l = strlen(s);
            if (strncmp(envp[i], s, l) == 0)
                ee++;
        }
        if (el) {
            if (api_ver >= 23)
                pfile = elf32_open(atoi(envp[i] + l));
            if (pfile < 0) {
                /* 64bit elf */
                OPEN_DYN();
                stubinfo.elfload_arg = atoi(envp[i] + l);
                dyn = 1;
                /* Not setting emb_ov flags: pfile closed, no payload,
                 * second loader will re-check everything anyway.
                 * But we set pl32 below to reserve space for upl. */
            } else {
                /* 32bit elf */
                dj32 = 1;
                ioops = &hops;
                ops = &elf_ops;
            }
            dosops->_dos_close(ifile);
            ifile = -1;  // load dynamically via API
        }
        if (ee) {
            unsigned rd;

            dosops->_dos_close(ifile);
            ifile = -1;
            rc = dosops->_dos_open(envp[i] + l, O_RDONLY, &pfile);
            if (rc || pfile < 0) {
                error("cannot open %s\n", envp[i] + l);
                exit(EXIT_FAILURE);
            }
            rc = dosops->_dos_read(pfile, buf, BUF_SIZE, &rd);
            if (rc) {
                dosops->_dos_close(pfile);
                error("stub: read() failure\n");
                exit(EXIT_FAILURE);
            }
            if (rd != BUF_SIZE) {
                dosops->_dos_close(pfile);
                error("stub: read(%i)=%i, wrong exe file\n", BUF_SIZE, rd);
                exit(EXIT_FAILURE);
            }
            if (!IS_ELF(buf)) {
                dosops->_dos_close(pfile);
                error("stub: not an ELF: %s\n", envp[i] + l);
                exit(EXIT_FAILURE);
            }
            dosops->_dos_seek(pfile, 0, SEEK_SET);
            if (ELF_IS64(buf)) {
                ifile = pfile;
                pfile = -1;
                OPEN_DYN();
                stubinfo.elfload_arg = atoi(envp[i] + l);
                dyn = 1;

                /* calling second ldr */
                stubinfo.flags = ((STFLG2_EMBOV) << 8);
                stubinfo.flags |= SHM_FLAGS;
                nsize = dosops->_dos_seek(ifile, 0, SEEK_END);
            } else {
                dj32 = 1;
                ops = &elf_ops;
            }
        }
        if (el || ee) {
            pl32 = 1;  // even for ELFLOAD we reserve space for pl32
            done = 1;
            stubinfo.flags |= SIFLG_ELFEXEC;
            break;
        }
    }
    while (!done) {
        unsigned rd;
#if STUB_DEBUG
        int cnt = 0;
#endif

        stub_debug("Expecting header at 0x%lx\n", coffset);
        rc = dosops->_dos_read(ifile, buf, BUF_SIZE, &rd);
        if (rc) {
            error("stub: read() failure\n");
            exit(EXIT_FAILURE);
        }
        if (rd != BUF_SIZE) {
            error("stub: read(%i)=%i, wrong exe file\n", BUF_SIZE, rd);
            exit(EXIT_FAILURE);
        }
        if (buf[0] == 'M' && buf[1] == 'Z' && buf[8] == 4 && buf[9] == 0) {
            /* lfanew */
            uint32_t offs;
            int moff, sz;

            stub_ver = buf[0x3b];
#if STUB_DEBUG
            cnt++;
#endif
            stubinfo.stubinfo_ver |= stub_ver << 16;
            if (stub_ver >= 7)
                stubinfo.flags |= SIFLG_SPLITPL;
            stub_debug("Found exe header %i at 0x%lx\n", cnt, coffset);
            memcpy(&offs, &buf[0x3c], sizeof(offs));
            if (!(buf[FLG2_OFF] & STFLG2_C32PL) && (stub_ver < 8 ||
                    (buf[FLG1_OFF] & STFLG1_ELF)))
                dyn++;
            else
                pfile = ifile;
            if (HAS_32PL(buf))
                pl32++;
            if (buf[FLG2_OFF] & STFLG2_DJ32) {
                done = 1;
                dj32 = 1;
                ops = &elf_ops;
            }
            if (stub_ver >= 8 && (buf[FLG1_OFF] & STFLG1_ELF)) {
                done = 1;
                if (dyn)
                    OPEN_DYN();
                else
                    ops = &elf_ops;
            }

            coffset = offs;
            if (stub_ver <= 7) {
                noffset = offs;
                if (!dyn) {
                    memcpy(&coffsize, &buf[0x1c], sizeof(coffsize));
                    noffset += coffsize;
                }
                memcpy(&nsize, &buf[0x20 - dyn * 4], sizeof(nsize));
            } else {
                if ((sz = find_size(get_type0(buf), buf)) != -1)
                    coffsize = sz;
                if ((moff = find_offs(OT_E64, buf)) != -1)
                    noffset = offs + moff;
                if ((sz = find_size(OT_E64, buf)) != -1)
                    nsize = sz;
                if ((moff = find_offs(OT_D64, buf)) != -1)
                    noffset2 = offs + moff;
                if ((sz = find_size(OT_D64, buf)) != -1)
                    nsize2 = sz;
                if ((!(buf[FLG1_OFF] & STFLG1_ELF) || !dyn) && !coffsize)
                    exit(EXIT_FAILURE);
            }

            memcpy(&stubinfo.flags, &buf[STFLAGS_OFF], 2);
            if (stub_ver >= 6) {
                uint32_t nmoffs;
                memcpy(&nmoffs, &buf[0x28], sizeof(nmoffs));
                if (nmoffs) {
                    dosops->_dos_seek(ifile, noffset + nmoffs, SEEK_SET);
                    rc = dosops->_dos_read(ifile, stubinfo.payload2_name,
                            16, &rd);
                    stubinfo.payload2_name[rd] = '\0';
                }
            } else {
                error("unsupported stub version %i\n", stub_ver);
                return -1;
            }
            stub_debug("suppl name %s\n", stubinfo.payload2_name[0] ?
                    stubinfo.payload2_name : "<not set>");
        } else if (buf[0] == 0x4c && buf[1] == 0x01) { /* it's a COFF */
            if (dyn) {
                /* undo mistake: no dyn with COFF */
                assert(stub_ver < 8);  // no such mistake for v8+
                dyn = 0;
                pl32 = 1;
                pfile = ifile;
            }
            done = 1;
            ops = &coff_ops;
        } else if (IS_ELF(buf)) { /* it's an ELF */
            int is_64 = ELF_IS64(buf);
            if (!coffset) {
                assert(!dyn);
                if (is_64) {
                    dyn++;
                    OPEN_DYN();
                    stubinfo.flags = (STFLG2_EMBOV) << 8;
                    stubinfo.flags |= SHM_FLAGS;
                    nsize = dosops->_dos_seek(ifile, 0, SEEK_END);
                } else {
                    stubinfo.flags = ((STFLG2_DJ32) << 8);
                    dj32 = 1;
                    pfile = ifile;
                    ops = &elf_ops;
                }
                pl32 = 1;
            } else if (stub_ver >= 8) {
                error("djstub: ELF at position %lx (expected COFF)\n", coffset);
                return -1;
            } else if (is_64 && (stub_ver < 7 || !dyn)) {
                error("djstub: 64bit ELF at position %lx\n", coffset);
                return -1;
            } else if (!dyn) {
                ops = &elf_ops;
            }
            done = 1;
        } else {
            error("not an exe %s at %lx\n", argv[0], coffset);
            exit(EXIT_FAILURE);
        }
        dosops->_dos_seek(ifile, coffset, SEEK_SET);
    }
    if (dyn && coffset && stub_ver < 8)
        OPEN_DYN();
    assert(ops);
    assert(pfile != -1);

    if (dj32)
        strncpy(stubinfo.magic, "dj32 (C) stsp", sizeof(stubinfo.magic));
    else
        strncpy(stubinfo.magic, "dj64 (C) stsp", sizeof(stubinfo.magic));
    stubinfo.size = sizeof(stubinfo);
    i = 3;
    while(envp && *envp) {
        i += strlen(*envp) + 1;
        envp++;
    }
    if (argv && argv[0])
        i += strlen(argv[0]) + 1;
    stub_debug("env size %i\n", i);
    stubinfo.env_size = i;
    stubinfo.minstack = 0x80000;
    stubinfo.minkeep = 0x4000;
    if (argv) {
        strncpy(stubinfo.argv0, _basename(argv[0]), sizeof(stubinfo.argv0));
        stubinfo.argv0[sizeof(stubinfo.argv0) - 1] = '\0';
    } else {
        stubinfo.argv0[0] = '\0';
    }
    /* basename seems unused and produces warning about missing 0-terminator */
//    strncpy(stubinfo.basename, _fname(argv0), sizeof(stubinfo.basename));
//    stubinfo.basename[sizeof(stubinfo.basename) - 1] = '\0';
    strncpy(stubinfo.dpmi_server, "CWSDPMI.EXE", sizeof(stubinfo.dpmi_server));
#define max(a, b) ((a) > (b) ? (a) : (b))
    stubinfo.psp_selector = psp_sel;
    /* DJGPP relies on ds_selector, cs_selector and ds_segment all mapping
     * the same real-mode memory block. */
    dosops->_dos_link_umb(1);
    db.rm = __dpmi_allocate_dos_memory(stubinfo.minkeep >> 4, &db.pm);
    dosops->_dos_link_umb(0);
    stub_debug("rm seg 0x%x\n", db.rm);
    stubinfo.ds_selector = db.pm;
    stubinfo.ds_segment = db.rm;
    /* create alias */
    stubinfo.cs_selector = __dpmi_create_alias_descriptor(db.pm);
    /* set descriptor access rights */
    __dpmi_set_descriptor_access_rights(stubinfo.cs_selector, 0x00fb);

    clnt_entry.selector = __dpmi_allocate_ldt_descriptors(1);
    clnt_ds = __dpmi_allocate_ldt_descriptors(1);

    register_dosops(ioops);
    handle = ops->read_headers(pfile);
    if (!handle)
        exit(EXIT_FAILURE);
    va = ops->get_va(handle);
    va_size = ops->get_length(handle);
    clnt_entry.offset32 = ops->get_entry(handle);
    stub_debug("va 0x%x va_size 0x%x entry 0x%x\n",
            va, va_size, clnt_entry.offset32);
    if (!dj32 && va_size > MB)
        exit(EXIT_FAILURE);
    /* if we load 2 payloads, use larger estimate */
    if (pl32) {
        stubinfo.initial_size = VA_SZ;
        stubinfo.upl_base = va;
        stubinfo.upl_size = VA_SZ;
    } else {
        stubinfo.initial_size = max(va_size, 0x10000);
    }
    info.size = PAGE_ALIGN(stubinfo.initial_size);
    /* allocate mem */
    rc = __dpmi_allocate_memory(&info);
    if (rc)
        exit(EXIT_FAILURE);
    stubinfo.memory_handle = info.handle;
    mem_lin = info.address;
    mem_base = mem_lin - va;
    stubinfo.mem_base = mem_base;
    stub_debug("mem_lin 0x%x mem_base 0x%x\n", mem_lin, mem_base);
    ops->read_sections(handle, lin2ptr(mem_lin), va, pfile, dyn ? 0 : coffset);
    ops->close(handle);
    /* host fd can clash with DOS fd, so check also ioops */
    if (dj32 && (pfile != ifile || ioops == &hops))
        __dos_close(pfile);
    unregister_dosops();

    /* set base */
    __dpmi_set_segment_base_address(clnt_entry.selector, mem_base);
    /* set descriptor access rights */
    __dpmi_set_descriptor_access_rights(clnt_entry.selector, 0xc0fb);
    /* set limit */
    __dpmi_set_segment_limit(clnt_entry.selector, 0xffffffff);
    /* set base */
    __dpmi_set_segment_base_address(clnt_ds, mem_base);
    /* set descriptor access rights */
    __dpmi_set_descriptor_access_rights(clnt_ds, 0xc0f3);
    /* set limit */
    __dpmi_set_segment_limit(clnt_ds, 0xffffffff);

    stubinfo_fs = __dpmi_allocate_ldt_descriptors(1);
    info.size = PAGE_ALIGN(sizeof(_GO32_StubInfo));
    __dpmi_allocate_memory(&info);
    __dpmi_set_segment_base_address(stubinfo_fs, info.address);
    __dpmi_set_segment_limit(stubinfo_fs, sizeof(_GO32_StubInfo) - 1);
    stubinfo_p = (_GO32_StubInfo *)lin2ptr(info.address);

    stubinfo.self_fd = ifile;
    stubinfo.self_offs = coffset;
    stubinfo.self_size = coffsize;
    stubinfo.payload_offs = noffset;
    stubinfo.payload_size = nsize;
    stubinfo.payload2_offs = noffset2;
    stubinfo.payload2_size = nsize2;
    dosops->_dos_seek(ifile, noffset, SEEK_SET);
    if (nsize > 0)
        stub_debug("Found payload of size %i at 0x%x\n", nsize, noffset);
    stubinfo.stubinfo_ver |= DJSTUB_VERSION | ((ver & 0xff00) << 16);

    unregister_dpmiops();

    memcpy(stubinfo_p, &stubinfo, sizeof(stubinfo));
    stub_debug("Jump to entry...\n");
    regs->fs = stubinfo_fs;
    regs->ds = clnt_ds;
    regs->cs = clnt_entry.selector;
    regs->eip = clnt_entry.offset32;
    return 0;
}
