redirect.cpp

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

#include <vector>
#include <string>
#include <unordered_map>

#include <elf.h>

#include "redirect.h"

int redirect_function(void* orig, void* target){
    // the redirect needs 14 bytes. 6 bytes for the jump instruction, and 8 for the address (using a relative near jump would only take 5 bytes, but limits the jump distance)
    size_t patchsize = 14;
    size_t pagesize = getpagesize();
    void* addr = (void*)((size_t)orig - ((size_t)orig%pagesize));
    size_t len = (size_t)orig + patchsize - (size_t)addr;
    int ret = mprotect(addr, len, PROT_READ | PROT_WRITE);
    if(ret){
        printf("redirect: failed to make code writable\n");
        return -1;
    }

    // actually patch now
    // this is a "jmp qword ptr [rip+0]"
    uint8_t instr[] = { 0xFF, 0x25, 0x00, 0x00, 0x00, 0x00 };
    memcpy(orig, instr, sizeof(instr));
    *(void**)((size_t)orig+sizeof(instr)) = target; // where to jump to is stored right after the instruction
    ret = mprotect(addr, len, PROT_READ | PROT_EXEC);
    if(ret){
        printf("redirect: failed to make code executable again\n");
        return -1;
    }
    return 0;
}


typedef struct lib_context {
    int fd;
    Elf* elf;
    std::unordered_map<std::string, size_t> section_indices;
    std::unordered_map<std::string, size_t> func_addresses;

    void* text_base;    // base address of the text segment. Not automatically set!
} lib_context;

void set_text_base(lib_context* ctx, void* base){
    ctx->text_base = base;
}

size_t get_function_offset(lib_context* ctx, const char* name){
    if(ctx->func_addresses.count(name) == 0){
        printf("Can't find function %s\n", name);
        return 0;
    }
    return ctx->func_addresses[name];
}

int redirect_named_function(lib_context* ctx, const char* name, void* target){
    if(!ctx->text_base){
        printf("Can't redirect named function: .text base address not set!\n");
        return 1;
    }
    size_t offset = get_function_offset(ctx, name);
    if(!offset){
        return 1;
    }
    redirect_function((void*)((size_t)ctx->text_base + offset), target);
    return 0;
}

lib_context* open_library(const char* path){
    lib_context* ctx = new lib_context();
    ctx->text_base = NULL;
    if(!ctx){
        return NULL;
    }
    ctx->fd = open(path, O_RDONLY);
    if(ctx->fd == -1){
        perror("Failed to open file\n");
        delete ctx;
        return NULL;
    }
    elf_version(EV_CURRENT);

    ctx->elf = elf_begin(ctx->fd, ELF_C_READ, NULL);
    if(!ctx->elf){
        printf("Failed to read elf: %s\n", elf_errmsg(elf_errno()));
        delete ctx;
        return NULL;
    }

    size_t section_name_section_index;
    elf_getshdrstrndx(ctx->elf, &section_name_section_index);
    Elf_Scn* section_string_table = elf_getscn(ctx->elf, section_name_section_index);
    Elf_Data* section_names_data = elf_getdata(section_string_table, NULL);

    Elf_Scn* current_scn = NULL;
    do{
        current_scn = elf_nextscn(ctx->elf, current_scn);
        if(!current_scn)
            break;
        Elf64_Shdr* header = elf64_getshdr(current_scn);
        ctx->section_indices[(char*)section_names_data->d_buf + header->sh_name] = elf_ndxscn(current_scn);
        //printf("Section %d has name %d (%s)\n", elf_ndxscn(current_scn), header->sh_name, (char*)section_names_data->d_buf + header->sh_name);
    } while(current_scn);
    //printf("symtab: %d strtab: %d\n", ctx->section_indices[".symtab"], ctx->section_indices[".strtab"]);

    if(ctx->section_indices.count(".symtab") == 0 || ctx->section_indices.count(".strtab") == 0 || ctx->section_indices.count(".text") == 0){
        printf("missing either the .symtab, .text or .strtab section!\n");
        delete ctx;
        return NULL;
    }

    Elf_Scn* strtab = elf_getscn(ctx->elf, ctx->section_indices[".strtab"]);
    Elf_Data* strtab_data = elf_getdata(strtab, NULL);

    Elf_Scn* symtab = elf_getscn(ctx->elf, ctx->section_indices[".symtab"]);
    Elf_Data* symtab_data = elf_getdata(symtab, NULL);
    size_t text_idx = ctx->section_indices[".text"];

    for(size_t i = 0; i < symtab_data->d_size / sizeof(Elf64_Sym); i++){
        Elf64_Sym* sym = &((Elf64_Sym*)symtab_data->d_buf)[i];
        // I only need functions, other stuff doesn't matter;
        if(ELF64_ST_TYPE(sym->st_info) == STT_FUNC){
            if(sym->st_shndx == text_idx){
                char* name_ptr = (char*)strtab_data->d_buf + sym->st_name;
                ctx->func_addresses[std::string(name_ptr)] = sym->st_value;
                //printf("Symbol %s at 0x%lx\n", name_ptr, sym->st_value);
            }

        }
    }
    
    //printf("hid_init at 0x%x\n", ctx->func_addresses["hid_init"]);

    //printf("section names: %d\n", section_name_section_index);

    return ctx;
}