/*
 * Copyright 2018, The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <android/log.h>
#include <nativehelper/JNIHelp.h>
#include <nativehelper/ScopedUtfChars.h>
#include <jni.h>
#include <pcap.h>
#include <stdlib.h>
#include <string>
#include <vector>

#include "disassembler.h"
#include "nativehelper/scoped_primitive_array.h"

#include "next/test_buf_allocator.h"
#include "apflib.h"

#define ARRAY_SIZE(x) (sizeof(x) / sizeof((x)[0]))
#define LOG_TAG "ApfJniUtils"


// JNI function acting as simply call-through to native APF interpreter.
static jint
com_android_server_ApfTest_apfSimulate(JNIEnv* env, jclass, jint apf_version,
                                       jbyteArray jprogram, jbyteArray jpacket,
                                       jbyteArray jdata, jint filter_age) {

    ScopedByteArrayRO packet(env, jpacket);
    uint32_t packet_len = (uint32_t)packet.size();
    uint32_t program_len = env->GetArrayLength(jprogram);
    uint32_t data_len = jdata ? env->GetArrayLength(jdata) : 0;
    // For APFv6+ we need to make sure ram_len is a multiple of 4 bytes,
    // so that the counters (which are indexed from the back) are aligned.
    // We also need to guarantee room for APFv6's 5 u32 counters (20 bytes),
    // and APFv6.1's 6 u32 counters (24 bytes).
    // Furthermore APFv6.1 has a 1024 byte minimum ram_len.
    uint32_t ram_len = program_len + data_len;
    if (apf_version >= 6000) {
        uint32_t builtin_counters = (apf_version >= 6100) ? 6 : 5;
        uint32_t required_data_len = builtin_counters * sizeof(uint32_t);
        if (data_len < required_data_len) ram_len = program_len + required_data_len;
        ram_len += 3; ram_len &= ~3;
        if (apf_version >= 6100 && ram_len < 1024) ram_len = 1024;
    }
    std::vector<uint32_t> buf((ram_len + 3) / 4, 0);
    jbyte* jbuf = reinterpret_cast<jbyte*>(buf.data());

    env->GetByteArrayRegion(jprogram, 0, program_len, jbuf);
    if (jdata) {
        // Merge program and data into a single buffer.
        env->GetByteArrayRegion(jdata, 0, data_len, jbuf + ram_len - data_len);
    }

    jint result = apf_run_generic(
        apf_version, buf.data(), program_len, ram_len,
        reinterpret_cast<const uint8_t *>(packet.get()), packet_len,
        filter_age << 14);

    if (jdata) {
        env->SetByteArrayRegion(jdata, 0, data_len, jbuf + ram_len - data_len);
    }

    return result;
}

class ScopedPcap {
  public:
    explicit ScopedPcap(pcap_t* pcap) : pcap_ptr(pcap) {}
    ~ScopedPcap() {
        pcap_close(pcap_ptr);
    }

    pcap_t* get() const { return pcap_ptr; };
  private:
    pcap_t* const pcap_ptr;
};

class ScopedFILE {
  public:
    explicit ScopedFILE(FILE* fp) : file(fp) {}
    ~ScopedFILE() {
        fclose(file);
    }

    FILE* get() const { return file; };
  private:
    FILE* const file;
};

static void throwException(JNIEnv* env, const std::string& error) {
    jclass newExcCls = env->FindClass("java/lang/IllegalStateException");
    if (newExcCls == 0) {
      abort();
      return;
    }
    env->ThrowNew(newExcCls, error.c_str());
}

static jstring com_android_server_ApfTest_compileToBpf(JNIEnv* env, jclass, jstring jfilter) {
    ScopedUtfChars filter(env, jfilter);
    std::string bpf_string;
    ScopedPcap pcap(pcap_open_dead(DLT_EN10MB, 65535));
    if (pcap.get() == NULL) {
        throwException(env, "pcap_open_dead failed");
        return NULL;
    }

    // Compile "filter" to a BPF program
    bpf_program bpf;
    if (pcap_compile(pcap.get(), &bpf, filter.c_str(), 0, PCAP_NETMASK_UNKNOWN)) {
        throwException(env, "pcap_compile failed");
        return NULL;
    }

    // Translate BPF program to human-readable format
    const struct bpf_insn* insn = bpf.bf_insns;
    for (uint32_t i = 0; i < bpf.bf_len; i++) {
        bpf_string += bpf_image(insn++, i);
        bpf_string += "\n";
    }

    return env->NewStringUTF(bpf_string.c_str());
}

static jboolean com_android_server_ApfTest_compareBpfApf(
    JNIEnv* env, jclass, jint apf_version, jstring jfilter,
    jstring jpcap_filename, jbyteArray japf_program) {
    ScopedUtfChars filter(env, jfilter);
    ScopedUtfChars pcap_filename(env, jpcap_filename);
    uint32_t program_len = env->GetArrayLength(japf_program);
    uint32_t data_len = (apf_version > 4) ? 20 : 0;
    uint32_t ram_len = program_len + data_len;
    if (apf_version > 4) { ram_len += 3; ram_len &= ~3; }
    std::vector<uint32_t> apf_program((ram_len + 3) / 4, 0);
    env->GetByteArrayRegion(japf_program, 0, program_len,
                            reinterpret_cast<jbyte*>(apf_program.data()));

    // Open pcap file for BPF filtering
    ScopedFILE bpf_fp(fopen(pcap_filename.c_str(), "rb"));
    char pcap_error[PCAP_ERRBUF_SIZE];
    ScopedPcap bpf_pcap(pcap_fopen_offline(bpf_fp.get(), pcap_error));
    if (bpf_pcap.get() == NULL) {
        throwException(env, "pcap_fopen_offline failed: " + std::string(pcap_error));
        return false;
    }

    // Open pcap file for APF filtering
    ScopedFILE apf_fp(fopen(pcap_filename.c_str(), "rb"));
    ScopedPcap apf_pcap(pcap_fopen_offline(apf_fp.get(), pcap_error));
    if (apf_pcap.get() == NULL) {
        throwException(env, "pcap_fopen_offline failed: " + std::string(pcap_error));
        return false;
    }

    // Compile "filter" to a BPF program
    bpf_program bpf;
    if (pcap_compile(bpf_pcap.get(), &bpf, filter.c_str(), 0, PCAP_NETMASK_UNKNOWN)) {
        throwException(env, "pcap_compile failed");
        return false;
    }

    // Install BPF filter on bpf_pcap
    if (pcap_setfilter(bpf_pcap.get(), &bpf)) {
        throwException(env, "pcap_setfilter failed");
        return false;
    }

    while (1) {
        pcap_pkthdr bpf_header, apf_header;
        // Run BPF filter to the next matching packet.
        const uint8_t* bpf_packet = pcap_next(bpf_pcap.get(), &bpf_header);

        // Run APF filter to the next matching packet.
        const uint8_t* apf_packet;
        do {
            apf_packet = pcap_next(apf_pcap.get(), &apf_header);
        } while (apf_packet != NULL && !apf_run_generic(apf_version,
                apf_program.data(), program_len, ram_len,
                apf_packet, apf_header.len, 0 /* filter_age */));

        // Make sure both filters matched the same packet.
        if (apf_packet == NULL && bpf_packet == NULL)
            break;
        if (apf_packet == NULL || bpf_packet == NULL)
            return false;
        if (apf_header.len != bpf_header.len ||
                apf_header.ts.tv_sec != bpf_header.ts.tv_sec ||
                apf_header.ts.tv_usec != bpf_header.ts.tv_usec ||
                memcmp(apf_packet, bpf_packet, apf_header.len))
            return false;
    }
    return true;
}

static jobjectArray com_android_server_ApfTest_disassembleApf(
    JNIEnv* env, jclass, jbyteArray jprogram) {
    uint32_t program_len = env->GetArrayLength(jprogram);
    std::vector<uint8_t> buf(program_len, 0);

    env->GetByteArrayRegion(jprogram, 0, program_len,
                            reinterpret_cast<jbyte*>(buf.data()));
    std::vector<std::string> disassemble_output;
    for (uint32_t pc = 0; pc < program_len;) {
      // TODO: Implement proper selection of APFv4 or APFv6 code for
      // disassembly.
      const disas_ret ret =
          apf_disassemble(buf.data(), program_len, &pc, true /* is_v6*/);
      disassemble_output.emplace_back(ret.content);
    }
    jclass stringClass = env->FindClass("java/lang/String");
    jobjectArray disassembleOutput =
        env->NewObjectArray(disassemble_output.size(), stringClass, nullptr);

    for (jsize i = 0; i < (jsize) disassemble_output.size(); i++) {
         jstring j_disassemble_output =
             env->NewStringUTF(disassemble_output[i].c_str());
         env->SetObjectArrayElement(disassembleOutput, i, j_disassemble_output);
         env->DeleteLocalRef(j_disassemble_output);
    }

    return disassembleOutput;
}

static jobjectArray com_android_server_ApfTest_getAllTransmittedPackets(JNIEnv* env,
                                                                        jclass) {
    jclass arrayListClass = env->FindClass("java/util/ArrayList");
    jmethodID arrayListConstructor = env->GetMethodID(arrayListClass, "<init>", "()V");
    jobject arrayList = env->NewObject(arrayListClass, arrayListConstructor);

    jmethodID addMethod = env->GetMethodID(arrayListClass, "add", "(Ljava/lang/Object;)Z");
    packet_buffer *ptr = head;
    while (ptr) {
        jbyteArray jdata = env->NewByteArray((jint) ptr->len);
        if (jdata == NULL) {
            return static_cast<jobjectArray>(arrayList);
        }

        env->SetByteArrayRegion(jdata, 0, (jint) ptr->len,
                                reinterpret_cast<jbyte*>(ptr->data));
        env->CallBooleanMethod(arrayList, addMethod, jdata);
        env->DeleteLocalRef(jdata);

        ptr = ptr->next;
    }

    env->DeleteLocalRef(arrayListClass);
    return static_cast<jobjectArray>(arrayList);
}

void com_android_server_ApfTest_resetTransmittedPacketMemory(JNIEnv, jclass) {
    packet_buffer* current = head;
    packet_buffer* tmp = NULL;
    while (current) {
        tmp = current->next;
        free(current);
        current = tmp;
    }

    head = NULL;
    tail = NULL;
}

extern "C" jint JNI_OnLoad(JavaVM* vm, void*) {
    JNIEnv *env;
    if (vm->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION_1_6) != JNI_OK) {
        __android_log_print(ANDROID_LOG_ERROR, LOG_TAG, "ERROR: GetEnv failed");
        return -1;
    }

    static JNINativeMethod gMethods[] = {
        {"apfSimulate", "(I[B[B[BI)I",
         (void *)com_android_server_ApfTest_apfSimulate},
        {"compileToBpf", "(Ljava/lang/String;)Ljava/lang/String;",
         (void *)com_android_server_ApfTest_compileToBpf},
        {"compareBpfApf", "(ILjava/lang/String;Ljava/lang/String;[B)Z",
         (void *)com_android_server_ApfTest_compareBpfApf},
        {"disassembleApf", "([B)[Ljava/lang/String;",
         (void *)com_android_server_ApfTest_disassembleApf},
        {"getAllTransmittedPackets", "()Ljava/util/List;",
         (void *)com_android_server_ApfTest_getAllTransmittedPackets},
        {"resetTransmittedPacketMemory", "()V",
         (void *)com_android_server_ApfTest_resetTransmittedPacketMemory},
    };

    jniRegisterNativeMethods(env, "android/net/apf/ApfJniUtils",
            gMethods, ARRAY_SIZE(gMethods));

    return JNI_VERSION_1_6;
}