/*

 *

 * Redistribution and use in source and binary forms, with or without

 * modification, are permitted provided that the following conditions

 * are met:

 *

 *   - Redistributions of source code must retain the above copyright

 *     notice, this list of conditions and the following disclaimer.

 *

 *   - Redistributions in binary form must reproduce the above copyright

 *     notice, this list of conditions and the following disclaimer in the

 *     documentation and/or other materials provided with the distribution.

 *

 *   - Neither the name of Oracle nor the names of its

 *     contributors may be used to endorse or promote products derived

 *     from this software without specific prior written permission.

 *

 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS

 * IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,

 * THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR

 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR

 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,

 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,

 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR

 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF

 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING

 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS

 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

 */

/*

 * This source code is provided to illustrate the usage of a given feature

 * or technique and has been deliberately simplified. Additional steps

 * required for a production-quality application, such as security checks,

 * input validation and proper error handling, might not be present in

 * this sample code.

 */

#include "stdlib.h"

#include "mtrace.h"

#include "java_crw_demo.h"

/* ------------------------------------------------------------------- */

/* Some constant maximum sizes */

#define MAX_TOKEN_LENGTH        16

#define MAX_THREAD_NAME_LENGTH  512

#define MAX_METHOD_NAME_LENGTH  1024

/* Some constant names that tie to Java class/method names.

 *    We assume the Java class whose static methods we will be calling

 *    looks like:

 *

 * public class Mtrace {

 *     private static int engaged;

 *     private static native void _method_entry(Object thr, int cnum, int mnum);

 *     public static void method_entry(int cnum, int mnum)

 *     {

 *         if ( engaged != 0 ) {

 *             _method_entry(Thread.currentThread(), cnum, mnum);

 *         }

 *     }

 *     private static native void _method_exit(Object thr, int cnum, int mnum);

 *     public static void method_exit(int cnum, int mnum)

 *     {

 *         if ( engaged != 0 ) {

 *             _method_exit(Thread.currentThread(), cnum, mnum);

 *         }

 *     }

 * }

 *

 *    The engaged field allows us to inject all classes (even system classes)

 *    and delay the actual calls to the native code until the VM has reached

 *    a safe time to call native methods (Past the JVMTI VM_START event).

 *

 */

#define MTRACE_class        Mtrace          /* Name of class we are using */

#define MTRACE_entry        method_entry    /* Name of java entry method */

#define MTRACE_exit         method_exit     /* Name of java exit method */

#define MTRACE_native_entry _method_entry   /* Name of java entry native */

#define MTRACE_native_exit  _method_exit    /* Name of java exit native */

#define MTRACE_engaged      engaged         /* Name of java static field */

/* C macros to create strings from tokens */

#define _STRING(s) #s

#define STRING(s) _STRING(s)

/* ------------------------------------------------------------------- */

/* Data structure to hold method and class information in agent */

typedef struct MethodInfo {

    const char *name;          /* Method name */

    const char *signature;     /* Method signature */

    int         calls;         /* Method call count */

    int         returns;       /* Method return count */

} MethodInfo;

typedef struct ClassInfo {

    const char *name;          /* Class name */

    int         mcount;        /* Method count */

    MethodInfo *methods;       /* Method information */

    int         calls;         /* Method call count for this class */

} ClassInfo;

/* Global agent data structure */

typedef struct {

    /* JVMTI Environment */

    jvmtiEnv      *jvmti;

    jboolean       vm_is_dead;

    jboolean       vm_is_started;

    /* Data access Lock */

    jrawMonitorID  lock;

    /* Options */

    char           *include;

    char           *exclude;

    int             max_count;

    /* ClassInfo Table */

    ClassInfo      *classes;

    jint            ccount;

} GlobalAgentData;

static GlobalAgentData *gdata;

/* Enter a critical section by doing a JVMTI Raw Monitor Enter */

static void

enter_critical_section(jvmtiEnv *jvmti)

{

    jvmtiError error;

    error = (*jvmti)->RawMonitorEnter(jvmti, gdata->lock);

    check_jvmti_error(jvmti, error, "Cannot enter with raw monitor");

}

/* Exit a critical section by doing a JVMTI Raw Monitor Exit */

static void

exit_critical_section(jvmtiEnv *jvmti)

{

    jvmtiError error;

    error = (*jvmti)->RawMonitorExit(jvmti, gdata->lock);

    check_jvmti_error(jvmti, error, "Cannot exit with raw monitor");

}

/* Get a name for a jthread */

static void

get_thread_name(jvmtiEnv *jvmti, jthread thread, char *tname, int maxlen)

{

    jvmtiThreadInfo info;

    jvmtiError      error;

    /* Make sure the stack variables are garbage free */

    (void)memset(&info,0, sizeof(info));

    /* Assume the name is unknown for now */

    (void)strcpy(tname, "Unknown");

    /* Get the thread information, which includes the name */

    error = (*jvmti)->GetThreadInfo(jvmti, thread, &info);

    check_jvmti_error(jvmti, error, "Cannot get thread info");

    /* The thread might not have a name, be careful here. */

    if ( info.name != NULL ) {

        int len;

        /* Copy the thread name into tname if it will fit */

        len = (int)strlen(info.name);

        if ( len < maxlen ) {

            (void)strcpy(tname, info.name);

        }

        /* Every string allocated by JVMTI needs to be freed */

        deallocate(jvmti, (void*)info.name);

    }

}

/* Qsort class compare routine */

static int

class_compar(const void *e1, const void *e2)

{

    ClassInfo *c1 = (ClassInfo*)e1;

    ClassInfo *c2 = (ClassInfo*)e2;

    if ( c1->calls > c2->calls ) return  1;

    if ( c1->calls < c2->calls ) return -1;

    return 0;

}

/* Qsort method compare routine */

static int

method_compar(const void *e1, const void *e2)

{

    MethodInfo *m1 = (MethodInfo*)e1;

    MethodInfo *m2 = (MethodInfo*)e2;

    if ( m1->calls > m2->calls ) return  1;

    if ( m1->calls < m2->calls ) return -1;

    return 0;

}

/* Callback from java_crw_demo() that gives us mnum mappings */

static void

mnum_callbacks(unsigned cnum, const char **names, const char**sigs, int mcount)

{

    ClassInfo  *cp;

    int         mnum;

    if ( cnum >= (unsigned)gdata->ccount ) {

        fatal_error("ERROR: Class number out of range\n");

    }

    if ( mcount == 0 ) {

        return;

    }

    cp           = gdata->classes + (int)cnum;

    cp->calls    = 0;

    cp->mcount   = mcount;

    cp->methods  = (MethodInfo*)calloc(mcount, sizeof(MethodInfo));

    if ( cp->methods == NULL ) {

        fatal_error("ERROR: Out of malloc memory\n");

    }

    for ( mnum = 0 ; mnum < mcount ; mnum++ ) {

        MethodInfo *mp;

        mp            = cp->methods + mnum;

        mp->name      = (const char *)strdup(names[mnum]);

        if ( mp->name == NULL ) {

            fatal_error("ERROR: Out of malloc memory\n");

        }

        mp->signature = (const char *)strdup(sigs[mnum]);

        if ( mp->signature == NULL ) {

            fatal_error("ERROR: Out of malloc memory\n");

        }

    }

}

/* Java Native Method for entry */

static void

MTRACE_native_entry(JNIEnv *env, jclass klass, jobject thread, jint cnum, jint mnum)

{

    enter_critical_section(gdata->jvmti); {

        /* It's possible we get here right after VmDeath event, be careful */

        if ( !gdata->vm_is_dead ) {

            ClassInfo  *cp;

            MethodInfo *mp;

            if ( cnum >= gdata->ccount ) {

                fatal_error("ERROR: Class number out of range\n");

            }

            cp = gdata->classes + cnum;

            if ( mnum >= cp->mcount ) {

                fatal_error("ERROR: Method number out of range\n");

            }

            mp = cp->methods + mnum;

            if ( interested((char*)cp->name, (char*)mp->name,

                            gdata->include, gdata->exclude)  ) {

                mp->calls++;

                cp->calls++;

            }

        }

    } exit_critical_section(gdata->jvmti);

}

/* Java Native Method for exit */

static void

MTRACE_native_exit(JNIEnv *env, jclass klass, jobject thread, jint cnum, jint mnum)

{

    enter_critical_section(gdata->jvmti); {

        /* It's possible we get here right after VmDeath event, be careful */

        if ( !gdata->vm_is_dead ) {

            ClassInfo  *cp;

            MethodInfo *mp;

            if ( cnum >= gdata->ccount ) {

                fatal_error("ERROR: Class number out of range\n");

            }

            cp = gdata->classes + cnum;

            if ( mnum >= cp->mcount ) {

                fatal_error("ERROR: Method number out of range\n");

            }

            mp = cp->methods + mnum;

            if ( interested((char*)cp->name, (char*)mp->name,

                            gdata->include, gdata->exclude)  ) {

                mp->returns++;

            }

        }

    } exit_critical_section(gdata->jvmti);

}

/* Callback for JVMTI_EVENT_VM_START */

static void JNICALL

cbVMStart(jvmtiEnv *jvmti, JNIEnv *env)

{

    enter_critical_section(jvmti); {

        jclass   klass;

        jfieldID field;

        int      rc;

        /* Java Native Methods for class */

        static JNINativeMethod registry[2] = {

            {STRING(MTRACE_native_entry), "(Ljava/lang/Object;II)V",

                (void*)&MTRACE_native_entry},

            {STRING(MTRACE_native_exit),  "(Ljava/lang/Object;II)V",

                (void*)&MTRACE_native_exit}

        };

        /* The VM has started. */

        stdout_message("VMStart\n");

        /* Register Natives for class whose methods we use */

        klass = (*env)->FindClass(env, STRING(MTRACE_class));

        if ( klass == NULL ) {

            fatal_error("ERROR: JNI: Cannot find %s with FindClass\n",

                        STRING(MTRACE_class));

        }

        rc = (*env)->RegisterNatives(env, klass, registry, 2);

        if ( rc != 0 ) {

            fatal_error("ERROR: JNI: Cannot register native methods for %s\n",

                        STRING(MTRACE_class));

        }

        /* Engage calls. */

        field = (*env)->GetStaticFieldID(env, klass, STRING(MTRACE_engaged), "I");

        if ( field == NULL ) {

            fatal_error("ERROR: JNI: Cannot get field from %s\n",

                        STRING(MTRACE_class));

        }

        (*env)->SetStaticIntField(env, klass, field, 1);

        /* Indicate VM has started */

        gdata->vm_is_started = JNI_TRUE;

    } exit_critical_section(jvmti);

}

/* Callback for JVMTI_EVENT_VM_INIT */

static void JNICALL

cbVMInit(jvmtiEnv *jvmti, JNIEnv *env, jthread thread)

{

    enter_critical_section(jvmti); {

        char  tname[MAX_THREAD_NAME_LENGTH];

        static jvmtiEvent events[] =

                { JVMTI_EVENT_THREAD_START, JVMTI_EVENT_THREAD_END };

        int        i;

        /* The VM has started. */

        get_thread_name(jvmti, thread, tname, sizeof(tname));

        stdout_message("VMInit %s\n", tname);

        /* The VM is now initialized, at this time we make our requests

         *   for additional events.

         */

        for( i=0; i < (int)(sizeof(events)/sizeof(jvmtiEvent)); i++) {

            jvmtiError error;

            /* Setup event  notification modes */

            error = (*jvmti)->SetEventNotificationMode(jvmti, JVMTI_ENABLE,

                                  events[i], (jthread)NULL);

            check_jvmti_error(jvmti, error, "Cannot set event notification");

        }

    } exit_critical_section(jvmti);

}

/* Callback for JVMTI_EVENT_VM_DEATH */

static void JNICALL

cbVMDeath(jvmtiEnv *jvmti, JNIEnv *env)

{

    enter_critical_section(jvmti); {

        jclass   klass;

        jfieldID field;

        /* The VM has died. */

        stdout_message("VMDeath\n");

        /* Disengage calls in MTRACE_class. */

        klass = (*env)->FindClass(env, STRING(MTRACE_class));

        if ( klass == NULL ) {

            fatal_error("ERROR: JNI: Cannot find %s with FindClass\n",

                        STRING(MTRACE_class));

        }

        field = (*env)->GetStaticFieldID(env, klass, STRING(MTRACE_engaged), "I");

        if ( field == NULL ) {

            fatal_error("ERROR: JNI: Cannot get field from %s\n",

                        STRING(MTRACE_class));

        }

        (*env)->SetStaticIntField(env, klass, field, 0);

        /* The critical section here is important to hold back the VM death

         *    until all other callbacks have completed.

         */

        /* Since this critical section could be holding up other threads

         *   in other event callbacks, we need to indicate that the VM is

         *   dead so that the other callbacks can short circuit their work.

         *   We don't expect any further events after VmDeath but we do need

         *   to be careful that existing threads might be in our own agent

         *   callback code.

         */

        gdata->vm_is_dead = JNI_TRUE;

        /* Dump out stats */

        stdout_message("Begin Class Stats\n");

        if ( gdata->ccount > 0 ) {

            int cnum;

            /* Sort table (in place) by number of method calls into class. */

            /*  Note: Do not use this table after this qsort! */

            qsort(gdata->classes, gdata->ccount, sizeof(ClassInfo),

                        &class_compar);

            /* Dump out gdata->max_count most called classes */

            for ( cnum=gdata->ccount-1 ;

                  cnum >= 0 && cnum >= gdata->ccount - gdata->max_count;

                  cnum-- ) {

                ClassInfo *cp;

                int        mnum;

                cp = gdata->classes + cnum;

                stdout_message("Class %s %d calls\n", cp->name, cp->calls);

                if ( cp->calls==0 ) continue;

                /* Sort method table (in place) by number of method calls. */

                /*  Note: Do not use this table after this qsort! */

                qsort(cp->methods, cp->mcount, sizeof(MethodInfo),

                            &method_compar);

                for ( mnum=cp->mcount-1 ; mnum >= 0 ; mnum-- ) {

                    MethodInfo *mp;

                    mp = cp->methods + mnum;

                    if ( mp->calls==0 ) continue;

                    stdout_message("\tMethod %s %s %d calls %d returns\n",

                        mp->name, mp->signature, mp->calls, mp->returns);

                }

            }

        }

        stdout_message("End Class Stats\n");

        (void)fflush(stdout);

    } exit_critical_section(jvmti);

}

/* Callback for JVMTI_EVENT_THREAD_START */

static void JNICALL

cbThreadStart(jvmtiEnv *jvmti, JNIEnv *env, jthread thread)

{

    enter_critical_section(jvmti); {

        /* It's possible we get here right after VmDeath event, be careful */

        if ( !gdata->vm_is_dead ) {

            char  tname[MAX_THREAD_NAME_LENGTH];

            get_thread_name(jvmti, thread, tname, sizeof(tname));

            stdout_message("ThreadStart %s\n", tname);

        }

    } exit_critical_section(jvmti);

}

/* Callback for JVMTI_EVENT_THREAD_END */

static void JNICALL

cbThreadEnd(jvmtiEnv *jvmti, JNIEnv *env, jthread thread)

{

    enter_critical_section(jvmti); {

        /* It's possible we get here right after VmDeath event, be careful */

        if ( !gdata->vm_is_dead ) {

            char  tname[MAX_THREAD_NAME_LENGTH];

            get_thread_name(jvmti, thread, tname, sizeof(tname));

            stdout_message("ThreadEnd %s\n", tname);

        }

    } exit_critical_section(jvmti);

}

/* Callback for JVMTI_EVENT_CLASS_FILE_LOAD_HOOK */

static void JNICALL

cbClassFileLoadHook(jvmtiEnv *jvmti, JNIEnv* env,

                jclass class_being_redefined, jobject loader,

                const char* name, jobject protection_domain,

                jint class_data_len, const unsigned char* class_data,

                jint* new_class_data_len, unsigned char** new_class_data)

{

    enter_critical_section(jvmti); {

        /* It's possible we get here right after VmDeath event, be careful */

        if ( !gdata->vm_is_dead ) {

            const char *classname;

            /* Name could be NULL */

            if ( name == NULL ) {

                classname = java_crw_demo_classname(class_data, class_data_len,

                        NULL);

                if ( classname == NULL ) {

                    fatal_error("ERROR: No classname inside classfile\n");

                }

            } else {

                classname = strdup(name);

                if ( classname == NULL ) {

                    fatal_error("ERROR: Out of malloc memory\n");

                }

            }

            *new_class_data_len = 0;

            *new_class_data     = NULL;

            /* The tracker class itself? */

            if ( interested((char*)classname, "", gdata->include, gdata->exclude)

                  &&  strcmp(classname, STRING(MTRACE_class)) != 0 ) {

                jint           cnum;

                int            system_class;

                unsigned char *new_image;

                long           new_length;

                ClassInfo     *cp;

                /* Get unique number for every class file image loaded */

                cnum = gdata->ccount++;

                /* Save away class information */

                if ( gdata->classes == NULL ) {

                    gdata->classes = (ClassInfo*)malloc(

                                gdata->ccount*sizeof(ClassInfo));

                } else {

                    gdata->classes = (ClassInfo*)