/*
 * Copyright (c) 2017, Google and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

#include <stdio.h>
#include <string.h>
#include "jvmti.h"

#ifdef __cplusplus
extern "C" {
#endif

#ifndef JNI_ENV_ARG

#ifdef __cplusplus
#define JNI_ENV_ARG(x, y) y
#define JNI_ENV_PTR(x) x
#else
#define JNI_ENV_ARG(x,y) x, y
#define JNI_ENV_PTR(x) (*x)
#endif

#endif

#define PASSED 0
#define FAILED 2

#define MAX_TRACES 400

static const char *EXC_CNAME = "java/lang/Exception";
static jvmtiEnv *jvmti = NULL;

static int check_error(jvmtiError err, const char* s) {
  if (err != JVMTI_ERROR_NONE) {
    printf("  ## %s error: %d\n", s, err);
    return 1;
  }
  return 0;
}

static int check_capability_error(jvmtiError err, const char* s) {
  if (err != JVMTI_ERROR_NONE) {
    if (err == JVMTI_ERROR_MUST_POSSESS_CAPABILITY) {
      return 0;
    }
    printf("  ## %s error: %d\n", s, err);
    return 1;
  }
  return 1;
}

static
jint throw_exc(JNIEnv *env, char *msg) {
  jclass exc_class = JNI_ENV_PTR(env)->FindClass(JNI_ENV_ARG(env, EXC_CNAME));

  if (exc_class == NULL) {
    printf("throw_exc: Error in FindClass(env, %s)\n", EXC_CNAME);
    return -1;
  }
  return JNI_ENV_PTR(env)->ThrowNew(JNI_ENV_ARG(env, exc_class), msg);
}

static jint Agent_Initialize(JavaVM *jvm, char *options, void *reserved);

JNIEXPORT
jint JNICALL Agent_OnLoad(JavaVM *jvm, char *options, void *reserved) {
  return Agent_Initialize(jvm, options, reserved);
}

JNIEXPORT
jint JNICALL Agent_OnAttach(JavaVM *jvm, char *options, void *reserved) {
  return Agent_Initialize(jvm, options, reserved);
}

JNIEXPORT
jint JNICALL JNI_OnLoad(JavaVM *jvm, void *reserved) {
  return JNI_VERSION_1_8;
}

JNIEXPORT void JNICALL OnVMInit(jvmtiEnv *jvmti, JNIEnv *jni_env, jthread thread) {
}

JNIEXPORT void JNICALL OnClassLoad(jvmtiEnv *jvmti_env, JNIEnv *jni_env,
                                              jthread thread, jclass klass) {
  // NOP.
}

JNIEXPORT void JNICALL OnClassPrepare(jvmtiEnv *jvmti_env, JNIEnv *jni_env,
                                      jthread thread, jclass klass) {
  // We need to do this to "prime the pump", as it were -- make sure
  // that all of the methodIDs have been initialized internally, for
  // AsyncGetCallTrace.
  jint method_count;
  jmethodID *methods = 0;
  jvmtiError err = (*jvmti)->GetClassMethods(jvmti, klass, &method_count, &methods);
  if ((err != JVMTI_ERROR_NONE) && (err != JVMTI_ERROR_CLASS_NOT_PREPARED)) {
    // JVMTI_ERROR_CLASS_NOT_PREPARED is okay because some classes may
    // be loaded but not prepared at this point.
    throw_exc(jni_env, "Failed to create method IDs for methods in class\n");
  }
}

static
jint Agent_Initialize(JavaVM *jvm, char *options, void *reserved) {
  jint res;

  res = JNI_ENV_PTR(jvm)->GetEnv(JNI_ENV_ARG(jvm, (void **) &jvmti),
                                 JVMTI_VERSION_9);
  if (res != JNI_OK || jvmti == NULL) {
    printf("    Error: wrong result of a valid call to GetEnv!\n");
    return JNI_ERR;
  }

  jvmtiEventCallbacks callbacks;
  memset(&callbacks, 0, sizeof(callbacks));

  callbacks.VMInit = &OnVMInit;
  callbacks.ClassLoad = &OnClassLoad;
  callbacks.ClassPrepare = &OnClassPrepare;

  jvmtiCapabilities caps;
  memset(&caps, 0, sizeof(caps));
  // Get line numbers, sample heap, and filename for the test.
  caps.can_get_line_numbers = 1;
  caps.can_sample_heap= 1;
  caps.can_get_source_file_name = 1;
  if (check_error((*jvmti)->AddCapabilities(jvmti, &caps),
                  "Add capabilities\n")){
    return JNI_ERR;
  }

  if (check_error((*jvmti)->SetEventCallbacks(jvmti, &callbacks,
                                              sizeof(jvmtiEventCallbacks)),
                  " Set Event Callbacks")) {
    return JNI_ERR;
  }
  if (check_error((*jvmti)->SetEventNotificationMode(jvmti, JVMTI_ENABLE,
                                                     JVMTI_EVENT_VM_INIT, NULL),
                  "Set Event for VM Init")) {
    return JNI_ERR;
  }
  if (check_error((*jvmti)->SetEventNotificationMode(jvmti, JVMTI_ENABLE,
                                                     JVMTI_EVENT_CLASS_LOAD, NULL),
                  "Set Event for Class Load")) {
    return JNI_ERR;
  }
  if (check_error( (*jvmti)->SetEventNotificationMode(jvmti, JVMTI_ENABLE,
                                                      JVMTI_EVENT_CLASS_PREPARE, NULL),
                  "Set Event for Class Prepare")) {
    return JNI_ERR;
  }

  return JNI_OK;
}

// Given a method and a location, this method gets the line number.
// Kind of expensive, comparatively.
static
jint get_line_number(jvmtiEnv *jvmti, jmethodID method, jlocation location) {
  // The location is -1 if the bci isn't known or -3 for a native method.
  if (location == -1 || location == -3) {
    return -1;
  }

  // Read the line number table.
  jvmtiLineNumberEntry *table_ptr = 0;
  jint line_number_table_entries;
  int jvmti_error = (*jvmti)->GetLineNumberTable(jvmti, method,
                                                 &line_number_table_entries,
                                                 &table_ptr);

  if (JVMTI_ERROR_NONE != jvmti_error) {
    return -1;
  }
  if (line_number_table_entries <= 0) {
    return -1;
  }
  if (line_number_table_entries == 1) {
    return table_ptr[0].line_number;
  }

  // Go through all the line numbers...
  jint last_location = table_ptr[0].start_location;
  int l;
  for (l = 1; l < line_number_table_entries; l++) {
    // ... and if you see one that is in the right place for your
    // location, you've found the line number!
    if ((location < table_ptr[l].start_location) &&
        (location >= last_location)) {
      return table_ptr[l - 1].line_number;
    }
    last_location = table_ptr[l].start_location;
  }

  if (location >= last_location) {
    return table_ptr[line_number_table_entries - 1].line_number;
  } else {
    return -1;
  }
}

typedef struct _ExpectedContentFrame {
  const char *name;
  const char *signature;
  const char *file_name;
  int line_number;
} ExpectedContentFrame;

static jint check_sample_content(JNIEnv *env,
                                 jvmtiStackTrace *trace,
                                 ExpectedContentFrame *expected,
                                 int expected_count) {
  int i;

  if (expected_count > trace->frame_count) {
    return 0;
  }

  for (i = 0; i < expected_count; i++) {
    // Get basic information out of the trace.
    int bci = trace->frames[i].location;
    jmethodID methodid = trace->frames[i].method;
    char *name = NULL, *signature = NULL, *file_name = NULL;

    if (bci < 0) {
      return 0;
    }

    // Transform into usable information.
    int line_number = get_line_number(jvmti, methodid, bci);
    (*jvmti)->GetMethodName(jvmti, methodid, &name, &signature, 0);

    jclass declaring_class;
    if (JVMTI_ERROR_NONE !=
        (*jvmti)->GetMethodDeclaringClass(jvmti, methodid, &declaring_class)) {
      return 0;
    }

    jvmtiError err = (*jvmti)->GetSourceFileName(jvmti, declaring_class,
                                                 &file_name);
    if (err != JVMTI_ERROR_NONE) {
      return 0;
    }

    // Compare now, none should be NULL.
    if (name == NULL) {
      return 0;
    }

    if (file_name == NULL) {
      return 0;
    }

    if (signature == NULL) {
      return 0;
    }

    if (strcmp(name, expected[i].name) ||
        strcmp(signature, expected[i].signature) ||
        strcmp(file_name, expected[i].file_name) ||
        line_number != expected[i].line_number) {
      return 0;
    }
  }

  return 1;
}

static jint compare_samples(JNIEnv* env, jvmtiStackTrace* traces, int trace_count,
                            ExpectedContentFrame* expected_content, size_t size) {
  // We expect the code to record correctly the bci, retrieve the line
  // number, have the right method and the class name of the first frames.
  int i;
  for (i = 0; i < trace_count; i++) {
    jvmtiStackTrace *trace = traces + i;
    if (check_sample_content(env, trace, expected_content, size)) {
      // At least one frame matched what we were looking for.
      return 1;
    }
  }

  return 0;
}

static jint check_samples(JNIEnv* env, ExpectedContentFrame* expected,
                          size_t size,
                          jvmtiError (*const get_traces)(jvmtiEnv*, jvmtiStackTraces*)) {
  jvmtiStackTraces traces;
  jvmtiError error = get_traces(jvmti, &traces);

  if (error != JVMTI_ERROR_NONE) {
    return 0;
  }

  int result = compare_samples(env, traces.stack_traces, traces.trace_count,
                            expected, size);
  (*jvmti)->ReleaseTraces(jvmti, &traces);
  return result;
}

static jint frames_exist_live(JNIEnv* env, ExpectedContentFrame* expected,
                            size_t size) {
  return check_samples(env, expected, size, (*jvmti)->GetLiveTraces);
}

static jint frames_exist_recent(JNIEnv* env, ExpectedContentFrame* expected,
                              size_t size) {
  return check_samples(env, expected, size, (*jvmti)->GetGarbageTraces);
}

static jint frames_exist_frequent(JNIEnv* env, ExpectedContentFrame* expected,
                              size_t size) {
  return check_samples(env, expected, size, (*jvmti)->GetFrequentGarbageTraces);
}

// Static native API for various tests.
static void fill_native_frames(JNIEnv* env, jobjectArray frames,
                               ExpectedContentFrame* native_frames, size_t size) {
  size_t i;
  for(i = 0; i < size; i++) {
    jobject obj = (*env)->GetObjectArrayElement(env, frames, i);
    jclass frame_class = (*env)->GetObjectClass(env, obj);
    jfieldID line_number_field_id = (*env)->GetFieldID(env, frame_class, "lineNumber", "I");
    int line_number = (*env)->GetIntField(env, obj, line_number_field_id);

    jfieldID string_id = (*env)->GetFieldID(env, frame_class, "method", "Ljava/lang/String;");
    jstring string_object = (jstring) (*env)->GetObjectField(env, obj, string_id);
    const char* method = (*env)->GetStringUTFChars(env, string_object, 0);

    string_id = (*env)->GetFieldID(env, frame_class, "fileName", "Ljava/lang/String;");
    string_object = (jstring) (*env)->GetObjectField(env, obj, string_id);
    const char* file_name = (*env)->GetStringUTFChars(env, string_object, 0);

    string_id = (*env)->GetFieldID(env, frame_class, "signature", "Ljava/lang/String;");
    string_object = (jstring) (*env)->GetObjectField(env, obj, string_id);
    const char* signature= (*env)->GetStringUTFChars(env, string_object, 0);

    native_frames[i].name = method;
    native_frames[i].file_name = file_name;
    native_frames[i].signature = signature;
    native_frames[i].line_number = line_number;
  }
}

static jint checkAnd(JNIEnv *env, jobjectArray frames, int live, int recent,
                     int frequent) {
  jobject loader = NULL;

  if (frames == NULL) {
    return 0;
  }

  // Start by transforming the frames into a C-friendly structure.
  jsize size = (*env)->GetArrayLength(env, frames);
  ExpectedContentFrame native_frames[size];
  fill_native_frames(env, frames, native_frames, size);

  if (jvmti == NULL) {
    throw_exc(env, "JVMTI client was not properly loaded!\n");
    return 0;
  }

  int result = 1;

  if (live) {
    result = frames_exist_live(env, native_frames, size);
  }

  if (recent) {
    result = result &&
        frames_exist_recent(env, native_frames, size);
  }

  if (frequent) {
    result = result &&
        frames_exist_frequent(env, native_frames, size);
  }

  return result;
}

static jint checkOr(JNIEnv *env, jobjectArray frames, int live, int recent,
                    int frequent) {
  jobject loader = NULL;

  if (frames == NULL) {
    return 0;
  }

  // Start by transforming the frames into a C-friendly structure.
  jsize size = (*env)->GetArrayLength(env, frames);
  ExpectedContentFrame native_frames[size];
  fill_native_frames(env, frames, native_frames, size);

  if (jvmti == NULL) {
    throw_exc(env, "JVMTI client was not properly loaded!\n");
    return 0;
  }

  int result = 0;

  if (live) {
    result = frames_exist_live(env, native_frames, size);
  }

  if (recent) {
    result = result ||
        frames_exist_recent(env, native_frames, size);
  }

  if (frequent) {
    result = result ||
        frames_exist_frequent(env, native_frames, size);
  }

  return result;
}

static jint checkAll(JNIEnv *env, jobjectArray frames) {
  return checkAnd(env, frames, 1, 1, 1);
}

static jint checkNone(JNIEnv *env, jobjectArray frames) {
  jobject loader = NULL;

  if (frames == NULL) {
    return 0;
  }

  // Start by transforming the frames into a C-friendly structure.
  jsize size = (*env)->GetArrayLength(env, frames);
  ExpectedContentFrame native_frames[size];
  fill_native_frames(env, frames, native_frames, size);

  if (jvmti == NULL) {
    throw_exc(env, "JVMTI client was not properly loaded!\n");
    return 0;
  }

  if ((!frames_exist_live(env, native_frames, size)) &&
      (!frames_exist_recent(env, native_frames, size)) &&
      (!frames_exist_frequent(env, native_frames, size))) {
    return 1;
  }
  return 0;
}

static void enable_sampling() {
  check_error((*jvmti)->StartHeapSampling(jvmti, 1<<19, MAX_TRACES),
              "Start Heap Sampling");
}

static void disable_sampling() {
  check_error((*jvmti)->StopHeapSampling(jvmti), "Stop Heap Sampling");
}

// HeapMonitorTest JNI.
JNIEXPORT jint JNICALL
Java_MyPackage_HeapMonitorTest_checkFrames(JNIEnv *env, jclass cls, jobjectArray frames) {
  // We want the frames in each part.
  if (!checkAll(env, frames)) {
    return FAILED;
  }
  return PASSED;
}

JNIEXPORT void JNICALL
Java_MyPackage_HeapMonitorTest_enableSampling(JNIEnv *env, jclass cls) {
  enable_sampling();
}

// HeapMonitorOnOffTest JNI.
JNIEXPORT jint JNICALL
Java_MyPackage_HeapMonitorOnOffTest_checkFrames(JNIEnv *env, jclass cls, jobjectArray frames) {
  // We want the frames in each part.
  if (!checkAll(env, frames)) {
    return FAILED;
  }
  return PASSED;
}

JNIEXPORT jint JNICALL
Java_MyPackage_HeapMonitorOnOffTest_checkWipeOut(JNIEnv *env, jclass cls, jobjectArray frames) {
  // We want the frames in none of the parts.
  if (!checkNone(env, frames)) {
    return FAILED;
  }
  return PASSED;
}

JNIEXPORT void JNICALL
Java_MyPackage_HeapMonitorOnOffTest_enableSampling(JNIEnv *env, jclass cls) {
  enable_sampling();
}

JNIEXPORT void JNICALL
Java_MyPackage_HeapMonitorOnOffTest_disableSampling(JNIEnv *env, jclass cls) {
  disable_sampling();
}

// HeapMonitorRecentTest JNI.
JNIEXPORT jint JNICALL
Java_MyPackage_HeapMonitorRecentTest_checkFrames(JNIEnv *env, jclass cls, jobjectArray frames) {
  // We want the frames in each part.
  if (!checkAll(env, frames)) {
    return FAILED;
  }
  return PASSED;
}

JNIEXPORT jint JNICALL
Java_MyPackage_HeapMonitorRecentTest_checkLiveOrRecentFrames(JNIEnv *env, jclass cls, jobjectArray frames) {
  if (checkOr(env, frames, 1, 1, 0)) {
    return FAILED;
  }
  return PASSED;
}

JNIEXPORT jint JNICALL
Java_MyPackage_HeapMonitorRecentTest_checkLiveAndRecentFrames(JNIEnv *env, jclass cls, jobjectArray frames) {
  if (checkAnd(env, frames, 1, 1, 0)) {
    return FAILED;
  }
  return PASSED;
}

JNIEXPORT void JNICALL
Java_MyPackage_HeapMonitorRecentTest_enableSampling(JNIEnv *env, jclass cls) {
  enable_sampling();
}

// HeapMonitorFrequentTest JNI.
JNIEXPORT jint JNICALL
Java_MyPackage_HeapMonitorFrequentTest_checkFrames(JNIEnv *env, jclass cls, jobjectArray frames) {
  // We want the frames in each part.
  if (!checkAll(env, frames)) {
    return FAILED;
  }
  return PASSED;
}

JNIEXPORT jint JNICALL
Java_MyPackage_HeapMonitorFrequentTest_checkFrequentFrames(JNIEnv *env, jclass cls, jobjectArray frames) {
  if (checkAnd(env, frames, 0, 0, 1)) {
    return PASSED;
  }
  return FAILED;
}

JNIEXPORT void JNICALL
Java_MyPackage_HeapMonitorFrequentTest_enableSampling(JNIEnv *env, jclass cls) {
  enable_sampling();
}

JNIEXPORT jboolean JNICALL
Java_MyPackage_HeapMonitorNoCapabilityTest_allSamplingMethodsFail(JNIEnv *env, jclass cls) {
  jvmtiCapabilities caps;
  memset(&caps, 0, sizeof(caps));
  caps.can_sample_heap= 1;
  if (check_error((*jvmti)->RelinquishCapabilities(jvmti, &caps),
                  "Add capabilities\n")){
    return 0;
  }

  if (check_capability_error((*jvmti)->StartHeapSampling(jvmti, 1<<19,
                                                         MAX_TRACES),
                             "Start Heap Sampling")) {
    return 0;
  }

  if (check_capability_error((*jvmti)->StopHeapSampling(jvmti),
                             "Stop Heap Sampling")) {
    return 0;
  }

  if (check_capability_error((*jvmti)->ReleaseTraces(jvmti, NULL),
                             "Release Traces")) {
    return 0;
  }

  if (check_capability_error((*jvmti)->GetHeapSamplingStats(jvmti, NULL),
                             "Get Heap Sampling Stats")) {
    return 0;
  }

  if (check_capability_error((*jvmti)->GetGarbageTraces(jvmti, NULL),
                             "Get Garbage Traces")) {
    return 0;
  }

  if (check_capability_error((*jvmti)->GetFrequentGarbageTraces(jvmti, NULL),
                             "Get Frequent Garbage Traces")) {
    return 0;
  }

  if (check_capability_error((*jvmti)->GetLiveTraces(jvmti, NULL),
                             "Get Live Traces")) {
    return 0;
  }

  // Calling enable sampling should fail now.
  return 1;
}

JNIEXPORT void JNICALL
Java_MyPackage_HeapMonitorStatTest_enableSampling(JNIEnv *env, jclass cls) {
  enable_sampling();
}

JNIEXPORT jint JNICALL
Java_MyPackage_HeapMonitorStatTest_statsNull(JNIEnv *env, jclass cls) {
  jvmtiHeapSamplingStats stats;
  check_error((*jvmti)->GetHeapSamplingStats(jvmti, &stats),
              "Heap Sampling Statistics");

  jvmtiHeapSamplingStats zero;
  memset(&zero, 0, sizeof(zero));
  return memcmp(&stats, &zero, sizeof(zero)) == 0;
}

#ifdef __cplusplus
}
#endif