/*JVM Thread Monitor (TMon)
 *@version:0.1
 *@date:Aug 10, 2006
 *@author:Li Ling
 * from https://www.ibm.com/developerworks/cn/java/j-lo-jvmti/
 * gcc TMon.c -I/usr/lib/jvm/java-8-oracle/include -I/usr/lib/jvm/java-8-oracle/include/linux -shared -fPIC -o ./libtestagent.so
 * clang++ TMon.cpp -I/usr/lib/jvm/java-8-oracle/include -I/usr/lib/jvm/java-8-oracle/include/linux -shared -fPIC -o ./libtestagent.so
 */



#include <string.h>
#include "jvmti.h"
#include <iostream>
#include <stdlib.h>

//#define DEBUG

#define checkError(error,msg) if(error!= JVMTI_ERROR_NONE){printf("[TMon] ERROR: can't %s, error code:%d\n",msg,error);exit(-1);}//return;}
#define checkNull(error,msg) if(error==NULL){printf("[TMon] FATAL ERROR: %s!",msg);exit(-1);}
#define debug(msg) printf("[TMon]: %s!\n",msg);
#define debugDetail(msg,code) printf("[TMon]: %s: %d!\n",msg,code);

#define THREAD_RESERVE 0
#define THREAD_NOTIFY 1
#define THREAD_WAIT 2
#define THREAD_START 3
#define THREAD_STOP 4
#define THREAD_SLEEP 5
#define THREAD_INTERRUPT 6
#define THREAD_JOIN 7
#define THREAD_RESUME 8
#define THREAD_SUSPEND 9
#define THREAD_YIELD 10

static jvmtiEnv *gb_jvmti = NULL;
static jvmtiCapabilities gb_capa;
static jrawMonitorID gb_lock;
static jvmtiError error;

//save the thread jump context, such as (thread A, start, thread B)
typedef struct thrContext
{
	jint caller;
	jint called;
	int type;
	char *action;
	thrContext *next;
}THR_CONTEXT;

//save the monitor info in each thread
typedef struct thrInfo
{
	jvmtiMonitorUsage *info_ptr;
	thrInfo *next;
}THR_INFO;

//head link
THR_CONTEXT *ctx_head=NULL,*ctx_tail=NULL;
THR_INFO *inf_head=NULL,*inf_tail=NULL;

int printLink(THR_CONTEXT *head)
{
	printf("\n*****************[TMon Status Dump]****************\n");
	for(;head!=NULL;)
	{
		printf("[TMon]: thread (%10d) %10s (%10d)\n",head->caller,head->action,head->called);
		head=head->next;
	}
	printf("*****************[TMon Status Dump]****************\n\n");
	return 0;
}

//insert elem to thread context link
int insertElem(THR_CONTEXT **head,THR_CONTEXT **tail,int caller,int called,char *action)
{
	THR_CONTEXT *current=NULL;
	current=(THR_CONTEXT*)malloc(sizeof(THR_CONTEXT));
	checkNull(current,"can't alloc mem");
	current->called=called;
	current->caller=caller;
	current->type=THREAD_RESERVE;
	current->action=action;
	current->next=NULL;

	if(*head==NULL)
	{
		*head=*tail=current;
	}
	else
	{
		(*tail)->next=current;
		*tail=(*tail)->next;
	}
	return 0;
}

//insert elem to thread monitor info link
int insertElem(THR_INFO **head,THR_INFO **tail,jvmtiMonitorUsage *info_ptr)
{
	THR_INFO *current=NULL;
	current=(THR_INFO*)malloc(sizeof(THR_INFO));
	checkNull(current,"can't alloc mem");
	current->info_ptr=info_ptr;
	current->next=NULL;
	if(*head==NULL)
	{
		*head=*tail=current;
	}
	else
	{
		(*tail)->next=current;
		*tail=(*tail)->next;
	}
	return 0;
}

int removeElem(THR_INFO **head,THR_INFO **tail)
{
	THR_INFO *current=NULL;
	if(*head==NULL)
	{
		return 0;
	}
	if(*head==*tail)
	{
		free(*head);
		*head=NULL;
		*tail=NULL;	
	}
	else
	{
		current=*head;
		*head=(*head)->next;
		*tail=NULL;
		free(current);
	}
	return 0;
}

//create synchronized region for each callback
static void enter_critical_section()
{
    error = gb_jvmti->RawMonitorEnter(gb_lock);
	checkError(error,"RawMonitorEnter");
}

//exit synchronized region
static void exit_critical_section()
{
    error = gb_jvmti->RawMonitorExit(gb_lock);
	checkError(error,"RawMonitorExit");
}

/*	uidThrCtx[][]
	usageFlag,uid1,thrID1,thrID2...
	usageFlag,uid2,thrID1,thrID2...
	...
  */
#define MaxUidNum 30
#define MaxThrNum 60
jint uidThrCtx[MaxUidNum][MaxThrNum];

void initUidThrCtx()
{
	for(int i=0;i<MaxUidNum;i++)
		for(int j=0;j<MaxThrNum;j++)
			uidThrCtx[i][j]=0;
}

void printUidThrCtx()
{
	printf("\n***************[TMon Demo Status Dump]****************\n");
	for(int i=0;i<MaxUidNum;i++)
	{
		for(int j=0;j<MaxThrNum;j++)
		{
			printf("%d,",uidThrCtx[i][j]);
		}
		printf("\n");
	}
}

//clear the thd link, if repeated uid accur
void clearLink(jint uid)
{
	for(int j=0;j<MaxUidNum;j++)
	{
		if(uidThrCtx[j][1]==uid)
			for(int i=2;i<MaxThrNum;i++)
			{
				uidThrCtx[j][i]=0;
			}
	}
}

//set the uidThr link expired, because of repeated tid
void setFlag(/*jint uid,*/jint thr)
{
	for(int i=0;i<MaxUidNum;i++)
	{
		if(uidThrCtx[i][0]==0/*&&uidThrCtx[i][1]==uid*/&&uidThrCtx[i][2]==thr)
			uidThrCtx[i][0]=1;//expired
	}
}

//remove the thr existed in uid link previously
void removeThr(jint thr)
{
	for(int i=0;i<MaxUidNum;i++)
		for(int j=2;j<MaxThrNum;j++)
			if(uidThrCtx[i][j]==thr)
				uidThrCtx[i][j]=0;
}

//add thr to uidThrCtx according to parent thr, if there is a thread jumping
void addThr(jint parentThr,jint childThr)
{
	//remove the old thr uid item
	removeThr(childThr);

	for(int i=0;i<MaxUidNum;i++)
		for(int j=2;j<MaxThrNum;j++)
		{
			if(uidThrCtx[i][j]==parentThr)
			{
				//maybe can insert directly
				int position=0;
				for(int k=2;k<MaxThrNum;k++)
				{
					if(uidThrCtx[i][k]==childThr)
					{
						return;
					}
					else if(position==0&&uidThrCtx[i][k]==0)
					{
						position=k;
					}
				}
				uidThrCtx[i][position]=childThr;
				return;
			}
		}
}

//find uid according to thr
jint findUid(jint thr)
{
	for(int i=0;i<MaxUidNum;i++)
		for(int j=2;j<MaxThrNum;j++)
			if(uidThrCtx[i][0]==0&&uidThrCtx[i][j]==thr)
				return uidThrCtx[i][1];
	return 0;
}

//get uid from jvm and insert into uidThrCtx
//bind user and tid context
void getUid(jthread thr,char *method,jint thrCode)
{
	if(strcmp(method,"beforeAdvice")!=0)
		return;

	//if thr has existed
	removeThr(thrCode);

	//if uidThr expired
	//setFlag(thrCode);

	jint uid=0;
	error=gb_jvmti->GetLocalInt(thr,0,1,&uid);
	checkError(error,"GetLocalObjectUid");

#ifdef DEBUG
	printf("[TMon]: get uid:%d!\n",uid);
#endif

	//reuse the old uid link
	clearLink(uid);
	for(int i=0;i<MaxUidNum;i++)
	{
		if(uidThrCtx[i][1]==uid)
		{
			uidThrCtx[i][2]=thrCode;
			return;
		}
	}

	//use new uid link
	for(int j=0;j<MaxUidNum;j++)
	{
		if(uidThrCtx[j][1]==0)
		{
			uidThrCtx[j][1]=uid;
			uidThrCtx[j][2]=thrCode;
			break;
		}
	}
}

//set uid to jvm
void setUid(jthread thr,char *method,jint thrCode)
{
	if(strcmp(method,"afterAdvice")!=0)
		return;

	error=gb_jvmti->SetLocalInt(thr,0,1,findUid(thrCode));
	checkError(error,"SetLocalObjectUid");

#ifdef DEBUG
	printf("[TMon]: inject uid ok!\n");
#endif

}

//this function will be invoked when a method returns
//it is used to get another part of the implicit context when "notify" exit
static void JNICALL callbackMethodExit(jvmtiEnv *jvmti_env, JNIEnv* env, jthread thr, jmethodID method,jboolean was_popped_by_exception,jvalue return_value)
{
	enter_critical_section();
	{
		//must be deallocate
		char *name=NULL,*sig=NULL,*gsig=NULL;
		jint thr_hash_code=0;

		error = gb_jvmti->GetMethodName(method, &name, &sig, &gsig);
		checkError(error,"GetMethodName");
		
		error = gb_jvmti->GetObjectHashCode(thr, &thr_hash_code);
		checkError(error,"GetObjectHashCode");
		
		//for demo
		if((strcmp(name,"notify")==0||strcmp(name,"notifyAll")==0)
			&&findUid(thr_hash_code)!=0)
		{

#ifdef DEBUG
			debug("method notify() or notifyAll() exit");
#endif

			jint owned_monitor_count=0;
			//must be deallocate
			jobject* owned_monitors_ptr=NULL;
			//get monitor objs owned by current thread
			error=gb_jvmti->GetOwnedMonitorInfo(thr,&owned_monitor_count,&owned_monitors_ptr);
			checkError(error,"GetOwnedMonitorInfo");

			for(int index=0;index<owned_monitor_count;index++)
			{
				//must be deallocate
				jvmtiMonitorUsage* info_ptr=NULL;
				info_ptr=(jvmtiMonitorUsage*)malloc(sizeof(jvmtiMonitorUsage));
				checkNull(info_ptr,"can't alloc mem");

				//get wait pools of each monitor obj
				error = gb_jvmti->GetObjectMonitorUsage(*(owned_monitors_ptr+index),info_ptr);
				checkError(error,"GetObjectMonitorUsage");

				//compare the two wait pools
				if(info_ptr->notify_waiter_count!=inf_head->info_ptr->notify_waiter_count)
				{
#ifdef DEBUG
					debug("find different monitor usage");
					debugDetail("inf_head->info_ptr->notify_waiter_count",inf_head->info_ptr->notify_waiter_count);
					debugDetail("info_ptr->notify_waiter_count",info_ptr->notify_waiter_count);
#endif
					for(int i=0;i<inf_head->info_ptr->notify_waiter_count;i++)
					{
						bool isObj=true;
						for(int j=0;j<info_ptr->notify_waiter_count;j++)
						{
							if(inf_head->info_ptr->notify_waiters+i!=info_ptr->notify_waiters+j)
							{
								isObj=false;
								break;
							}
						}
						if(isObj==true)
						{
							jint hash_code=0;
							error = gb_jvmti->GetObjectHashCode(*(inf_head->info_ptr->notify_waiters+i), &hash_code);
							checkError(error,"GetObjectHashCode");

							//for demo
							addThr(thr_hash_code,hash_code);

							//add new tail
							error =(jvmtiError)insertElem(&ctx_head,&ctx_tail,thr_hash_code,hash_code,name);
							checkError(error,"insertThrContext");
#ifdef DEBUG
							printf("[TMon]: thread (%10d) %10s  (%10d)\n",thr_hash_code,name,hash_code);
#endif
						}
						
					}
				}
				//free info_ptr
				free(inf_head->info_ptr);
				//remove old head
				removeElem(&inf_head,&inf_tail);
				error=gb_jvmti->Deallocate((unsigned char*)info_ptr);
				checkError(error,"Deallocate");
			}

			error=gb_jvmti->Deallocate((unsigned char*)owned_monitors_ptr);
			checkError(error,"Deallocate");
		}
		else
		{
		//***********
		//name is used now
		error=gb_jvmti->Deallocate((unsigned char*)name);
		checkError(error,"Deallocate");
		//***********
		}

		error=gb_jvmti->Deallocate((unsigned char*)sig);
		checkError(error,"Deallocate");
		error=gb_jvmti->Deallocate((unsigned char*)gsig);
		checkError(error,"Deallocate");
	}
	exit_critical_section(); 
}

//this function will be invoked when a method enters
//it is used to get the explicit context when "start" enters
//and part of the implicit context when "notify" enters
static void JNICALL callbackMethodEntry(jvmtiEnv *jvmti_env, JNIEnv* env, jthread thr, jmethodID method)
{
	enter_critical_section();
	{
		//must be deallocate
		char *name=NULL,*sig=NULL,*gsig=NULL;
		jint thr_hash_code=0;

		error = gb_jvmti->GetMethodName(method, &name, &sig, &gsig);
		checkError(error,"GetMethodName");

		error = gb_jvmti->GetObjectHashCode(thr, &thr_hash_code);
		checkError(error,"GetObjectHashCode");

		//for demo
		getUid(thr,name,thr_hash_code);
		setUid(thr,name,thr_hash_code);	
		//printf("re");

		//get the explicit context
		//for demo
		if((strcmp(name,"start")==0/*||strcmp(name,"interrupt")==0||
			strcmp(name,"join")==0||strcmp(name,"stop")==0||
			strcmp(name,"suspend")==0||strcmp(name,"resume")==0*/)
			&&findUid(thr_hash_code)!=0)
		{


			//must be deallocate
			jobject thd_ptr=NULL;
			error=gb_jvmti->GetLocalObject(thr,0,0,&thd_ptr);
			checkError(error,"GetLocalObject");
			jint hash_code=0;
			error = gb_jvmti->GetObjectHashCode(thd_ptr, &hash_code);
			checkError(error,"GetObjectHashCode");

			//for demo
			addThr(thr_hash_code,hash_code);

			//add new tail
			error =(jvmtiError)insertElem(&ctx_head,&ctx_tail,thr_hash_code,hash_code,name);
			checkError(error,"insertThrContext");
			error=gb_jvmti->Deallocate((unsigned char*)thd_ptr);
			checkError(error,"Deallocate");

#ifdef DEBUG
			printf("[TMon]: thread (%10d) %10s (%10d)\n",thr_hash_code,name,hash_code);
#endif

		}
		//get the implicit context
		else if((strcmp(name,"notify")==0||strcmp(name,"notifyAll")==0)
			&&findUid(thr_hash_code)!=0)
		{
			jint owned_monitor_count=0;
			//must be deallocate
			jobject *owned_monitors_ptr=NULL;
			error=gb_jvmti->GetOwnedMonitorInfo(thr,&owned_monitor_count,&owned_monitors_ptr);
			checkError(error,"GetOwnedMonitorInfo");

			for(int index=0;index<owned_monitor_count;index++)
			{
				//must be deallocate
				jvmtiMonitorUsage *info_ptr=NULL;
				info_ptr=(jvmtiMonitorUsage*)malloc(sizeof(jvmtiMonitorUsage));
				checkNull(info_ptr,"can't alloc mem");

				error = gb_jvmti->GetObjectMonitorUsage(*(owned_monitors_ptr+index),info_ptr);
				checkError(error,"GetObjectMonitorUsage");
				//add new tail
				error=(jvmtiError)insertElem(&inf_head,&inf_tail,info_ptr);
				checkError(error,"insertThrInfo");
			}

			error=gb_jvmti->Deallocate((unsigned char*)owned_monitors_ptr);
			checkError(error,"Deallocate");
		}
		else
		{
		//******name is used now
		error=gb_jvmti->Deallocate((unsigned char*)name);
		checkError(error,"Deallocate");
		//******
		}

		error=gb_jvmti->Deallocate((unsigned char*)sig);
		checkError(error,"Deallocate");
		error=gb_jvmti->Deallocate((unsigned char*)gsig);
		checkError(error,"Deallocate");
	}
	exit_critical_section();
}

//when JVM init, this function is invoked
JNIEXPORT jint JNICALL Agent_OnLoad(JavaVM *jvm, char *options, void *reserved)
{
	//for demo
	initUidThrCtx();

	jvmtiEventCallbacks callbacks;

	jint result = jvm->GetEnv((void **) &gb_jvmti, JVMTI_VERSION_1_0);

	if(result != JNI_OK || gb_jvmti==NULL)
	{
		printf("ERROR: Unable to access JVMTI!");
		return JNI_ERR;
	}
	
	//enable the capabilities of jvm
	memset(&gb_capa,0,sizeof(jvmtiCapabilities));
	
	gb_capa.can_signal_thread = 1;
	
	gb_capa.can_get_owned_monitor_info = 1;

	gb_capa.can_generate_method_exit_events = 1;
	
	gb_capa.can_generate_method_entry_events = 1;
	
	gb_capa.can_generate_exception_events = 1;
	
	gb_capa.can_generate_vm_object_alloc_events = 1;
	
	gb_capa.can_tag_objects = 1;  
	
	gb_capa.can_generate_all_class_hook_events=1;

	gb_capa.can_generate_native_method_bind_events=1;

	gb_capa.can_access_local_variables=1;

	gb_capa.can_get_monitor_info=1;

	error = gb_jvmti->AddCapabilities(&gb_capa);
	
	if(error != JVMTI_ERROR_NONE)
	{
		printf("ERROR: Can't get JVMTI capabilities");
		return JNI_ERR;
	}
	
	//regedit the callbacks
	memset(&callbacks,0,sizeof(jvmtiEventCallbacks));

	callbacks.MethodEntry = &callbackMethodEntry;

	callbacks.MethodExit = &callbackMethodExit;

	error = gb_jvmti->SetEventCallbacks(&callbacks, sizeof(callbacks));
	
	if(error != JVMTI_ERROR_NONE)
	{
		printf("ERROR: Can't set jvmti callback!");
		return JNI_ERR;
	}
	
	error = gb_jvmti->CreateRawMonitor("TMon", &gb_lock);
	
	if(error != JVMTI_ERROR_NONE)
	{
		printf("ERROR: Can't create raw monitor!");
		return JNI_ERR;
	}
	
	//register the event notification for each event
	error = gb_jvmti->SetEventNotificationMode(JVMTI_ENABLE,JVMTI_EVENT_VM_INIT, (jthread)NULL);
	
       error = gb_jvmti->SetEventNotificationMode(JVMTI_ENABLE,JVMTI_EVENT_EXCEPTION, (jthread)NULL);
	
	error = gb_jvmti->SetEventNotificationMode(JVMTI_ENABLE,JVMTI_EVENT_NATIVE_METHOD_BIND, (jthread)NULL);
	
	error = gb_jvmti->SetEventNotificationMode(JVMTI_ENABLE,JVMTI_EVENT_CLASS_FILE_LOAD_HOOK, (jthread)NULL);

	error = gb_jvmti->SetEventNotificationMode(JVMTI_ENABLE,JVMTI_EVENT_METHOD_ENTRY, (jthread)NULL);

	error = gb_jvmti->SetEventNotificationMode(JVMTI_ENABLE,JVMTI_EVENT_METHOD_EXIT, (jthread)NULL);
	
	return JNI_OK;
	
}

JNIEXPORT void JNICALL Agent_OnUnload(JavaVM *vm)
{

#ifdef DEBUG
	printLink(ctx_head);
	printUidThrCtx();
#endif
	printUidThrCtx();
}