#include <linux/elfcore.h>
#include <linux/elf.h>
#include <linux/kdebug.h>
#include <linux/slab.h>
#include <linux/ioport.h>
#include <asm/stacktrace.h>
#include <asm/irq_regs.h>
#include <asm/cacheflush.h>

#include "conf.h"
#include "net.h"

#define NOTE_HEAD_BYTES sizeof(struct elf_note)
#define CORE_NOTE_NAME "CORE"
#define CORE_NOTE_NAME_BYTES sizeof(CORE_NOTE_NAME)
#define CORE_NOTE_DESC_BYTES sizeof(struct elf_prstatus)
#define NOTE_BYTES ALIGN(NOTE_HEAD_BYTES + CORE_NOTE_NAME_BYTES + CORE_NOTE_DESC_BYTES, 4)

/* Alignment required for elf header segment */
#define ELF_CORE_HEADER_ALIGN   4096

/*
 * Special Flag only used in avm_coredump to denote segments that are
 * inside other load segments, meaning that their content is already transmitted
 *
 * Therefore those segments need special handeling to caluclate the correct
 * poffset
 */
#define PF_INSIDE 0x10

struct crash_elf_data {
	/* Pointer to elf header */
	void *ehdr;
	/* Pointer to next phdr */
	void *bufp;
};

struct ram_region {
    // Start and end in phy space
    u64 start;
    u64 end;
    // offset in file
    int offset;
};

static u8 *append_elf_note(u8 *buf, char *name, unsigned type, void *data,
			    size_t data_len)
{
	struct elf32_note note;

	note.n_namesz = strlen(name) + 1;
	note.n_descsz = data_len;
	note.n_type   = type;
	memcpy(buf, &note, sizeof(note));
	buf += ALIGN(sizeof(note), 4);
	memcpy(buf, name, note.n_namesz);
	buf += ALIGN(note.n_namesz, 4);
	memcpy(buf, data, note.n_descsz);
	buf += ALIGN(note.n_descsz, 4);

	return buf;
}

static void crash_save_cpu(struct pt_regs *regs, unsigned int cpu)
{
    struct elf_prstatus prstatus;

    if (regs == NULL) {
        return;
    }

    memset(&prstatus, 0, sizeof(prstatus));
    prstatus.pr_pid = cpu;
    elf_core_copy_kernel_regs(&prstatus.pr_reg, regs);
    append_elf_note(conf->coredump.crash_notes + NOTE_BYTES * cpu, CORE_NOTE_NAME, NT_PRSTATUS, &prstatus, sizeof(prstatus));
}

static int elf_add_ram(struct ram_region *reg, void *arg)
{
	struct crash_elf_data *ced = arg;
	Elf32_Ehdr *ehdr;
	Elf32_Phdr *phdr;

	ehdr = ced->ehdr;

	phdr = ced->bufp;
	ced->bufp += sizeof(Elf32_Phdr);

	phdr->p_type = PT_LOAD;
	phdr->p_flags = PF_R|PF_W|PF_X;
	phdr->p_offset  = reg->offset;

	phdr->p_paddr = reg->start;
	phdr->p_vaddr = (unsigned long) __va(reg->start);
	phdr->p_filesz = phdr->p_memsz = reg->end - reg->start + 1;
	phdr->p_align = 0;
	ehdr->e_phnum++;
	//printk(KERN_ERR "Crash PT_LOAD elf header. phdr=%p vaddr=0x%x, paddr=0x%x, sz=0x%x e_phnum=%d p_offset=0x%x\n",
	//		phdr, phdr->p_vaddr, phdr->p_paddr, phdr->p_filesz,
	//		ehdr->e_phnum, phdr->p_offset);

	return 0;
}

static int prepare_elf32_headers(void *buf, unsigned long sz)
{
    Elf32_Ehdr *ehdr;
    Elf32_Phdr *phdr;
    unsigned long nr_cpus = num_possible_cpus();
    unsigned char *bufp;
    unsigned int i;
    unsigned int notes_addr;
    int ret;
	struct crash_elf_data ced;

    bufp = buf;
    ehdr = (Elf32_Ehdr *)bufp;
    bufp += sizeof(Elf32_Ehdr);
    memcpy(ehdr->e_ident, ELFMAG, SELFMAG);
    ehdr->e_ident[EI_CLASS] = ELFCLASS32;
#if defined(CONFIG_CPU_BIG_ENDIAN)
    ehdr->e_ident[EI_DATA] = ELFDATA2MSB;
#else
    ehdr->e_ident[EI_DATA] = ELFDATA2LSB;
#endif
    ehdr->e_ident[EI_VERSION] = EV_CURRENT;
    ehdr->e_ident[EI_OSABI] = ELF_OSABI;
    memset(ehdr->e_ident + EI_PAD, 0, EI_NIDENT - EI_PAD);
    ehdr->e_type = ET_CORE;
    ehdr->e_machine = ELF_ARCH;
    ehdr->e_version = EV_CURRENT;
    ehdr->e_phoff = sizeof(Elf32_Ehdr);
    ehdr->e_ehsize = sizeof(Elf32_Ehdr);
    ehdr->e_phentsize = sizeof(Elf32_Phdr);

    /* Prepare one phdr of type PT_NOTE for each present cpu */
    phdr = (Elf32_Phdr *)bufp;
    bufp += sizeof(Elf32_Phdr);
    phdr->p_type = PT_NOTE;
    phdr->p_flags = PF_INSIDE;
    notes_addr = __pa(conf->coredump.crash_notes);
    phdr->p_offset = phdr->p_paddr = notes_addr;
    phdr->p_filesz = phdr->p_memsz = NOTE_BYTES * nr_cpus;
    (ehdr->e_phnum)++;

    /* Prepare PT_LOAD headers for system ram chunks. */
    ced.ehdr = ehdr;
    ced.bufp = bufp;
    for(i = 0; i < conf->coredump.num_ram_regions; i++) {
        ret = elf_add_ram(conf->coredump.ram_regions + i, &ced);
        if (ret < 0)
            return ret;
    }

    return 0;
}

/*
 * ELF files contain p_offset in each program header describing
 * the offset into the file. Currently the p_offset describes the offset
 * into the physical ram. However in the final file, this is prepended by
 * the elf header, therefore we need to advance the offset by the elf
 * headers length.
 */
static void core_fixup_offsets(void *elfcorehdr, unsigned int len) {
    Elf32_Ehdr *ehdr;
    Elf32_Phdr *phdr, *phdr2;
    unsigned char *bufp, *bufp2;
    unsigned int i, j;

    bufp = elfcorehdr;
    ehdr = (Elf32_Ehdr*)bufp;
    bufp += ehdr->e_phoff;

    // First we fix the offsets of load segments as they represent the memory we're transmitting
    for (i = 0; i < ehdr->e_phnum; i++, bufp += ehdr->e_phentsize) {
        phdr = (Elf32_Phdr*)bufp;
        if (phdr->p_flags & PF_INSIDE) {
            continue;
        }
        phdr->p_offset += len;

    }

    bufp = elfcorehdr + ehdr->e_phoff;
    // Secondly we fixup the file offset of note segements that are inside one of those load segments
    for (i = 0; i < ehdr->e_phnum; i++, bufp += ehdr->e_phentsize) {
        phdr = (Elf32_Phdr*)bufp;
        if ((phdr->p_flags & PF_INSIDE) == 0) {
            continue;
        }

        // Find corresponding LOAD
        bufp2 = elfcorehdr + ehdr->e_phoff;
        for (j = 0; j < ehdr->e_phnum; j++, bufp2 += ehdr->e_phentsize) {
            phdr2 = (Elf32_Phdr*)bufp2;
            if (phdr2->p_flags & PF_INSIDE) {
                continue;
            }

            // Note (phdr) is within the LOAD (phdr2)
            if (phdr2->p_paddr <= phdr->p_paddr && phdr->p_paddr <= phdr2->p_paddr + phdr2->p_memsz) {
                break;
            }
        }

        if (j >= ehdr->e_phnum) {
            pr_err(LOG_PREFIX "Found no phy space for phdr %d\n", i);
            continue;
        }

        phdr->p_offset = phdr2->p_offset + (phdr->p_paddr - phdr2->p_paddr);
    }

    // Third: we remove the internal PF_INSIDE flag
    bufp = elfcorehdr + ehdr->e_phoff;
    for (i = 0; i < ehdr->e_phnum; i++, bufp += ehdr->e_phentsize) {
        phdr = (Elf32_Phdr*)bufp;
        phdr->p_flags &= ~PF_INSIDE;

    }
}

static int avm_coredump_set_regs(struct notifier_block *this, unsigned long cmd, void *ptr) {
	if (cmd == DIE_OOPS) {
        struct die_args *args = (struct die_args *)ptr;

        conf->coredump.regs = args->regs;
    }

    return NOTIFY_OK;
}

static int send_buf(u8 *buf, u32 size) {
    int ret;
    u8 *tbuf;
    u32 available_size;
    u64 block_num = cpu_to_be64(conf->coredump.block_num);

    avm_coredump_net_prepare_pkg(conf->coredump_port, &tbuf, &available_size);

    memcpy(tbuf, &block_num, sizeof(u64));
    memcpy(tbuf + sizeof(u64), buf, size);

    ret = avm_coredump_net_send_pkg(sizeof(u64) + size);

    if (ret == 0) {
        conf->coredump.block_num += 1;
        conf->coredump.send_bytes += size;
    }

    return ret;
}

int avm_coredump_coredump_do(void) {
	int ret;
    unsigned int i, j;
    unsigned char *buf;
    u64 addr;
#ifdef CONFIG_SMP
    unsigned int cpu;
    unsigned int self_cpu = smp_processor_id();
#endif

    // TODO this may not be the best way to check if __irq_regs exists
#ifdef CONFIG_SMP
    for_each_present_cpu(cpu) {
        if (cpu == self_cpu && conf->coredump.regs != NULL) {
            crash_save_cpu(conf->coredump.regs, cpu);
        } else {
            crash_save_cpu(per_cpu(__irq_regs, cpu), cpu);
        }
    }
#else
    // Don't know how to save, so just save current context
    crash_save_cpu(conf->coredump.regs, 0);
#endif

    ret = prepare_elf32_headers(conf->coredump.elfhdr, conf->coredump.elfhdr_size);

    if (ret != 0) {
        pr_err(LOG_PREFIX "prepare_elf_header failed\n");
    }

    core_fixup_offsets(conf->coredump.elfhdr, conf->coredump.elfhdr_size);

    if (conf->family == AF_INET6) {
        pr_err(LOG_PREFIX "Sending elf core to %pI6c (port=%d, mac=%pM)...\n",
                conf->target_ip, conf->coredump_port, conf->target_mac);
    }
    else {
        pr_err(LOG_PREFIX "Sending elf core to %pI4 (port=%d, mac=%pM)...\n",
                conf->target_ip, conf->coredump_port, conf->target_mac);
    }

    // Make sure all ram content is current and present in the core file
    flush_cache_all();

    for (j = 0; j < 3; j++) {
        conf->coredump.block_num = 0;
        conf->coredump.send_bytes = 0;

        buf = conf->coredump.elfhdr;
        for (i = 0; i < conf->coredump.elfhdr_size; i += CHUNK_SIZE, buf += CHUNK_SIZE) {
            ret = send_buf(buf, CHUNK_SIZE);

            if (ret != 0) {
                goto fail;
            }
        }

        for (i = 0; i < conf->coredump.num_ram_regions; i++) {
            struct ram_region *reg = conf->coredump.ram_regions + i;

            addr = reg->start;
            for (; addr < reg->end; addr += CHUNK_SIZE) {
                ret = send_buf(__va(addr), CHUNK_SIZE);

                if (ret != 0) {
                    goto fail;
                }
            }
        }
    }

    pr_err(LOG_PREFIX "Send %llu bytes\n", conf->coredump.send_bytes);

    return 0;

fail:
    pr_err(LOG_PREFIX "Sending failed: %d\n", ret);
    return ret;
}

static struct notifier_block oops_block = {
    .notifier_call = avm_coredump_set_regs,
    // We need to be more important than the default die handler that
    // calls panic. Otherwise we wont be able to intercept the crashed
    // register set.
    .priority = 1,
};

static int count_ram_region(u64 start, u64 end, void *arg) {
    conf->coredump.num_ram_regions++;

    return 0;
}

static int prepare_ram_region(u64 start, u64 end, void* arg) {
    struct ram_region *reg = conf->coredump.ram_regions + conf->coredump.i_ram_region;

    reg->start = start;
    reg->end = end;

    // In the file, the ram regions will be right behind another without any holes
    if (conf->coredump.i_ram_region == 0) {
        reg->offset = 0;
    } else {
        struct ram_region *prev = conf->coredump.ram_regions + conf->coredump.i_ram_region - 1;
        reg->offset = prev->offset + prev->end - prev->start + 1;
    }
    pr_err(LOG_PREFIX "Prepare ram region 0x%08llx->0x%08llx (ofs=0x%08x)\n", reg->start, reg->end, reg->offset);

    conf->coredump.i_ram_region++;

    return 0;
}

int avm_coredump_coredump_init(void) {
    int ret;

    // Crash notes must be physically continous
    conf->coredump.crash_notes = kzalloc(NOTE_BYTES * num_possible_cpus(), GFP_KERNEL);
    if (conf->coredump.crash_notes == NULL) {
        pr_err(LOG_PREFIX "Could not allocate crash notes");
        return -ENOMEM;
    }

    ret = walk_system_ram_res(0, -1, NULL, count_ram_region);

    if (ret != 0) {
        goto fail;
    }

    conf->coredump.ram_regions = kzalloc(sizeof(struct ram_region) * conf->coredump.num_ram_regions, GFP_KERNEL);
    if (conf->coredump.ram_regions == NULL) {
        pr_err(LOG_PREFIX "Could not allocate ram region descriptions\n");
        return -ENOMEM;
    }

    ret = walk_system_ram_res(0, -1, NULL, prepare_ram_region);

    if (ret != 0) {
        goto fail;
    }

    conf->coredump.elfhdr_size = ALIGN(sizeof(Elf32_Ehdr) + 4 * sizeof(Elf32_Phdr), ELF_CORE_HEADER_ALIGN);
    conf->coredump.elfhdr = kzalloc(conf->coredump.elfhdr_size, GFP_KERNEL);
    if (conf->coredump.elfhdr == NULL) {
        pr_err(LOG_PREFIX "Could not allocate elf header\n");
        return -ENOMEM;
    }

    register_die_notifier(&oops_block);

fail:
    return ret;
}

void avm_coredump_coredump_exit(void) {
    kfree(conf->coredump.ram_regions);
    unregister_die_notifier(&oops_block);
    kfree(conf->coredump.crash_notes);
}