#include <substrate.h>

#include <mach-o/dyld.h>

#include <mach/mach_init.h>
#include <mach/vm_map.h>

#include <sys/mman.h>
#include <unistd.h>

#include <CoreFoundation/CoreFoundation.h>
#include <Foundation/Foundation.h>

template <typename Type_>
static void nlset(Type_ &function, struct nlist *nl, size_t index) {
    struct nlist &name(nl[index]);
    uintptr_t value(name.n_value);
    if ((name.n_desc & N_ARM_THUMB_DEF) != 0)
        value |= 0x00000001;
    function = reinterpret_cast<Type_>(value);
}

static bool DidAddImage = false;

#define libCGFreeType "/System/Library/Frameworks/CoreGraphics.framework/Resources/libCGFreetype.A.dylib"
#define CGFreeType "/System/Library/Frameworks/CoreGraphics.framework/Resources/CGFreetype"

#define FailFast(message) do { \
    NSLog(@"CVE-2010-1797: %s", message); \
    return; \
} while (false)

static inline bool T$32bit$i(uint16_t ic) {
    return ((ic & 0xe000) == 0xe000 && (ic & 0x1800) != 0x0000);
}

#define T$b(im) /* b im */ \
    (0xe000 | (im & 0x7ff))
#define T$nop /* nop */ \
    (0x46c0)

// XXX: these are technically variables, but have not changed since /forever/
static const unsigned cff_op_random = 31;
static const unsigned cff_op_drop = 35;
static const unsigned cff_op_exch = 36;
static const unsigned cff_op_dup = 39;
static const unsigned cff_op_store = 42;

static unsigned r_args;
static unsigned r_op;
static const unsigned r_op_ = 1;
static unsigned r_return;

//static const unsigned s_stack = 5;
static unsigned s_decoder;
static const unsigned d_stack = 24;

static void *AfterLoop;
static void *DupCase;
static void *FailReturn;

__attribute__((__naked__))
void BeforeLoop() {
    __asm__ volatile ("sub sp, sp, #4");
    __asm__ volatile ("stmdb sp!, {sp}");
    __asm__ volatile ("stmdb sp!, {r0, r1, r2, r3, r4, r5, r6, lr}");
    //__asm__ volatile ("sub sp, sp, #4");

    volatile struct {
        uintptr_t r[7];
        void *lr;
        uintptr_t *sp;
        void *pc;
    } $;

    $.sp += 1;
    $.pc = $.r[r_op] + r_op_ == cff_op_random ? AfterLoop : DupCase;

    //if (($.r[r_args] - $.sp[s_stack]) / sizeof(uintptr_t) >= 48) {
    if (($.r[r_args] - $.sp[s_decoder]) / sizeof(uintptr_t) - d_stack >= 48) {
        $.r[r_return] = 0x82; // Stack_Overflow
        $.pc = FailReturn;
        NSLog(@"Prevented PDF Exploit");
    }

    /*if ($.r[r_op] + r_op_ != cff_op_random)
        NSLog(@"op=%u", $.r[r_op] + r_op_);*/

    //__asm__ volatile ("add sp, sp, #4");
    __asm__ volatile ("ldmia sp!, {r0, r1, r2, r3, r4, r5, r6, lr}");
    __asm__ volatile ("add sp, sp, #4");
    __asm__ volatile ("ldmia sp!, {pc}");
}

static void Simplify(uint16_t *&pc) {
    pc = reinterpret_cast<uint16_t *>(reinterpret_cast<uintptr_t>(pc) & ~0x1);
}

static void FixThumb(uint16_t *store, uint16_t *drop, uint16_t *exch, uint16_t *rand, uint16_t *dup) {
    //NSLog(@"store=%p, drop=%p, rand=%p, dup=%p", store, drop, rand, dup);

    Simplify(store);
    Simplify(drop);
    Simplify(exch);
    Simplify(rand);
    Simplify(dup);

    FailReturn = reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(store + 1) | 0x1);

    uint16_t mov(*store);
    if ((mov & 0xf800) != 0x2000)
        FailFast("failed to detect mov");
    r_return = (mov & 0x0700) >> 8;

    uint16_t ldrsp(*drop);
    if ((ldrsp & 0xf800) != 0x9800)
        FailFast("failed to detect ldrsp");
    s_decoder = ldrsp & 0x000f;

    void *begin(reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(rand) | 0x1));
    MSHookFunction(begin, reinterpret_cast<void *>(&BeforeLoop), &AfterLoop);

    uint16_t ldr(*exch);
    if ((ldr & 0xf800) != 0x6800)
        FailFast("failed to detect ldr");
    r_args = (ldr & 0x0038) >> 3;

    DupCase = reinterpret_cast<void *>(reinterpret_cast<uintptr_t>(dup) | 0x1);
}

static uint32_t *GetBranch(uint32_t *pc) {
    union {
        uint32_t value;

        struct {
            int32_t imm24 : 24;
            uint32_t l : 1;
            uint32_t : 3;
            uint32_t cond : 4;
        };
    } bits = {pc[0]};

    return pc + 2 + bits.imm24;
}

static uint16_t *GetBranch(uint16_t *pc) {
    union {
        uint16_t value;

        struct {
            uint16_t imm6 : 6;
            uint16_t cond : 4;
            uint16_t s : 1;
            uint16_t : 5;
        };
    } bits = {pc[0]};

    union {
        uint16_t value;

        struct {
            uint16_t imm11 : 11;
            uint16_t j2 : 1;
            uint16_t a : 1;
            uint16_t j1 : 1;
            uint16_t : 2;
        };
    } exts = {pc[1]};

    intptr_t jump(1);
    jump |= exts.imm11 << 1;
    jump |= bits.imm6 << 12;

    if (exts.a) {
        jump |= bits.s << 24;
        jump |= (~(bits.s ^ exts.j1) & 0x1) << 23;
        jump |= (~(bits.s ^ exts.j2) & 0x1) << 22;
        jump |= bits.cond << 18;
        jump <<= 7;
        jump >>= 7;
    } else {
        jump |= bits.s << 20;
        jump |= exts.j2 << 19;
        jump |= exts.j1 << 18;
        jump <<= 11;
        jump >>= 11;
    }

    return reinterpret_cast<uint16_t *>(reinterpret_cast<uintptr_t>(pc) + 4 + jump);
}

static void FixARM(uintptr_t address) {
    uint32_t *base(reinterpret_cast<uint32_t *>(address));
    for (unsigned index(0); index != 1024; ++index) {
        if ((base[index] & 0xfde00000) != 0x90800000)
            continue;
        //uint32_t *pc(base + index + 2);

        uint32_t cmp(base[index - 1]);
        if ((cmp & 0xfdf00000) != 0xe1500000)
            FailFast("failed to detect cmp");

        uint32_t imm(cmp & 0x000000ff);
        if (imm != 50)
            continue;

        r_op = (cmp & 0x000f0000) >> 16;

        uint32_t *start(base + index + 1);

        uint32_t *store(GetBranch(start + cff_op_store));
        uint32_t *drop(GetBranch(start + cff_op_drop));
        uint32_t *exch(GetBranch(start + cff_op_exch));
        uint32_t *rand(GetBranch(start + cff_op_random));
        uint32_t *dup(GetBranch(start + cff_op_dup));

        //NSLog(@"store=%p, drop=%p, exch=%p, rand=%p, dup=%p", store, drop, exch, rand, dup);

        FailReturn = store + 1;

        uint32_t mov(*store);
        if ((mov & 0xfde00000) != 0xe1a00000)
            FailFast("failed to detect mov");
        r_return = (mov & 0x0000f000) >> 12;

        uint32_t ldrsp(*drop);
        if ((ldrsp & 0xfc100000) != 0xe4100000)
            FailFast("failed to detect (ldr)sp");
        s_decoder = ldrsp & 0x0000000f;

        void *begin(reinterpret_cast<void *>(rand));
        MSHookFunction(begin, reinterpret_cast<void *>(&BeforeLoop), &AfterLoop);

        uint32_t ldr(*exch);
        if ((ldrsp & 0xfc100000) != 0xe4100000)
            FailFast("failed to detect ldr");
        r_args = (ldr & 0x0000f000) >> 12;

        MSHookFunction(dup, reinterpret_cast<void *>(&BeforeLoop), &DupCase);
    }
}

static void OnAddImage(const struct mach_header* mh, intptr_t vmaddr_slide) {
    if (DidAddImage)
        return;

    const char *library(NULL);

    uint32_t count(_dyld_image_count());
    for (uint32_t index(0); index != count; ++index)
        if (_dyld_get_image_header(index) == mh) {
            const char *name(_dyld_get_image_name(index));
            if (strcmp(name, libCGFreeType) == 0)
                library = libCGFreeType;
            if (strcmp(name, CGFreeType) == 0)
                library = CGFreeType;
        }

    if (library == NULL)
        return;

    DidAddImage = true;

    uint16_t *_cff_decoder_parse_charstrings;

    struct nlist nl[2];
    memset(nl, 0, sizeof(nl));
    nl[0].n_un.n_name = (char *) "_cff_decoder_parse_charstrings";
    nlist(library, nl);
    nlset(_cff_decoder_parse_charstrings, nl, 0);
    if (_cff_decoder_parse_charstrings == NULL)
        return;

    uintptr_t address(reinterpret_cast<uintptr_t>(_cff_decoder_parse_charstrings));
    if ((address & 0x1) == 0) {
        FixARM(address);
        return;
    }

    uint16_t *base(reinterpret_cast<uint16_t *>(address & ~0x1));

    mach_port_t self(mach_task_self());
    int page(getpagesize());

    for (unsigned index(0); index != 1024; ++index) {
        if ((base[index] & 0xfff0) != 0xe8d0)
            continue;
        uint16_t *pc(base + index + 2);

        bool wide(T$32bit$i(base[index - 2]));
        uint16_t cmp(base[index - (wide ? 3 : 2)]);
        if ((cmp & 0xf800) != 0x2800)
            FailFast("failed to detect cmp");

        uint16_t imm(cmp & 0x00ff);
        if (imm != 50)
            continue;

        r_op = (cmp & 0x0700) >> 8;

        uint16_t *start(reinterpret_cast<uint16_t *>(reinterpret_cast<uintptr_t>(pc) | 0x1));

        uint16_t *store(start + pc[cff_op_store - 1]);
        uint16_t *drop(start + pc[cff_op_drop - 1]);
        uint16_t *exch(start + pc[cff_op_exch - 1]);
        uint16_t *rand(start + pc[cff_op_random - 1]);
        uint16_t *dup(start + pc[cff_op_dup - 1]);

        FixThumb(store, drop, exch, rand, dup);

        uintptr_t address(reinterpret_cast<uintptr_t>(pc + cff_op_dup - 1));
        uintptr_t base(address / page * page);

        if (kern_return_t error = vm_protect(self, base, page, FALSE, VM_PROT_READ | VM_PROT_WRITE | VM_PROT_COPY))
            fprintf(stderr, "MS:Error:vm_protect():%d\n", error);
        else {
            pc[cff_op_dup - 1] = pc[cff_op_random - 1];
            if (kern_return_t error = vm_protect(self, base, page, FALSE, VM_PROT_READ | VM_PROT_EXECUTE | VM_PROT_COPY))
                fprintf(stderr, "MS:Error:vm_protect():%d\n", error);
        }

        return;
    }

    for (unsigned index(0); index != 1024; ++index) {
        if ((base[index] & 0xff00) != 0x4600)
            continue;
        if ((base[index] & 0x0007) != 0x0007)
            continue;

        uint16_t cmp;
        for (unsigned reverse(0); reverse != 7; ++reverse) {
            cmp = base[index - reverse - 1];
            if ((cmp & 0xf800) == 0x2800)
                goto cmp;
        }

        FailFast("failed to detect cmp");

      cmp:
        uint16_t imm(cmp & 0x00ff);
        if (imm != 50)
            continue;

        r_op = (cmp & 0x0700) >> 8;

        uint16_t *start(base + index - 1);

        uint16_t *store(GetBranch(start + cff_op_store * 2));
        uint16_t *drop(GetBranch(start + cff_op_drop * 2));
        uint16_t *exch(GetBranch(start + cff_op_exch * 2));
        uint16_t *rand(GetBranch(start + cff_op_random * 2));
        uint16_t *dup(GetBranch(start + cff_op_dup * 2));

        FixThumb(store, drop, exch, rand, dup);

        uintptr_t address(reinterpret_cast<uintptr_t>(start + cff_op_dup * 2));
        uintptr_t base(address / page * page);

        // XXX: this is wrong(ish)
        page *= 2;

        if (kern_return_t error = vm_protect(self, base, page, FALSE, VM_PROT_READ | VM_PROT_WRITE | VM_PROT_COPY))
            fprintf(stderr, "MS:Error:vm_protect():%d\n", error);
        else {
            start[cff_op_dup * 2] = T$b((cff_op_random - cff_op_dup - 1) * 2);
            start[cff_op_dup * 2 + 1] = T$nop;
            if (kern_return_t error = vm_protect(self, base, page, FALSE, VM_PROT_READ | VM_PROT_EXECUTE | VM_PROT_COPY))
                fprintf(stderr, "MS:Error:vm_protect():%d\n", error);
        }

        return;
    }
}

MSInitialize {
    _dyld_register_func_for_add_image(&OnAddImage);
}