net.bytebuddy.dynamic.loading.ClassInjector Java Examples

static Class injectClass(@Nullable ClassLoader targetClassLoader, @Nullable ProtectionDomain pd, String className, boolean isBootstrapClass) throws IOException, ClassNotFoundException {
    if (targetClassLoader == null) {
        if (isBootstrapClass) {
            return Class.forName(className, false, null);
        } else {
            throw new UnsupportedOperationException("Cannot load non-bootstrap class from bootstrap class loader");
        }
    }

    ClassInjector classInjector;
    if (targetClassLoader == ClassLoader.getSystemClassLoader()) {
        classInjector = ClassInjector.UsingReflection.ofSystemClassLoader();
    } else {
        classInjector = new ClassInjector.UsingReflection(targetClassLoader, pd, PackageDefinitionStrategy.NoOp.INSTANCE,
            true);
    }
    final byte[] classBytes = getAgentClassBytes(className);
    final TypeDescription typeDesc =
        new TypeDescription.Latent(className, 0, null, Collections.<TypeDescription.Generic>emptyList());
    Map<TypeDescription, byte[]> typeMap = new HashMap<>();
    typeMap.put(typeDesc, classBytes);
    return classInjector.inject(typeMap).values().iterator().next();
}
 

@Test
@SuppressWarnings("unchecked")
public void testMinimalRegistrationIndependentType() throws Exception {
    Annotation eagerAnnotation = mock(AuxiliaryType.SignatureRelevant.class);
    when(eagerAnnotation.annotationType()).thenReturn((Class) AuxiliaryType.SignatureRelevant.class);
    TypeDescription independent = mock(TypeDescription.class), dependent = mock(TypeDescription.class);
    when(independent.getDeclaredAnnotations()).thenReturn(new AnnotationList.ForLoadedAnnotations(eagerAnnotation));
    when(dependent.getDeclaredAnnotations()).thenReturn(new AnnotationList.Empty());
    Map<TypeDescription, byte[]> map = new HashMap<TypeDescription, byte[]>();
    map.put(independent, QUX);
    map.put(dependent, BAZ);
    when(dynamicType.getAuxiliaryTypes()).thenReturn(map);
    ClassInjector classInjector = mock(ClassInjector.class);
    when(injectionStrategy.resolve(classLoader, protectionDomain)).thenReturn(classInjector);
    when(classInjector.inject(Collections.singletonMap(independent, QUX)))
            .thenReturn(Collections.<TypeDescription, Class<?>>singletonMap(independent, Foo.class));
    LoadedTypeInitializer loadedTypeInitializer = mock(LoadedTypeInitializer.class);
    when(dynamicType.getLoadedTypeInitializers()).thenReturn(Collections.singletonMap(independent, loadedTypeInitializer));
    AgentBuilder.InitializationStrategy.Minimal.INSTANCE.register(dynamicType, classLoader, protectionDomain, injectionStrategy);
    verify(classInjector).inject(Collections.singletonMap(independent, QUX));
    verifyNoMoreInteractions(classInjector);
    verify(loadedTypeInitializer).onLoad(Foo.class);
    verifyNoMoreInteractions(loadedTypeInitializer);
}
 

private AgentBuilder withBootstrapAttaching(AgentBuilder agentBuilder) {
    if(moduleDescription.shouldInjectInBootstrap()){
        Logger.info(() -> "Bootstrap Injection activated.");
        agentBuilder = agentBuilder.with(new InjectionStrategy.UsingUnsafe.OfFactory(ClassInjector.UsingUnsafe.Factory.resolve(instrumentation)));
    }
    return agentBuilder;
}
 

@Test
public void testUsingUnsafeFactory() throws Exception {
    ClassInjector.UsingUnsafe.Factory factory = mock(ClassInjector.UsingUnsafe.Factory.class);
    ClassInjector classInjector = mock(ClassInjector.class);
    when(factory.make(classLoader, protectionDomain)).thenReturn(classInjector);
    assertThat(new AgentBuilder.InjectionStrategy.UsingUnsafe.OfFactory(factory).resolve(classLoader, protectionDomain), is(classInjector));
}
 

public static ClassLoadingStrategy<ClassLoader> resolve(Class<?> base) throws Exception {
    if (ClassInjector.UsingReflection.isAvailable()) {
        return ClassLoadingStrategy.Default.INJECTION;
    } else if (ClassInjector.UsingLookup.isAvailable()) {
        return ClassLoadingStrategy.UsingLookup.of(
                Class.forName("java.lang.invoke.MethodHandles")
                        .getMethod("privateLookupIn", Class.class, Class.forName("java.lang.invoke.MethodHandles$Lookup"))
                        .invoke(null, base, Class.forName("java.lang.invoke.MethodHandles").getMethod("lookup").invoke(null)));
    } else {
        throw new AssertionError("No injection strategy available");
    }
}
 

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());
}
 

public static void inject(File folder, Instrumentation instrumentation, List<String>  allClasses) {
    ClassInjector.UsingUnsafe.Factory.resolve(instrumentation)
            .make(null, null)
            .injectRaw(getTypeDefinitions(allClasses));
}
 

@Test
@ClassReflectionInjectionAvailableRule.Enforce
public void testUsingReflection() throws Exception {
    assertThat(AgentBuilder.InjectionStrategy.UsingReflection.INSTANCE.resolve(classLoader, protectionDomain),
            hasPrototype((ClassInjector) new ClassInjector.UsingReflection(classLoader, protectionDomain)));
}
 

@Test
@ClassUnsafeInjectionAvailableRule.Enforce
public void testUsingUnsafe() throws Exception {
    assertThat(AgentBuilder.InjectionStrategy.UsingUnsafe.INSTANCE.resolve(classLoader, protectionDomain),
            hasPrototype((ClassInjector) new ClassInjector.UsingUnsafe(classLoader, protectionDomain)));
}
 

@Test
public void testBootstrapInjectionBootClassLoader() throws Exception {
    assertThat(new AgentBuilder.InjectionStrategy.UsingInstrumentation(mock(Instrumentation.class), mock(File.class))
            .resolve(ClassLoadingStrategy.BOOTSTRAP_LOADER, protectionDomain), instanceOf(ClassInjector.UsingInstrumentation.class));
}
 

@Test
@ClassReflectionInjectionAvailableRule.Enforce
public void testBootstrapInjectionNonBootClassLoader() throws Exception {
    assertThat(new AgentBuilder.InjectionStrategy.UsingInstrumentation(mock(Instrumentation.class), mock(File.class))
            .resolve(classLoader, protectionDomain), instanceOf(ClassInjector.UsingReflection.class));
}
 

public Statement apply(Statement base, FrameworkMethod method, Object target) {
    return !ClassInjector.UsingReflection.isAvailable() && method.getAnnotation(Enforce.class) != null
            ? new NoOpStatement()
            : base;
}
 

public Statement apply(Statement base, FrameworkMethod method, Object target) {
    return !ClassInjector.UsingUnsafe.isAvailable() && method.getAnnotation(Enforce.class) != null
            ? new NoOpStatement()
            : base;
}