/*
 * Copyright (c) 2003-2004 Amit Singh. All Rights Reserved.
 * http://www.osxbook.com
 * http://www.kernelthread.com
 */

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

#define BRANCH_MOPCODE 0x12 /* Branch instruction's major opcode        */
#define DEFAULT_STUBSZ 128  /* large enough size for a function stub    */
#define DEFAULT_FRR_BYTES 4 /* first n bytes of the "original" function */

/*
 * We will be using one of the PowerPC bx (Branch) family of instructions,
 * specifically "b" - unconditional branch.
 * 
 * b    target_addr (AA = 0 LK = 0)
 * ba   target_addr (AA = 1 LK = 0)
 * bl   target_addr (AA = 0 LK = 1)
 * bla  target_addr (AA = 1 LK = 1)
 *
 * 0     5 6                                                          29 30 31
 * +------+-------------------------------------------------------------+-+-+
 * |______|_____________________________________________________________|_|_|
 *   OP                                 LI                              AA LK
 *
 *     target_addr specifies the branch target address
 *
 * If AA = 0, then the branch target address is the sum of LI || 0b00
 * sign-extended and the address of this instruction.
 *
 * If AA = 1, then the branch target address is the value LI || 0b00
 * sign-extended.
 *
 * If LK = 1, then the effective address of the instruction following
 * the branch instruction is placed into the link register.
 */
typedef struct branch_s {
    unsigned int OP: 6;  /* bits 0 - 5, primary opcode    */
    unsigned int LI: 24; /* bits 6 - 29, LI               */
    unsigned int AA: 1;  /* bit 30, absolute address      */
    unsigned int LK: 1;  /* bit 31, link or not           */
} branch_t;

/* 
 * Each instance of re-routing has the following data structure associated
 * with it. A pointer to a frr_data_t is returned by the "install"
 * function. The "remove" function takes the same pointer as argument.
 */
typedef struct frr_data_s {
    void *f_orig; /* "original" function                  */
    void *f_new;  /* user-provided "new" function         */
    void *f_stub; /* stub to call "original" inside "new" */
    char f_bytes[DEFAULT_FRR_BYTES]; /* bytes from f_orig */
} frr_data_t;

/*
 * Given an "original" function and a "new" function, frr_install()
 * re-routes so that anybody calling "original" will actually be
 * calling "new". Inside "new", it is possible to call "original"
 * through a stub.
 */
frr_data_t *
frr_install(void *original, void *new)
{
    int ret = -1;
    branch_t branch;
    frr_data_t *FRR = (frr_data_t *)0;

    /* allocate space for meta-data */
    FRR = (frr_data_t *)malloc(sizeof(frr_data_t));
    if (!FRR)
        return FRR;

    FRR->f_orig = original;
    FRR->f_new = new;

    /* allocate space for the stub to call the original function */
    FRR->f_stub = (char *)malloc(DEFAULT_STUBSZ);
    if (!FRR->f_stub) {
        free(FRR);
        FRR = (frr_data_t *)0;
        return FRR;
    }

    /* Prepare to write to the first 4 bytes of "original" */
    ret = mprotect(FRR->f_orig, 4, PROT_READ|PROT_WRITE|PROT_EXEC);
    if (ret != 0)
        goto ERROR;

    /* Prepare to populate the stub and make it executable */
    ret = mprotect(FRR->f_stub, DEFAULT_STUBSZ, PROT_READ|PROT_WRITE|PROT_EXEC);
    if (ret != 0)
        goto ERROR;

    memcpy(FRR->f_bytes, (char *)FRR->f_orig, DEFAULT_FRR_BYTES);

    /* Create unconditional branch from "original" to "new" */
    branch.OP = BRANCH_MOPCODE;
    branch.LI = (unsigned int)FRR->f_new >> 2;
    branch.AA = 1;
    branch.LK = 0;

    memcpy((char *)FRR->f_orig, (char *)&branch, 4);

    memcpy((char *)FRR->f_stub, (char *)FRR->f_bytes, DEFAULT_FRR_BYTES);

    /* Create unconditional branch from "stub" to "original" */
    branch.LI = (unsigned int)(FRR->f_orig + 4) >> 2;

    memcpy((char *)FRR->f_stub + DEFAULT_FRR_BYTES, (char *)&branch, 4);

    /*
     * Must flush cache(s). For demonstration purposes, this suffices.
     */
    msync(FRR->f_stub, 4096, MS_INVALIDATE);

    return FRR;

    ERROR:
    if (FRR && FRR->f_stub)
        free(FRR->f_stub);
    if (FRR)
        free(FRR);
    return FRR;
}

int
frr_remove(frr_data_t *FRR)
{
    int ret;

    if (!FRR)
        return 0;

    ret = mprotect(FRR->f_orig, 4, PROT_READ|PROT_WRITE|PROT_EXEC);
    if (ret != 0)
        return -1;

    memcpy((char *)FRR->f_orig, FRR->f_bytes, DEFAULT_FRR_BYTES);

    if (FRR && FRR->f_stub)
        free(FRR->f_stub);

    if (FRR)
        free(FRR);

    FRR = (frr_data_t *)0;

    /*
     * Perhaps use dcbf/icbf to flush cache(s). For demonstration purposes,
     * calling a random function suffices. I don't think this is foolproof.
     */
    sync();

    return 0;
}

/* EXAMPLE USAGE */
int
function(int i, char *s)
{
    int ret;
    char *m = s;
    if (!s)
        m = "(null)";

    printf("function()    : called as function(%d, %s).\n", i, m);
    ret = i + 1;
    printf("function()    : returning %d.\n", ret);

    return ret;
}

int (* _function)(int, char *);

int
function_new(int i, char *s)
{
    int ret = -1;
    char *m = s;
    if (!s)
        m = "(null)";

    printf("function_new(): called as function_new(%d, %s).\n", i, m);

    if (_function) {
        printf("function_new(): calling _function(%d, %s).\n", i, m);
        ret = _function(i, s);
        printf("function_new(): _function(%d, %s) returned %d.\n", i, m, ret);
    } else {
        printf("function_new(): _function missing.\n");
    }

    printf("function_new(): returning %d.\n", ret);

    return ret;
}

int
main(int argc, char **argv)
{
    int ret;
    int arg_i = 2;
    char *arg_s = "Hello, World!";
    frr_data_t *FRR;

    int(*_function)(int, char*) = 0;

    printf("[Clean]\n");
    printf("main()        : calling function(%d, %s).\n", arg_i, arg_s);
    ret = function(arg_i, arg_s);
    printf("main()        : function(%d, %s) returned %d.\n",
           arg_i, arg_s, ret);

    FRR = frr_install(function, function_new);
    if (FRR)
        _function = FRR->f_stub;
    else {
        fprintf(stderr, "main(): frr_install failed.\n");
        return 1;
    }

    printf("\n[Rerouting installed]\n");
    printf("main()        : calling function(%d, %s).\n", arg_i, arg_s);
    ret = function(arg_i, arg_s);
    printf("main()        : function(%d, %s) returned %d.\n",
           arg_i, arg_s, ret);

    ret = frr_remove(FRR);
    if (ret != 0) {
        fprintf(stderr, "main(): frr_remove failed.\n");
        return 1;
    }

    printf("\n[Rerouting removed]\n");
    printf("main()        : calling function(%d, %s).\n", arg_i, arg_s);
    ret = function(arg_i, arg_s);
    printf("main()        : function(%d, %s) returned %d.\n",
           arg_i, arg_s, ret);

    return 0;
}

/* END */