#
# Copyright (C) 2023-2023 Intel Corporation.
# SPDX-License-Identifier: MIT
#

import pin

output_file = open("script_malloctrace_probe.out", "w")

# Note: MallocAfter and Arg1Before functions are not defined here,
# Client must either register them as remote functions or deliver their definitions through scripts fragments

def image_instrumentation_cb(img):
    mallocRtn = pin.RTN_FindByName(img, "malloc")
    if(pin.RTN_Valid(mallocRtn)):
        assert(pin.RTN_IsSafeForProbedInsertion(mallocRtn))
        pin.RTN_InsertCallProbed(mallocRtn, pin.IPOINT_BEFORE, Arg1Before, pin.IARG_PYOBJ, "malloc", pin.IARG_FUNCARG_ENTRYPOINT_VALUE, 0)
        
        proto_malloc_after = pin.PROTO_Allocate(pin.PIN_PARG_POINTER, pin.CALLINGSTD_DEFAULT, "malloc", pin.PIN_PARG_SIZE_T)
        pin.RTN_InsertCallProbed(mallocRtn, pin.IPOINT_AFTER, MallocAfter, pin.IARG_PROTOTYPE, proto_malloc_after, pin.IARG_FUNCRET_EXITPOINT_VALUE)
        pin.PROTO_Free(proto_malloc_after)
	
    freeRtn = pin.RTN_FindByName(img, "free")
    if (pin.RTN_Valid(freeRtn)):
        assert(pin.RTN_IsSafeForProbedInsertion(freeRtn))
        pin.RTN_InsertCallProbed(freeRtn, pin.IPOINT_BEFORE, Arg1Before, pin.IARG_PYOBJ, "free", pin.IARG_FUNCARG_ENTRYPOINT_VALUE, 0)

pin.IMG_AddInstrumentFunction(image_instrumentation_cb)
