Byte Buddy PackageInjector

PackageInjector.java

import com.google.common.reflect.ClassPath;
import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.ClassFileLocator;
import net.bytebuddy.dynamic.loading.ClassInjector;

import java.io.File;
import java.io.IOException;
import java.lang.instrument.Instrumentation;
import java.nio.file.Files;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import static java.util.Objects.requireNonNull;
import static java.util.stream.Collectors.toMap;

@SuppressWarnings("UnstableApiUsage")
public class PackageInjector {

    private final Set<Class<?>> classes;

    public PackageInjector() {
        this.classes = new HashSet<>();
    }

    public PackageInjector(String packageName) throws IOException {
        this(Thread.currentThread().getContextClassLoader(), packageName);
    }

    public PackageInjector(ClassLoader classLoader, String packageName) throws IOException {
        this();
        addPackage(classLoader, packageName);
    }

    public void addPackage(ClassLoader classLoader, String packageName) throws IOException {
        requireNonNull(classLoader, "classLoader");
        requireNonNull(packageName, "packageName");

        classes.addAll(
            ClassPath.from(classLoader)
                .getTopLevelClassesRecursive(packageName).stream()
                .map(ClassPath.ClassInfo::load)
                .collect(Collectors.toSet()));
    }

    private Map<? extends TypeDescription, byte[]> getTypes() {
        return classes.stream().collect(
                toMap(TypeDescription.ForLoadedType::new, ClassFileLocator.ForClassLoader::read));
    }

    public void inject(ClassInjector.UsingInstrumentation.Target target) throws IOException {
        inject(target, ByteBuddyAgent.install());
    }

    public void inject(ClassInjector.UsingInstrumentation.Target target, Instrumentation instrumentation) throws IOException {
        File tempDir = Files.createTempDirectory("slimeInjector" + hashCode()).toFile();
        tempDir.deleteOnExit();

        inject(tempDir, target, instrumentation);
    }

    public void inject(File folder, ClassInjector.UsingInstrumentation.Target target, Instrumentation instrumentation) {
        ClassInjector injector = ClassInjector.UsingInstrumentation.of(folder, target, instrumentation);

        injector.inject(getTypes());
    }
}