package org.graalvm.compiler.replacements.test.classfile;
import static org.junit.Assume.assumeTrue;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.lang.management.ManagementFactory;
import java.lang.reflect.Method;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.ProtectionDomain;
import java.util.jar.Attributes;
import java.util.jar.JarEntry;
import java.util.jar.JarOutputStream;
import java.util.jar.Manifest;
import javax.tools.ToolProvider;
import org.junit.Assert;
import org.junit.Test;
import org.graalvm.compiler.api.replacements.ClassSubstitution;
import org.graalvm.compiler.api.replacements.MethodSubstitution;
import org.graalvm.compiler.bytecode.BytecodeProvider;
import org.graalvm.compiler.core.test.GraalCompilerTest;
import org.graalvm.compiler.nodes.graphbuilderconf.GraphBuilderConfiguration;
import org.graalvm.compiler.nodes.graphbuilderconf.InvocationPlugins;
import org.graalvm.compiler.nodes.graphbuilderconf.InvocationPlugins.Registration;
import jdk.vm.ci.meta.ResolvedJavaMethod;
public class RedefineIntrinsicTest extends GraalCompilerTest {
public static class Original {
public static String getValue() {
return "original";
}
}
@ClassSubstitution(Original.class)
private static class Intrinsic {
@MethodSubstitution
public static String getValue() {
return "intrinsic";
}
}
@Override
protected GraphBuilderConfiguration editGraphBuilderConfiguration(GraphBuilderConfiguration conf) {
InvocationPlugins invocationPlugins = conf.getPlugins().getInvocationPlugins();
BytecodeProvider replacementBytecodeProvider = getReplacements().getReplacementBytecodeProvider();
Registration r = new Registration(invocationPlugins, Original.class, replacementBytecodeProvider);
r.registerMethodSubstitution(Intrinsic.class, "getValue");
return super.editGraphBuilderConfiguration(conf);
}
public static String callOriginalGetValue() {
return Original.getValue();
}
public static String callIntrinsicGetValue() {
return Intrinsic.getValue();
}
@Test
public void test() throws Throwable {
Object receiver = null;
Object[] args = {};
Assert.assertEquals("original", Original.getValue());
Assert.assertEquals("intrinsic", Intrinsic.getValue());
ResolvedJavaMethod callOriginalGetValue = getResolvedJavaMethod("callOriginalGetValue");
ResolvedJavaMethod callIntrinsicGetValue = getResolvedJavaMethod("callIntrinsicGetValue");
testAgainstExpected(callOriginalGetValue, new Result("intrinsic", null), receiver, args);
testAgainstExpected(callIntrinsicGetValue, new Result("intrinsic", null), receiver, args);
redefineIntrinsic();
Assert.assertEquals("original", Original.getValue());
Assert.assertEquals("redefined", Intrinsic.getValue());
testAgainstExpected(callOriginalGetValue, new Result("intrinsic", null), receiver, args);
}
static void add(JarOutputStream jar, Class<?> c) throws IOException {
String name = c.getName();
String classAsPath = name.replace('.', '/') + ".class";
jar.putNextEntry(new JarEntry(classAsPath));
InputStream stream = c.getClassLoader().getResourceAsStream(classAsPath);
int nRead;
byte[] buf = new byte[1024];
while ((nRead = stream.read(buf, 0, buf.length)) != -1) {
jar.write(buf, 0, nRead);
}
jar.closeEntry();
}
static void redefineIntrinsic() throws Exception {
Manifest manifest = new Manifest();
manifest.getMainAttributes().put(Attributes.Name.MANIFEST_VERSION, "1.0");
Attributes mainAttrs = manifest.getMainAttributes();
mainAttrs.putValue("Agent-Class", RedefinerAgent.class.getName());
mainAttrs.putValue("Can-Redefine-Classes", "true");
mainAttrs.putValue("Can-Retransform-Classes", "true");
Path jar = Files.createTempFile("myagent", ".jar");
try {
JarOutputStream jarStream = new JarOutputStream(new FileOutputStream(jar.toFile()), manifest);
add(jarStream, RedefinerAgent.class);
add(jarStream, Redefiner.class);
jarStream.close();
loadAgent(jar);
} finally {
Files.deleteIfExists(jar);
}
}
public static void loadAgent(Path agent) throws Exception {
String vmName = ManagementFactory.getRuntimeMXBean().getName();
int p = vmName.indexOf('@');
assumeTrue("VM name not in <pid>@<host> format: " + vmName, p != -1);
String pid = vmName.substring(0, p);
ClassLoader cl = ToolProvider.getSystemToolClassLoader();
Class<?> c = Class.forName("com.sun.tools.attach.VirtualMachine", true, cl);
Method attach = c.getDeclaredMethod("attach", String.class);
Method loadAgent = c.getDeclaredMethod("loadAgent", String.class, String.class);
Method detach = c.getDeclaredMethod("detach");
Object vm = attach.invoke(null, pid);
loadAgent.invoke(vm, agent.toString(), "");
detach.invoke(vm);
}
public static class RedefinerAgent {
public static void agentmain(@SuppressWarnings("unused") String args, Instrumentation inst) throws Exception {
if (inst.isRedefineClassesSupported() && inst.isRetransformClassesSupported()) {
inst.addTransformer(new Redefiner(), true);
Class<?>[] allClasses = inst.getAllLoadedClasses();
for (int i = 0; i < allClasses.length; i++) {
Class<?> c = allClasses[i];
if (c == Intrinsic.class) {
inst.retransformClasses(new Class<?>[]{c});
}
}
}
}
}
static class Redefiner implements ClassFileTransformer {
@Override
public byte[] transform(ClassLoader cl, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
if (Intrinsic.class.equals(classBeingRedefined)) {
String cf = new String(classfileBuffer);
int i = cf.indexOf("intrinsic");
Assert.assertTrue("cannot find \"intrinsic\" constant in " + Intrinsic.class.getSimpleName() + "'s class file", i > 0);
System.arraycopy("redefined".getBytes(), 0, classfileBuffer, i, "redefined".length());
}
return classfileBuffer;
}
}
}