It's a demo of function inline hook for riscv64

hook_function_riscv.c

/*
*  This program is a demo of function inline hook for riscv64,
*  Please compile and test in riscv64, WITHOUT any compile optimization
*
*  Function sub will be hooked by hooked_sub, when invoke function sub,
*  hooked_sub will be invoked first, then it can decide whether to invoke the 
*  original sub or not.
*/

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

/*
* We want the function to be page aligned, required by mprotect.
* And we want the function length longer than 26 Byte
*/
__attribute__ ((optimize("align-functions"))) int sub(int a, int b)
{
        int aa = a;
        int bb = b;
        return aa - bb;
}

int hooked_sub(int a, int b)
{
        void *code_space_addr;
        __asm__ volatile ("la %0, origin_sub_code_space\n\t":"=r"(code_space_addr));
        int (*origin_sub)(int, int) = (int (*)(int, int))(code_space_addr);
        printf("We can view parameters in hooked_sub: %d, %d\n", a, b);
        printf("Now resume origin sub\n");
        return (*origin_sub)(a, b);
}

asm (
        "hook_trampoline:\n\t"
        "ld    t1, 8(sp)\n\t"
        "addi  sp, sp, 16\n\t"
        "j     hooked_sub\n\t"

        ".align 12\n\t"  // We want page aligned, also required by mprotect 
        "origin_sub_code_space:\n\t"

        ".word 0\n\t"
        ".word 0\n\t"
        ".word 0\n\t"
        ".word 0\n\t"
        ".word 0\n\t"
        ".word 0\n\t"
        ".word 0\n\t"

        ".word 0\n\t"
        ".word 0\n\t"
        ".word 0\n\t"
        ".word 0\n\t"
        ".word 0\n\t"
        ".word 0\n\t"
        ".word 0\n\t"
);

int disass_target(void *target)
{
        uint16_t insn_prob;
        int off = 0;

        while (off < 26) {
                insn_prob = *(uint16_t *)(target + off);
                off += ((insn_prob & 0x3) == 0x3 ? 4 : 2);
        }
        /* No other alternatives, must be a bug if any */
        assert(off == 26 || off == 28);
        return off;
}

void patch_function()
{
        int real_len;
        void *code_space_addr, *hook_trampoline_addr, *return_addr = &sub + 22;
        __asm__ volatile ("la %0, origin_sub_code_space\n\t":"=r"(code_space_addr));
        __asm__ volatile ("la %0, hook_trampoline\n\t":"=r"(hook_trampoline_addr));

/*
* Fill in the origin_sub_code_space in asm{}
* Long jmp code, 26B:
        "addi  sp, sp, -16"
        "sd    t1, 8(sp)"
        "auipc t1, 0"
        "ld    t1, 10(t1)"
        "jr    t1"
        ".word low"
        ".word high"
        "ld    t1, 8(sp)"
        "addi  sp, sp, 16"

* We will copy 28B from sub to origin_sub_code_space, we need to check
* if 28B is exactly the instruction boundary, or in the middle of the
* instruction, so disass_target() is used for this purpose.
*/
        memcpy(code_space_addr, &sub, 28);
        real_len = disass_target(code_space_addr);
        if (real_len == 26) {
                /* 
                * If 26B is the boundary, then we fill 27~28 to be NOP
                * instruction. If 28 is the boundary, we don't do anything.
                * Now 1~28B are all valid instructions.
                */
                *(char *)(code_space_addr + 26) = '\x01';
                *(char *)(code_space_addr + 27) = '\x00';
        }
        memcpy(code_space_addr + 28, "\x41\x11\x1a\xe4\x17\x03\x00\x00\x03\x33\xa3\x00\x02\x83"
                "\x00\x00\x00\x00\x00\x00\x00\x00\x22\x63\x41\x01", 26);
        memcpy(code_space_addr + 42, &return_addr, 8);
/*
* Patch the start of sub function
* Same long jmp code
*/
        memcpy(&sub, "\x41\x11\x1a\xe4\x17\x03\x00\x00\x03\x33\xa3\x00\x02\x83"
                "\x00\x00\x00\x00\x00\x00\x00\x00\x22\x63\x41\x01", 26);
        if (real_len != 26) {
                /*
                * If 28B is the boundary, 1~26 are overwritten by our long jump
                * code, which left 27~28 to be invalid instruction. So we
                * overwrite it to be NOP, to let 1~28B to be valid instructions.
                */
                *((char *)(&sub) + 26) = '\x01';
                *((char *)(&sub) + 27) = '\x00';
        }
        memcpy(&sub + 14, &hook_trampoline_addr, 8);
}

/*
* Before real patch, we need to change permissions of patched-areas 
*/
int remove_mem_protect()
{
        void *sub_code_space_addr = NULL;
        __asm__ volatile("la %0, origin_sub_code_space\n\t":"=r"(sub_code_space_addr));
        int ret1 = mprotect(sub_code_space_addr, 56, PROT_READ|PROT_WRITE|PROT_EXEC);
        int ret2 = mprotect(&sub, 28, PROT_READ|PROT_WRITE|PROT_EXEC);
        return (ret1|ret2);
}

int main()
{
        if (remove_mem_protect()) {
                printf("mprotect error!\n");
                return -1;
        }
        printf("***************************************************\n");
        printf("This is the result of unhooked sub(2, 100): %d\n", sub(2, 100));
        printf("This is the result of unhooked sub(50, 2): %d\n", sub(50, 2));
        printf("***************************************************\n");
        patch_function();
        printf("This is the result of hooked sub(2, 100): %d\n", sub(2, 100));
        printf("This is the result of hooked sub(50, 2): %d\n", sub(50, 2));
        printf("***************************************************\n");
        return 0;
}