package me.coley.recaf.workspace;

import me.coley.recaf.util.IOUtil;
import org.objectweb.asm.Type;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.instrument.Instrumentation;
import java.util.*;

/**
 * Importable instrumentation resource.
 *
 * @author Matt
 */
public class InstrumentationResource extends JavaResource {
	private final Instrumentation instrumentation;

	/**
	 * Constructs an instrumentation resource.
	 *
	 * @param instrumentation
	 * 		Instrumentation instance provided by the tools api.
	 */
	public InstrumentationResource(Instrumentation instrumentation) {
		super(ResourceKind.INSTRUMENTATION);
		this.instrumentation = instrumentation;
	}

	@Override
	protected Map<String, byte[]> loadClasses() throws IOException {
		Map<String, byte[]> classes = new HashMap<>();
		// iterate over loaded classes
		ByteArrayOutputStream out = new ByteArrayOutputStream();
		byte[] buffer = new byte[4096];
		for(Class<?> c : instrumentation.getAllLoadedClasses()) {
			String name = Type.getInternalName(c);
			// skip specified prefixes
			if(shouldSkip(name))
				continue;
			String path = name.concat(".class");
			ClassLoader loader = c.getClassLoader();
			InputStream in;
			if (loader == null) {
				in = ClassLoader.getSystemResourceAsStream(path);
			} else {
				in = loader.getResourceAsStream(path);
			}
			try(InputStream __ = in) {
				if(in != null) {
					out.reset();
					classes.put(name, IOUtil.toByteArray(in, out, buffer));
				}
			}
		}
		return classes;
	}

	@Override
	protected Map<String, byte[]> loadFiles() {
		return Collections.emptyMap();
	}

	@Override
	public String toString() {
		return "Instrumentation";
	}
}

@@ -1397,58 +1397,45 @@
 %0A%09%09%09
-try(InputStream __ = in) %7B%0A%09%09%09%09if(in != null) %7B
+if (in == null) %7B%0A%09%09%09%09continue;
 %0A%09%09
+%09%7D%0A
 %09%09%09o
@@ -1447,16 +1447,24 @@
 set();%0A%09
+%09%09try %7B%0A
 %09%09%09%09clas
@@ -1518,18 +1518,49 @@
 r));%0A%09%09%09
-%09
 %7D
+ finally %7B%0A%09%09%09%09IOUtil.close(in);
 %0A%09%09%09%7D%0A%09%09

Use IOUtil#close in InstrumentationResource