package unpatched;

import java.io.FileInputStream;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardCopyOption;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ForkJoinPool;
import java.util.stream.IntStream;

public class Main {

    private static final ForkJoinPool THREAD_POOL = new ForkJoinPool(10);

    public static native void loadLibUnpatched();

    public static native void execute(String libPath);

    private static Void executeFlow() {
        try {
            // Create a temp library file
            Path libraryPath = Files.createTempFile("temp", "my_lib.so");
            libraryPath.toFile().deleteOnExit();

            InputStream in = new FileInputStream("${PWD}/my_lib.so");
            assert in != null;
            Files.copy(in, libraryPath, StandardCopyOption.REPLACE_EXISTING);

            execute(libraryPath.toString());
            libraryPath.toFile().delete();
        } catch (Exception exception) {
            throw new RuntimeException(exception);
        }
        return null;
    }

    public static void execute() {
        for (int i = 0; i < 1000; i++) {
            List<CompletableFuture<Void>> futures = new ArrayList<>();
            IntStream.range(0, 100).forEach(ignore -> futures.add(CompletableFuture.supplyAsync(Main::executeFlow, THREAD_POOL)));
            futures.forEach(CompletableFuture::join);
        }
    }

    public static void main(String[] args) throws InterruptedException {
        System.load("${PWD}/unpatched_Main.so");
        loadLibUnpatched();

        Thread thread = new Thread(() -> {
            try {
                Thread.sleep(1000);
            } catch (Exception e) {}
            System.load("${PWD}/my_lib.so");
        });
        thread.start();
        execute();
        thread.join();
    }
}

/* DO NOT EDIT THIS FILE - it is machine generated */
#include <jni.h>
/* Header for class unpatched_Main */

#ifndef _Included_unpatched_Main
#define _Included_unpatched_Main
#ifdef __cplusplus
extern "C" {
#endif
/*
 * Class:     unpatched_Main
 * Method:    loadLibUnpatched
 * Signature: ()V
 */
JNIEXPORT void JNICALL Java_unpatched_Main_loadLibUnpatched
  (JNIEnv *, jclass);

/*
 * Class:     unpatched_Main
 * Method:    execute
 * Signature: (Ljava/lang/String;)V
 */
JNIEXPORT void JNICALL Java_unpatched_Main_execute
  (JNIEnv *, jclass, jstring);

#ifdef __cplusplus
}
#endif
#endif

#include "unpatched_Main.h"
#include <dlfcn.h>
#include <stdlib.h>

typedef void* (*my_dlopen_t)(char*, int);
typedef void* (*my_malloc_t)(size_t);
typedef void (*my_free_t)(void*);

my_dlopen_t local_dlopen = NULL;

JNIEXPORT void JNICALL Java_unpatched_Main_loadLibUnpatched(JNIEnv* jenv, jclass cls) {
    void *ptr = dlopen("my_dlopen_lib.so", RTLD_NOW | RTLD_GLOBAL);
    dlclose(ptr);
    ptr = dlopen("my_dlopen_lib.so", RTLD_NOW | RTLD_GLOBAL);

    local_dlopen = (my_dlopen_t) dlsym(ptr, "my_dlopen");
}

JNIEXPORT void JNICALL Java_unpatched_Main_execute(JNIEnv* jenv, jclass cls, jstring path) {

    // Convert jstring to char*
    const char* nativePath = jenv->GetStringUTFChars(path, NULL);

    void* lib = local_dlopen((char*)nativePath, RTLD_NOW | RTLD_GLOBAL);

    my_malloc_t local_malloc = (my_malloc_t) dlsym(lib, "my_malloc");
    my_free_t local_free = (my_free_t) dlsym(lib, "my_free");

    local_free(local_malloc(199999999));

    dlclose(lib);
    
    // IMPORTANT: Release the string when done to prevent memory leaks
    jenv->ReleaseStringUTFChars(path, nativePath);
}

#ifndef MYLIB_H
#define MYLIB_H

#define EXPORT __attribute__((visibility("default")))

#include "stdlib.h"


EXPORT void* my_malloc(size_t size);

EXPORT void my_free(void* ptr);

#endif

#include "lib.h"

EXPORT void* my_malloc(size_t size) {
    return malloc(size);
}

EXPORT void my_free(void* ptr) {
    free(ptr);
}

#ifndef MY_DLOPEN_LIB_H
#define MY_DLOPEN_LIB_H

#define EXPORT __attribute__((visibility("default")))


EXPORT void* my_dlopen(char* path, int options);

#endif

#include "lib_dlopen.h"
#include <dlfcn.h>
#include <stdio.h>

EXPORT void* my_dlopen(char* path, int options) {
    printf("my_dlopen: %s\n", path);
    return dlopen(path, options);
}