// jfs_probe_ctl — userland helper for /dev/jfs_probe.
// Run as root.
//
// Usage:
//   ./jfs_probe_ctl mycred
//     - prints current task's cred kva, pfn, slab cache, size
//
//   ./jfs_probe_ctl kva2pfn 0xffff888100000000
//     - converts kva to pfn
//
//   ./jfs_probe_ctl pageinfo <pfn>
//     - dumps single-page info
//
//   ./jfs_probe_ctl adj <pfn>
//     - dumps adjacency: pfn-1, pfn, pfn+1
//
//   ./jfs_probe_ctl dump-cache <name>
//     - dumps all pages of named cache (e.g. cred_jar, jfs_ip,
//       kmalloc-cg-1k, vm_area_struct, filp); see dmesg

#define _GNU_SOURCE
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/types.h>
#include <unistd.h>

#define JFS_PROBE_IOC_MAGIC 'J'
#define JFS_PROBE_IOC_KVA_PFN     _IOWR(JFS_PROBE_IOC_MAGIC, 1, struct kva_pfn)
#define JFS_PROBE_IOC_PAGE_INFO   _IOWR(JFS_PROBE_IOC_MAGIC, 2, struct page_info)
#define JFS_PROBE_IOC_ADJACENCY   _IOWR(JFS_PROBE_IOC_MAGIC, 3, struct adjacency)
#define JFS_PROBE_IOC_MY_CRED     _IOR (JFS_PROBE_IOC_MAGIC, 4, struct my_cred)
#define JFS_PROBE_IOC_DUMP_CACHE  _IOW (JFS_PROBE_IOC_MAGIC, 5, char[64])
#define JFS_PROBE_IOC_WALK_FD_MAP _IOW (JFS_PROBE_IOC_MAGIC, 6, int)
#define JFS_PROBE_IOC_WALK_IPIMAP _IOW (JFS_PROBE_IOC_MAGIC, 7, int)
#define JFS_PROBE_IOC_VICTIMS_ALLOC _IOWR(JFS_PROBE_IOC_MAGIC, 8, struct victims_alloc)
#define JFS_PROBE_IOC_VICTIMS_CHECK _IOWR(JFS_PROBE_IOC_MAGIC, 9, struct victims_check)
#define JFS_PROBE_IOC_VICTIMS_FREE  _IO  (JFS_PROBE_IOC_MAGIC, 10)
#define JFS_PROBE_IOC_READ_BYTES    _IOWR(JFS_PROBE_IOC_MAGIC, 11, struct read_bytes)
#define JFS_PROBE_IOC_VICTIMS_PUNCH _IOWR(JFS_PROBE_IOC_MAGIC, 12, struct victims_punch)
struct victims_punch {
    uint32_t stride, freed_count, remaining, _pad;
};

#define VICTIM_SIZE 208
struct victims_alloc {
    uint32_t want_count, got_count, stride_pat, _pad;
    uint64_t first_kva, first_pfn, cache_ptr;
    uint32_t cache_size, _pad2;
};
struct victims_check {
    uint32_t max_report, smashed_count, reported_count, _pad;
    uint64_t report_buf;
};
struct victim_report {
    uint32_t slot_index, first_bad_offset;
    uint64_t victim_kva, victim_pfn;
    uint8_t  bytes[VICTIM_SIZE];
};
struct read_bytes {
    uint64_t kva;
    uint32_t len, ok;
    uint64_t user_buf;
};

struct kva_pfn { uint64_t kva; uint64_t pfn; int32_t ok; };
struct page_info {
    uint64_t pfn;
    char     kind[16];
    uint64_t cache_ptr;
    uint32_t cache_size;
    uint32_t migrate_type;
    uint64_t page_flags;
    int32_t  refcount;
    int32_t  mapcount;
    int32_t  in_slab;
    int32_t  in_lru;
    int32_t  ok;
};
struct adjacency {
    uint64_t pfn;
    struct page_info prev, self, next;
};
struct my_cred {
    uint64_t cred_kva, cred_pfn;
    uint64_t cred_cache_ptr;
    uint32_t cred_size, my_uid;
    uint64_t task_kva, mm_kva;
};

static const char *mt_str(int32_t mt)
{
    if (mt < 0) return "(unresolved)";
    switch (mt) {
    case 0: return "UNMOVABLE";
    case 1: return "MOVABLE";
    case 2: return "RECLAIMABLE";
    case 3: return "HIGHATOMIC";
    case 4: return "CMA";
    case 5: return "ISOLATE";
    default: return "?";
    }
}

static void print_pi(const char *tag, struct page_info *pi)
{
    if (!pi->ok) {
        printf("  %s: pfn=0x%" PRIx64 " INVALID\n", tag, pi->pfn);
        return;
    }
    if (pi->in_slab) {
        printf("  %s: pfn=0x%-8" PRIx64 " kind=slab     cache=0x%" PRIx64 " sz=%-5u mt=%-12s refc=%d\n",
               tag, pi->pfn, pi->cache_ptr, pi->cache_size, mt_str(pi->migrate_type), pi->refcount);
    } else {
        printf("  %s: pfn=0x%-8" PRIx64 " kind=%-9s                          mt=%-12s refc=%d mapc=%d lru=%d\n",
               tag, pi->pfn, pi->kind, mt_str(pi->migrate_type), pi->refcount, pi->mapcount, pi->in_lru);
    }
}

int main(int argc, char **argv)
{
    int fd, r;
    if (argc < 2) {
        fprintf(stderr,
            "usage: %s {mycred | kva2pfn <kva> | pageinfo <pfn> | adj <pfn> | "
            "dump-cache <name> | victims-alloc <N> | victims-check [maxreport] | "
            "victims-free | read <kva> <len>}\n",
            argv[0]);
        return 1;
    }
    fd = open("/dev/jfs_probe", O_RDWR);
    if (fd < 0) { perror("open /dev/jfs_probe"); return 1; }

    if (!strcmp(argv[1], "mycred")) {
        struct my_cred mc;
        memset(&mc, 0, sizeof(mc));
        if ((r = ioctl(fd, JFS_PROBE_IOC_MY_CRED, &mc)) < 0) { perror("ioctl"); return 1; }
        printf("uid=%u\n", mc.my_uid);
        printf("task_kva       = 0x%" PRIx64 "\n", mc.task_kva);
        printf("mm_kva         = 0x%" PRIx64 "\n", mc.mm_kva);
        printf("cred_kva       = 0x%" PRIx64 "\n", mc.cred_kva);
        printf("cred_pfn       = 0x%" PRIx64 "\n", mc.cred_pfn);
        printf("cred_size      = %u\n", mc.cred_size);
        printf("cred_cache_ptr = 0x%" PRIx64 "  (match in /proc/slabinfo)\n", mc.cred_cache_ptr);
        return 0;
    }
    if (!strcmp(argv[1], "kva2pfn")) {
        if (argc < 3) { fprintf(stderr, "need kva\n"); return 1; }
        struct kva_pfn k = { .kva = strtoull(argv[2], NULL, 0) };
        if ((r = ioctl(fd, JFS_PROBE_IOC_KVA_PFN, &k)) < 0) { perror("ioctl"); return 1; }
        printf("kva 0x%" PRIx64 " -> pfn 0x%" PRIx64 " (ok=%d)\n", k.kva, k.pfn, k.ok);
        return 0;
    }
    if (!strcmp(argv[1], "pageinfo")) {
        if (argc < 3) { fprintf(stderr, "need pfn\n"); return 1; }
        struct page_info pi = { .pfn = strtoull(argv[2], NULL, 0) };
        if ((r = ioctl(fd, JFS_PROBE_IOC_PAGE_INFO, &pi)) < 0) { perror("ioctl"); return 1; }
        print_pi("self", &pi);
        return 0;
    }
    if (!strcmp(argv[1], "adj")) {
        if (argc < 3) { fprintf(stderr, "need pfn\n"); return 1; }
        struct adjacency a = { .pfn = strtoull(argv[2], NULL, 0) };
        if ((r = ioctl(fd, JFS_PROBE_IOC_ADJACENCY, &a)) < 0) { perror("ioctl"); return 1; }
        printf("adjacency around pfn 0x%" PRIx64 ":\n", a.pfn);
        print_pi("prev", &a.prev);
        print_pi("self", &a.self);
        print_pi("next", &a.next);
        return 0;
    }
    if (!strcmp(argv[1], "dump-cache")) {
        if (argc < 3) { fprintf(stderr, "need cache pointer (hex)\n"); return 1; }
        uint64_t target = strtoull(argv[2], NULL, 0);
        if ((r = ioctl(fd, JFS_PROBE_IOC_DUMP_CACHE, &target)) < 0) { perror("ioctl"); return 1; }
        printf("dump-cache 0x%" PRIx64 " -> see dmesg\n", target);
        return 0;
    }
    if (!strcmp(argv[1], "walk")) {
        if (argc < 3) { fprintf(stderr, "need path\n"); return 1; }
        int tfd = open(argv[2], O_RDONLY);
        if (tfd < 0) { perror("open target"); return 1; }
        if ((r = ioctl(fd, JFS_PROBE_IOC_WALK_FD_MAP, &tfd)) < 0) { perror("ioctl"); close(tfd); return 1; }
        printf("walk %s -> %d folios (see dmesg)\n", argv[2], r);
        close(tfd);
        return 0;
    }
    if (!strcmp(argv[1], "walk-ipimap")) {
        // path of any file/dir in JFS mount; module walks the imap inode's mapping
        if (argc < 3) { fprintf(stderr, "need path inside JFS mount\n"); return 1; }
        int tfd = open(argv[2], O_RDONLY);
        if (tfd < 0) { perror("open target"); return 1; }
        if ((r = ioctl(fd, JFS_PROBE_IOC_WALK_IPIMAP, &tfd)) < 0) { perror("ioctl"); close(tfd); return 1; }
        printf("walk-ipimap via %s -> %d folios (see dmesg)\n", argv[2], r);
        close(tfd);
        return 0;
    }
    if (!strcmp(argv[1], "victims-alloc")) {
        if (argc < 3) { fprintf(stderr, "need N\n"); return 1; }
        struct victims_alloc va = { .want_count = (uint32_t)atoi(argv[2]) };
        if (argc >= 4) va.stride_pat = (uint32_t)strtoul(argv[3], NULL, 0);
        if ((r = ioctl(fd, JFS_PROBE_IOC_VICTIMS_ALLOC, &va)) < 0) { perror("ioctl"); return 1; }
        printf("got=%u/%u first_kva=0x%" PRIx64 " first_pfn=0x%" PRIx64 " cache=0x%" PRIx64 " size=%u\n",
               va.got_count, va.want_count, va.first_kva, va.first_pfn, va.cache_ptr, va.cache_size);
        return 0;
    }
    if (!strcmp(argv[1], "victims-check")) {
        uint32_t maxr = argc >= 3 ? (uint32_t)atoi(argv[2]) : 16;
        struct victim_report *reports = calloc(maxr, sizeof(*reports));
        if (!reports) { perror("calloc"); return 1; }
        struct victims_check vc = { .max_report = maxr,
                                    .report_buf = (uint64_t)(uintptr_t)reports };
        if ((r = ioctl(fd, JFS_PROBE_IOC_VICTIMS_CHECK, &vc)) < 0) { perror("ioctl"); return 1; }
        printf("smashed=%u reported=%u\n", vc.smashed_count, vc.reported_count);
        for (uint32_t i = 0; i < vc.reported_count; i++) {
            struct victim_report *vr = &reports[i];
            printf("slot=%u first_bad=%u kva=0x%" PRIx64 " pfn=0x%" PRIx64 "\n",
                   vr->slot_index, vr->first_bad_offset,
                   vr->victim_kva, vr->victim_pfn);
            for (int row = 0; row < VICTIM_SIZE; row += 16) {
                printf("  +0x%03x:", row);
                for (int j = 0; j < 16; j++) printf(" %02x", vr->bytes[row + j]);
                printf("\n");
            }
        }
        free(reports);
        return 0;
    }
    if (!strcmp(argv[1], "victims-free")) {
        if ((r = ioctl(fd, JFS_PROBE_IOC_VICTIMS_FREE)) < 0) { perror("ioctl"); return 1; }
        printf("victims freed\n");
        return 0;
    }
    if (!strcmp(argv[1], "victims-punch")) {
        if (argc < 3) { fprintf(stderr, "need stride\n"); return 1; }
        struct victims_punch vp = { .stride = (uint32_t)atoi(argv[2]) };
        if ((r = ioctl(fd, JFS_PROBE_IOC_VICTIMS_PUNCH, &vp)) < 0) { perror("ioctl"); return 1; }
        printf("punched stride=%u freed=%u remaining=%u\n",
               vp.stride, vp.freed_count, vp.remaining);
        return 0;
    }
    if (!strcmp(argv[1], "read")) {
        if (argc < 4) { fprintf(stderr, "need kva and len\n"); return 1; }
        uint64_t kva = strtoull(argv[2], NULL, 0);
        uint32_t len = (uint32_t)strtoul(argv[3], NULL, 0);
        uint8_t *buf = calloc(len ? len : 1, 1);
        if (!buf) { perror("calloc"); return 1; }
        struct read_bytes rb = { .kva = kva, .len = len,
                                 .user_buf = (uint64_t)(uintptr_t)buf };
        if ((r = ioctl(fd, JFS_PROBE_IOC_READ_BYTES, &rb)) < 0) { perror("ioctl"); return 1; }
        if (!rb.ok) { fprintf(stderr, "read invalid (kva not virt_addr_valid)\n"); return 1; }
        for (uint32_t i = 0; i < len; i += 16) {
            printf("  +0x%03x:", i);
            for (uint32_t j = 0; j < 16 && i + j < len; j++) printf(" %02x", buf[i + j]);
            printf("\n");
        }
        free(buf);
        return 0;
    }
    fprintf(stderr, "unknown subcmd %s\n", argv[1]);
    return 1;
}