#include<iostream>
#include<stdio.h>

/*
 * Demonstrates an use after free c++ internally by deleting an object
 * and resetting its vtable pointing to legitimate functions
 */

// set pointer size and integer size dependent on architechture
#if __x86_64__
    /* 64-bit */
    const char* arch = "x86_64";
    int ptrSize = 8;
    typedef long long myInt;    // 8 byte
#else
    /* 32-bit */
    const char* arch = "x86";
    int ptrSize = 4;
    typedef int myInt;          // 4 byte
#endif
//

using namespace std;


// BASE CLASS
class Addition{
    protected:
        // member variables
        int s1,s2;
    public:
        // virtual member function will be implemented in derived class
        virtual int add(int a,int b){
        };
        int i;
};
//

// inherit new class
class Add: public Addition{
    public:
        // constructor:
        /*
        Add(int a, int b){
            s1 = a;
            s2 = b;
        };
        */
        void setvals(int a, int b){
            s1 = a;
            s2 = b;
        }
        // implement virtual function
        virtual int TRIGGER(){
            cout<<s1 + s2<<"\n";
        };
        int j;
};
//

// use this function in vtable which will be injected
void inject(){
    printf("%s\n","I've been executed");
}

int main(){

    printf("We're on a %s architecture\n\n",arch);

    // create object of inherited class
    //Add* A = new Add(4,5);
    Add* A = new Add;       //heap allocation happens here
    A->setvals(4,5);


    // add attributes
    A->i = 6;
    A->j = 7;

    // get vtable address
    myInt vtablePTR = NULL;
    vtablePTR = *((myInt*)A);

    printf("vtable: %.16x\n", vtablePTR);
    //function 1
    printf("%.16x : %.16x\n", vtablePTR, *((myInt*)vtablePTR));
    //function 2
    printf("%.16x : %.16x\n", vtablePTR + ptrSize, *((myInt*)(vtablePTR + ptrSize)));

    // call instantiated virtual function (legitimate)
    A->TRIGGER();

    // FREE OBJECT !
    delete A;
    //A = NULL;     // this prevents use after free

    // get objects address as it is still valid
    myInt* ObjPTR = (myInt*)A;

    // let object point to old vtable (inject old vtable)
    *ObjPTR = vtablePTR;

    // NOT legitimate
    A->TRIGGER();

    myInt vtable_func = (myInt)inject;

    // create new vtable
    myInt vtable_new[2] = {vtable_func, vtable_func};

    // let object point to new vtable (inject new vtable)
    *ObjPTR = (myInt)vtable_new;

    // USE AFTER FREE (execute injected function in vtable)
    A->TRIGGER();

    __asm__(
        "nop;"
        "nop;"
        "nop;"
    );


   // USE AFTER FREE with data members
    printf("%i %i\n", A->i, A->j);

}