mirror of
https://github.com/wassname/ray.git
synced 2026-07-10 00:35:08 +08:00
[JavaWorker] Enable java worker support (#2094)
* Enable java worker support -------------------------- This commit includes a tailored version of the Java worker implementation from Ant Financial. The changes for build system, python module, src module and arrow are in other commits, this commit consists of the following modules: - java/api: Ray API definition - java/common: utilities - java/hook: binary rewrite of the Java byte-code for remote execution - java/runtime-common: common implementation of the runtime in worker - java/runtime-dev: a pure-java mock implementation of the runtime for fast development - java/runtime-native: a native implementation of the runtime - java/test: various tests Contributors for this work: Guyang Song, Peng Cao, Senlin Zhu,Xiaoying Chu, Yiming Yu, Yujie Liu, Zhenyu Guo * change the format of java help document from markdown to RST * update the vesion of Arrow for java worker * adapt the new version of plasma java client from arrow which use byte[] instead of custom type * add java worker test to ci * add the example module for better usage guide
This commit is contained in:
Yujie Liu
committed by
Philipp Moritz
Philipp Moritz
parent
74cca3b284
commit
a8d3c057c1
@@ -0,0 +1,88 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0"
|
||||
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
|
||||
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
|
||||
<parent>
|
||||
<groupId>org.ray.parent</groupId>
|
||||
<artifactId>ray-superpom</artifactId>
|
||||
<version>1.0</version>
|
||||
</parent>
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
|
||||
<groupId>org.ray</groupId>
|
||||
<artifactId>ray-hook</artifactId>
|
||||
<name>java api hook for ray</name>
|
||||
<description>java api hook for ray</description>
|
||||
<url></url>
|
||||
|
||||
<packaging>jar</packaging>
|
||||
|
||||
<dependencies>
|
||||
|
||||
<!-- https://mvnrepository.com/artifact/org.ow2.asm/asm -->
|
||||
<dependency>
|
||||
<groupId>org.ow2.asm</groupId>
|
||||
<artifactId>asm</artifactId>
|
||||
<version>6.0</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>commons-codec</groupId>
|
||||
<artifactId>commons-codec</artifactId>
|
||||
<version>1.4</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>commons-io</groupId>
|
||||
<artifactId>commons-io</artifactId>
|
||||
<version>2.5</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.ray</groupId>
|
||||
<artifactId>ray-common</artifactId>
|
||||
<version>1.0</version>
|
||||
</dependency>
|
||||
|
||||
</dependencies>
|
||||
|
||||
<build>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<artifactId>maven-jar-plugin</artifactId>
|
||||
<version>2.5</version>
|
||||
<configuration>
|
||||
<archive>
|
||||
<manifestEntries>
|
||||
<Premain-Class>org.ray.hook.Agent</Premain-Class>
|
||||
<Can-Retransform-Classes>true</Can-Retransform-Classes>
|
||||
</manifestEntries>
|
||||
</archive>
|
||||
</configuration>
|
||||
</plugin>
|
||||
|
||||
<plugin>
|
||||
<artifactId>maven-shade-plugin</artifactId>
|
||||
<version>3.1.0</version>
|
||||
<executions>
|
||||
<execution>
|
||||
<phase>package</phase>
|
||||
<goals>
|
||||
<goal>shade</goal>
|
||||
</goals>
|
||||
<configuration>
|
||||
<relocations>
|
||||
<relocation>
|
||||
<pattern>org.objectweb.asm</pattern>
|
||||
<shadedPattern>agent.org.objectweb.asm</shadedPattern>
|
||||
</relocation>
|
||||
</relocations>
|
||||
<createDependencyReducedPom>false</createDependencyReducedPom>
|
||||
</configuration>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
</project>
|
||||
@@ -0,0 +1,49 @@
|
||||
package org.ray.hook;
|
||||
|
||||
import java.util.Set;
|
||||
import org.objectweb.asm.ClassReader;
|
||||
import org.objectweb.asm.ClassWriter;
|
||||
|
||||
public class ClassAdapter {
|
||||
|
||||
public static class Result {
|
||||
|
||||
public byte[] classBuffer;
|
||||
public Set<MethodId> changedMethods;
|
||||
}
|
||||
|
||||
public static Result hookClass(ClassLoader loader, String className, byte[] classfileBuffer) {
|
||||
// we have to comment out this quick filter as this is not accurate
|
||||
// e.g., org/ray/api/test/ActorTest$Adder.class is skipped!!!
|
||||
// even worse, this is non-deterministic...
|
||||
/*
|
||||
if (detectBody.contains("org/ray/hook/")) {
|
||||
return classfileBuffer;
|
||||
}
|
||||
*/
|
||||
|
||||
ClassReader reader = new ClassReader(classfileBuffer);
|
||||
ClassWriter writer = new ClassWriter(reader, 0);
|
||||
ClassDetectVisitor pre = new ClassDetectVisitor(loader, writer, className);
|
||||
byte[] result;
|
||||
|
||||
reader.accept(pre, ClassReader.SKIP_FRAMES);
|
||||
if (pre.detectedMethods().isEmpty() && pre.actorCalls() == 0) {
|
||||
result = classfileBuffer;
|
||||
} else {
|
||||
if (pre.actorCalls() > 0) {
|
||||
reader = new ClassReader(writer.toByteArray());
|
||||
}
|
||||
|
||||
writer = new ClassWriter(reader, ClassWriter.COMPUTE_FRAMES);
|
||||
reader.accept(new ClassOverrideVisitor(writer, className, pre.detectedMethods()),
|
||||
ClassReader.SKIP_FRAMES);
|
||||
result = writer.toByteArray();
|
||||
}
|
||||
|
||||
Result rr = new Result();
|
||||
rr.changedMethods = pre.detectedMethods();
|
||||
rr.classBuffer = result;
|
||||
return rr;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
package org.ray.hook;
|
||||
|
||||
import java.security.InvalidParameterException;
|
||||
import java.util.HashSet;
|
||||
import java.util.Set;
|
||||
import org.objectweb.asm.AnnotationVisitor;
|
||||
import org.objectweb.asm.ClassVisitor;
|
||||
import org.objectweb.asm.Handle;
|
||||
import org.objectweb.asm.MethodVisitor;
|
||||
import org.objectweb.asm.Opcodes;
|
||||
import org.objectweb.asm.Type;
|
||||
|
||||
/**
|
||||
* rewrite phase 1
|
||||
*/
|
||||
public class ClassDetectVisitor extends ClassVisitor {
|
||||
|
||||
final String className;
|
||||
final Set<MethodId> rayMethods = new HashSet<>();
|
||||
boolean isActor = false;
|
||||
static int count = 0;
|
||||
int actorCalls = 0;
|
||||
final ClassLoader loader;
|
||||
|
||||
public int actorCalls() {
|
||||
return actorCalls;
|
||||
}
|
||||
|
||||
public ClassDetectVisitor(ClassLoader loader, ClassVisitor origin, String className) {
|
||||
super(Opcodes.ASM6, origin);
|
||||
this.className = className;
|
||||
this.loader = loader;
|
||||
}
|
||||
|
||||
public Set<MethodId> detectedMethods() {
|
||||
return rayMethods;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void visitInnerClass(String name, String outerName,
|
||||
String innerName, int access) {
|
||||
// System.err.println("visist inner class " + outerName + "$" + innerName);
|
||||
super.visitInnerClass(name, outerName, innerName, access);
|
||||
}
|
||||
|
||||
@Override
|
||||
public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
|
||||
if (desc.contains("Lorg/ray/api/RayRemote;")) {
|
||||
isActor = true;
|
||||
}
|
||||
return super.visitAnnotation(desc, visible);
|
||||
}
|
||||
|
||||
private void visitRayMethod(int access, String name, String mdesc) {
|
||||
if (name.equals("<init>")) {
|
||||
return;
|
||||
}
|
||||
|
||||
MethodId m = new MethodId(className, name, mdesc, (access & Opcodes.ACC_STATIC) != 0, loader);
|
||||
rayMethods.add(m);
|
||||
//System.err.println("Visit " + m.toString());
|
||||
count++;
|
||||
}
|
||||
|
||||
@Override
|
||||
public MethodVisitor visitMethod(int access, String name, String mdesc, String signature,
|
||||
String[] exceptions) {
|
||||
//System.out.println("Visit " + className + "." + name);
|
||||
if (isActor && (access & Opcodes.ACC_PUBLIC) != 0) {
|
||||
visitRayMethod(access, name, mdesc);
|
||||
}
|
||||
|
||||
MethodVisitor origin = super.visitMethod(access, name, mdesc, signature, exceptions);
|
||||
return new MethodVisitor(this.api, origin) {
|
||||
@Override
|
||||
public AnnotationVisitor visitAnnotation(String adesc, boolean visible) {
|
||||
//handle rayRemote annotation
|
||||
if (adesc.contains("Lorg/ray/api/RayRemote;")) {
|
||||
visitRayMethod(access, name, mdesc);
|
||||
}
|
||||
return super.visitAnnotation(adesc, visible);
|
||||
}
|
||||
|
||||
private boolean isValidCallParameterOrReturnType(Type t) {
|
||||
if (t.equals(Type.VOID_TYPE)) {
|
||||
return false;
|
||||
}
|
||||
if (t.equals(Type.BOOLEAN_TYPE)) {
|
||||
return false;
|
||||
}
|
||||
if (t.equals(Type.CHAR_TYPE)) {
|
||||
return false;
|
||||
}
|
||||
if (t.equals(Type.BYTE_TYPE)) {
|
||||
return false;
|
||||
}
|
||||
if (t.equals(Type.SHORT_TYPE)) {
|
||||
return false;
|
||||
}
|
||||
if (t.equals(Type.INT_TYPE)) {
|
||||
return false;
|
||||
}
|
||||
if (t.equals(Type.FLOAT_TYPE)) {
|
||||
return false;
|
||||
}
|
||||
if (t.equals(Type.LONG_TYPE)) {
|
||||
return false;
|
||||
}
|
||||
if (t.equals(Type.DOUBLE_TYPE)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void visitInvokeDynamicInsn(String name, String desc, Handle bsm,
|
||||
Object... bsmArgs) {
|
||||
|
||||
// fix all actor calls from InvokeVirtual to InvokeStatic
|
||||
if (desc.contains("org/ray/api/funcs/RayFunc_")) {
|
||||
int count = bsmArgs.length;
|
||||
for (int i = 0; i < count; ++i) {
|
||||
Object arg = bsmArgs[i];
|
||||
// System.err.println(arg.getClass().getName() + " " + arg.toString());
|
||||
if (arg.getClass().equals(Handle.class)) {
|
||||
Handle h = (Handle) arg;
|
||||
if (h.getTag() == Opcodes.H_INVOKEVIRTUAL) {
|
||||
String dsptr = h.getDesc();
|
||||
|
||||
Type[] argTypes = Type.getArgumentTypes(dsptr);
|
||||
for (Type argt : argTypes) {
|
||||
if (!isValidCallParameterOrReturnType(argt)) {
|
||||
throw new InvalidParameterException(
|
||||
"cannot use primitive parameter type '" + argt.getClassName()
|
||||
+ "' in method " + h.getOwner() + "." + h.getName());
|
||||
}
|
||||
}
|
||||
Type retType = Type.getReturnType(dsptr);
|
||||
if (!isValidCallParameterOrReturnType(retType)) {
|
||||
throw new InvalidParameterException(
|
||||
"cannot use primitive return type '" + retType.getClassName() + "' in method "
|
||||
+ h.getOwner() + "." + h.getName());
|
||||
}
|
||||
|
||||
dsptr = "(L" + h.getOwner() + ";" + dsptr.substring(1);
|
||||
Handle newh = new Handle(
|
||||
Opcodes.H_INVOKESTATIC,
|
||||
h.getOwner(),
|
||||
h.getName() + MethodId.getFunctionIdPostfix,
|
||||
dsptr,
|
||||
h.isInterface());
|
||||
bsmArgs[i] = newh;
|
||||
//System.err.println("Change ray.call from " + h + " -> " + newh + ", isInterface = " + h.isInterface());
|
||||
++actorCalls;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
super.visitInvokeDynamicInsn(name, desc, bsm, bsmArgs);
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,223 @@
|
||||
package org.ray.hook;
|
||||
|
||||
import java.util.Set;
|
||||
import org.objectweb.asm.ClassVisitor;
|
||||
import org.objectweb.asm.Label;
|
||||
import org.objectweb.asm.MethodVisitor;
|
||||
import org.objectweb.asm.Opcodes;
|
||||
|
||||
/**
|
||||
* rewrite phase 2
|
||||
*/
|
||||
public class ClassOverrideVisitor extends ClassVisitor {
|
||||
|
||||
final String className;
|
||||
final Set<MethodId> rayRemoteMethods;
|
||||
MethodVisitor ClinitVisitor;
|
||||
|
||||
// init the static added field in <clinit>
|
||||
// static {
|
||||
// assign value to _hashOf_XXX
|
||||
// }
|
||||
class StaticBlockVisitor extends MethodVisitor {
|
||||
|
||||
StaticBlockVisitor(MethodVisitor mv) {
|
||||
super(Opcodes.ASM6, mv);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void visitCode() {
|
||||
super.visitCode();
|
||||
|
||||
// assign value for added hash fields within <clinit>
|
||||
for (MethodId m : rayRemoteMethods) {
|
||||
byte[] hash = m.getSha1Hash();
|
||||
insertByteArray(hash);
|
||||
mv.visitFieldInsn(Opcodes.PUTSTATIC, className, m.getStaticHashValueFieldName(), "[B");
|
||||
|
||||
System.out.println("assign field: " + m.getStaticHashValueFieldName() + " = " + MethodId
|
||||
.toHexHashString(hash));
|
||||
}
|
||||
}
|
||||
|
||||
private void insertByteArray(byte[] bytes) {
|
||||
int length = bytes.length;
|
||||
assert (length < Short.MAX_VALUE);
|
||||
mv.visitIntInsn(Opcodes.SIPUSH, length);
|
||||
mv.visitIntInsn(Opcodes.NEWARRAY, Opcodes.T_BYTE);
|
||||
mv.visitInsn(Opcodes.DUP);
|
||||
for (int i = 0; i < length; ++i) {
|
||||
mv.visitIntInsn(Opcodes.BIPUSH, i);
|
||||
mv.visitIntInsn(Opcodes.BIPUSH, bytes[i]);
|
||||
mv.visitInsn(Opcodes.BASTORE);
|
||||
if (i < (length - 1)) {
|
||||
mv.visitInsn(Opcodes.DUP);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
public ClassOverrideVisitor(ClassVisitor origin, String className,
|
||||
Set<MethodId> rayRemoteMethods) {
|
||||
super(Opcodes.ASM6, origin);
|
||||
this.className = className;
|
||||
this.rayRemoteMethods = rayRemoteMethods;
|
||||
this.ClinitVisitor = null;
|
||||
}
|
||||
|
||||
@Override
|
||||
public MethodVisitor visitMethod(int access, String name, String desc, String signature,
|
||||
String[] exceptions) {
|
||||
if ("<clinit>".equals(name) && ClinitVisitor == null) {
|
||||
MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
|
||||
ClinitVisitor = new StaticBlockVisitor(mv);
|
||||
return ClinitVisitor;// dispatch the ASM modifications (assign values to the preComputedxxx static field) to the ClinitVisitor
|
||||
}
|
||||
|
||||
ClassVisitor this_ = this;
|
||||
MethodId m = new MethodId(className, name, desc, (access & Opcodes.ACC_STATIC) != 0, null);
|
||||
if (rayRemoteMethods.contains(m)) {
|
||||
if (m.isStaticMethod()) {
|
||||
return new MethodVisitor(api,
|
||||
super.visitMethod(access, name, desc, signature, exceptions)) {
|
||||
@Override
|
||||
public void visitCode() {
|
||||
// step 1: add a field for the function id of this method
|
||||
System.out.println("add field: " + m.getStaticHashValueFieldName());
|
||||
String fieldName = m.getStaticHashValueFieldName();
|
||||
this_.visitField(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, fieldName, "[B",
|
||||
null, null);
|
||||
|
||||
// step 2: rewrite current method so if MethodSwitcher returns true, returns the added function id directly
|
||||
// else call the original method
|
||||
mv.visitFieldInsn(Opcodes.GETSTATIC, className, fieldName, "[B");
|
||||
mv.visitMethodInsn(Opcodes.INVOKESTATIC, "org/ray/hook/runtime/MethodSwitcher",
|
||||
"execute",
|
||||
"([B)Z", false);
|
||||
Label dorealwork = new Label();
|
||||
mv.visitJumpInsn(Opcodes.IFEQ, dorealwork);
|
||||
properReturn(sayReturnType(desc), mv);// proper return on
|
||||
// different types
|
||||
mv.visitLabel(dorealwork);
|
||||
mv.visitCode();// real work
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
// non-static
|
||||
else {
|
||||
return super.visitMethod(access, name, desc, signature, exceptions);
|
||||
}
|
||||
} else {
|
||||
return super.visitMethod(access, name, desc, signature, exceptions);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public void visitEnd() {
|
||||
if (ClinitVisitor == null) { // works fine
|
||||
// Create an empty static block and let our method
|
||||
// visitor modify it the same way it modifies an
|
||||
// existing static block
|
||||
{
|
||||
MethodVisitor mv = super.visitMethod(Opcodes.ACC_STATIC, "<clinit>", "()V", null, null);
|
||||
mv = new StaticBlockVisitor(mv);
|
||||
mv.visitCode();
|
||||
mv.visitInsn(Opcodes.RETURN);
|
||||
mv.visitMaxs(0, 0);
|
||||
mv.visitEnd();
|
||||
}
|
||||
}
|
||||
|
||||
// for each non-static method, create a method for returning hash
|
||||
for (MethodId mid : this.rayRemoteMethods) {
|
||||
if (!mid.isStaticMethod()) {
|
||||
|
||||
// step 1: create a new method called method_name_function_id()
|
||||
System.out.println("add method: " + mid.getIdMethodName());
|
||||
MethodVisitor mv = super.visitMethod(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC,
|
||||
mid.getIdMethodName(), mid.getIdMethodDesc(), null, null);
|
||||
mv.visitCode();
|
||||
Label l0 = new Label();
|
||||
mv.visitLabel(l0);
|
||||
|
||||
// step 2: add a new static field as the function id of this method
|
||||
System.out.println("add field: " + mid.getStaticHashValueFieldName());
|
||||
String fieldName = mid.getStaticHashValueFieldName();
|
||||
this.visitField(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, fieldName, "[B",
|
||||
null, null);
|
||||
mv.visitFieldInsn(Opcodes.GETSTATIC, className, fieldName, "[B");
|
||||
|
||||
// step 3: call method switcher, and returns the function id when the mode is rewritten
|
||||
mv.visitMethodInsn(Opcodes.INVOKESTATIC, "org/ray/hook/runtime/MethodSwitcher", "execute",
|
||||
"([B)Z", false);
|
||||
mv.visitInsn(Opcodes.POP);
|
||||
Label l1 = new Label();
|
||||
mv.visitLabel(l1);
|
||||
properReturn(sayReturnType(mid.getMethodDesc()), mv);
|
||||
|
||||
Label l2 = new Label();
|
||||
mv.visitLabel(l2);
|
||||
|
||||
org.objectweb.asm.Type[] args = org.objectweb.asm.Type
|
||||
.getArgumentTypes(mid.getIdMethodDesc());
|
||||
int argCount = args.length;
|
||||
/*
|
||||
for (int i = 0; i < argCount; ++i) {
|
||||
String ldsptr = args[i].getDescriptor();
|
||||
if (!ldsptr.endsWith(";"))
|
||||
ldsptr = "L" + ldsptr + ";";
|
||||
mv.visitLocalVariable("arg" + i, ldsptr, null, l0, l2, i);
|
||||
}
|
||||
*/
|
||||
mv.visitMaxs(2, argCount);
|
||||
mv.visitEnd();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private void properReturn(String returnType, MethodVisitor mv) {
|
||||
int returnCode;
|
||||
Object returnValue = null;
|
||||
switch (returnType) {
|
||||
case "V":
|
||||
mv.visitInsn(Opcodes.RETURN);
|
||||
return;
|
||||
case "I":
|
||||
case "B":
|
||||
case "S":
|
||||
case "Z": // int byte short boolean
|
||||
returnCode = Opcodes.IRETURN;
|
||||
returnValue = 0;
|
||||
break;
|
||||
case "J": // long
|
||||
returnCode = Opcodes.LRETURN;
|
||||
returnValue = 0L;
|
||||
break;
|
||||
case "D": // double
|
||||
returnCode = Opcodes.DRETURN;
|
||||
returnValue = 0D;
|
||||
break;
|
||||
case "F": // float
|
||||
returnCode = Opcodes.FRETURN;
|
||||
returnValue = 0F;
|
||||
break;
|
||||
default: // reference
|
||||
returnCode = Opcodes.ARETURN;
|
||||
break;
|
||||
}
|
||||
if (returnValue != null) {
|
||||
mv.visitLdcInsn(returnValue);
|
||||
} else {
|
||||
mv.visitInsn(Opcodes.ACONST_NULL);
|
||||
}
|
||||
mv.visitInsn(returnCode);
|
||||
}
|
||||
|
||||
private String sayReturnType(String str) {
|
||||
int left = str.lastIndexOf(")") + 1;
|
||||
return str.substring(left);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,201 @@
|
||||
package org.ray.hook;
|
||||
|
||||
import java.io.File;
|
||||
import java.io.FileNotFoundException;
|
||||
import java.io.FileOutputStream;
|
||||
import java.io.IOException;
|
||||
import java.io.PrintWriter;
|
||||
import java.lang.reflect.Method;
|
||||
import java.security.NoSuchAlgorithmException;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collection;
|
||||
import java.util.Enumeration;
|
||||
import java.util.List;
|
||||
import java.util.Scanner;
|
||||
import java.util.function.BiConsumer;
|
||||
import java.util.jar.JarEntry;
|
||||
import java.util.jar.JarFile;
|
||||
import java.util.jar.JarOutputStream;
|
||||
import java.util.zip.DataFormatException;
|
||||
import org.apache.commons.io.FileUtils;
|
||||
import org.apache.commons.io.IOUtils;
|
||||
import org.apache.commons.io.filefilter.DirectoryFileFilter;
|
||||
import org.apache.commons.io.filefilter.RegexFileFilter;
|
||||
import org.ray.hook.runtime.JarLoader;
|
||||
import org.ray.hook.runtime.LoadedFunctions;
|
||||
import org.ray.util.logger.RayLog;
|
||||
|
||||
/**
|
||||
* rewrite jars to new jars with methods marked using Ray annotations
|
||||
*/
|
||||
public class JarRewriter {
|
||||
|
||||
private static final String FUNCTIONS_FILE = "ray.functions.txt";
|
||||
|
||||
public static void main(String[] args)
|
||||
throws IOException, SecurityException, DataFormatException {
|
||||
if (args.length == 1) {
|
||||
LoadedFunctions funcs = load(args[0], null);
|
||||
for (MethodId mi : funcs.functions) {
|
||||
System.err.println(mi.getIdMethodDesc());
|
||||
Method m = mi.load();
|
||||
String logInfo = "load: " + m.getDeclaringClass().getName() + "." + m.getName();
|
||||
RayLog.core.info(logInfo);
|
||||
}
|
||||
return;
|
||||
} else if (args.length < 2) {
|
||||
System.err.println("org.ray.hook.JarRewriter source-jar-dir dest-jar-dir");
|
||||
System.exit(1);
|
||||
}
|
||||
|
||||
rewrite(args[0], args[1]);
|
||||
}
|
||||
|
||||
public static void rewrite(String fromDir, String toDir) throws IOException, DataFormatException {
|
||||
File fromDirFile = new File(fromDir);
|
||||
File toDirFileTmp = new File(toDir + ".tmp");
|
||||
File toDirFile = new File(toDir);
|
||||
|
||||
File[] topFiles = fromDirFile.listFiles();
|
||||
if (topFiles.length != 1 || !topFiles[0].isDirectory()) {
|
||||
throw new DataFormatException("There should be a top dir in the Ray app zip file.");
|
||||
}
|
||||
String topDir = topFiles[0].getName();
|
||||
|
||||
if (toDirFileTmp.exists()) {
|
||||
FileUtils.deleteDirectory(toDirFileTmp);
|
||||
}
|
||||
//toDirFileTmp.mkdir();
|
||||
FileUtils.copyDirectory(fromDirFile, toDirFileTmp);
|
||||
|
||||
PrintWriter functionCollector = new PrintWriter(toDir + ".tmp/" + FUNCTIONS_FILE, "UTF-8");
|
||||
|
||||
// get all jars
|
||||
Collection<File> files = FileUtils.listFiles(
|
||||
fromDirFile,
|
||||
new RegexFileFilter(".*\\.jar"),
|
||||
DirectoryFileFilter.DIRECTORY
|
||||
);
|
||||
|
||||
// load and rewrite
|
||||
int prefixLength = fromDirFile.getAbsolutePath().length() + topDir.length() + 2;
|
||||
for (File appJar : files) {
|
||||
String fromPath = appJar.getAbsolutePath();
|
||||
if (fromPath.substring(prefixLength).contains("/")) {
|
||||
functionCollector.close();
|
||||
throw new DataFormatException("There should not be any subdir"
|
||||
+ " containing jar file in the top dir of the Ray app zip file.");
|
||||
}
|
||||
JarFile jar = new JarFile(appJar.getAbsolutePath());
|
||||
String to = fromPath
|
||||
.replaceFirst(fromDirFile.getAbsolutePath(), toDirFileTmp.getAbsolutePath());
|
||||
rewrite(jar, to, (l, m) -> functionCollector.println(m.toEncodingString()));
|
||||
jar.close();
|
||||
}
|
||||
|
||||
// rename the whole dir
|
||||
functionCollector.close();
|
||||
|
||||
if (toDirFile.exists()) {
|
||||
FileUtils.deleteDirectory(toDirFile);
|
||||
}
|
||||
|
||||
FileUtils.moveDirectory(toDirFileTmp, toDirFile);
|
||||
}
|
||||
|
||||
public static LoadedFunctions load(String dir, String baseDir)
|
||||
throws FileNotFoundException, SecurityException {
|
||||
List<String> functions = JarRewriter.getRewrittenFunctions(dir);
|
||||
LoadedFunctions efuncs = new LoadedFunctions();
|
||||
efuncs.loader = JarLoader.loadJars(dir, false);
|
||||
|
||||
for (String func : functions) {
|
||||
MethodId mid = new MethodId(func, efuncs.loader);
|
||||
efuncs.functions.add(mid);
|
||||
}
|
||||
|
||||
if (baseDir != null && !baseDir.equals("")) {
|
||||
List<String> baseFunctions = JarRewriter.getRewrittenFunctions(baseDir);
|
||||
for (String func : baseFunctions) {
|
||||
MethodId mid = new MethodId(func, efuncs.loader);
|
||||
efuncs.functions.add(mid);
|
||||
}
|
||||
}
|
||||
|
||||
return efuncs;
|
||||
}
|
||||
|
||||
public static LoadedFunctions loadBase(String baseDir)
|
||||
throws FileNotFoundException, SecurityException {
|
||||
List<String> functions = JarRewriter.getRewrittenFunctions(baseDir);
|
||||
LoadedFunctions efuncs = new LoadedFunctions();
|
||||
efuncs.loader = null;
|
||||
|
||||
for (String func : functions) {
|
||||
MethodId mid = new MethodId(func, efuncs.loader);
|
||||
efuncs.functions.add(mid);
|
||||
}
|
||||
|
||||
return efuncs;
|
||||
}
|
||||
|
||||
public static List<String> getRewrittenFunctions(String rewrittenDir)
|
||||
throws FileNotFoundException {
|
||||
ArrayList<String> functions = new ArrayList<>();
|
||||
Scanner s = new Scanner(new File(rewrittenDir + "/" + FUNCTIONS_FILE));
|
||||
while (s.hasNext()) {
|
||||
String f = s.next();
|
||||
if (!f.startsWith("(")) {
|
||||
functions.add(f);
|
||||
}
|
||||
}
|
||||
s.close();
|
||||
|
||||
return functions;
|
||||
}
|
||||
|
||||
public static void rewrite(JarFile from, String to, BiConsumer<ClassLoader, MethodId> consumer)
|
||||
throws IOException {
|
||||
|
||||
FileOutputStream ofStream = new FileOutputStream(to);
|
||||
JarOutputStream ojStream = new JarOutputStream(ofStream);
|
||||
Enumeration<JarEntry> e = from.entries();
|
||||
String className;
|
||||
|
||||
while (e.hasMoreElements()) {
|
||||
JarEntry je = e.nextElement();
|
||||
byte[] jeBytes = IOUtils.toByteArray(from.getInputStream(je));
|
||||
|
||||
//System.err.println("XXXXXX " + from.getName() + " :: " + je.getName());
|
||||
if (!je.isDirectory() && je.getName().endsWith(".class")) {
|
||||
className = je.getName().substring(0, je.getName().length() - ".class".length());
|
||||
|
||||
//System.err.println("XXXXXX " + from.getName() + " :: " + je.getName() + " - " + className);
|
||||
ClassAdapter.Result result = ClassAdapter.hookClass(null, className, jeBytes);
|
||||
if (result.classBuffer != jeBytes) {
|
||||
String logInfo = "Rewrite class " + className + " from " + jeBytes.length + " bytes to "
|
||||
+ result.classBuffer.length + " bytes ";
|
||||
RayLog.core.info(logInfo);
|
||||
}
|
||||
|
||||
if (result.changedMethods != null) {
|
||||
for (MethodId m : result.changedMethods) {
|
||||
consumer.accept(null, m);
|
||||
}
|
||||
}
|
||||
|
||||
je = new JarEntry(je.getName());
|
||||
je.setTime(System.currentTimeMillis());
|
||||
je.setSize(result.classBuffer.length);
|
||||
jeBytes = result.classBuffer;
|
||||
}
|
||||
|
||||
ojStream.putNextEntry(je);
|
||||
ojStream.write(jeBytes);
|
||||
//ojStream.closeEntry();
|
||||
}
|
||||
|
||||
ojStream.close();
|
||||
ofStream.close();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
package org.ray.hook;
|
||||
|
||||
import java.lang.reflect.Field;
|
||||
import java.lang.reflect.Method;
|
||||
import java.nio.ByteBuffer;
|
||||
import java.nio.ByteOrder;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import org.apache.commons.codec.digest.DigestUtils;
|
||||
import org.objectweb.asm.Type;
|
||||
import org.ray.util.logger.RayLog;
|
||||
|
||||
/**
|
||||
* Represent a Method in a Class
|
||||
*/
|
||||
public class MethodId {
|
||||
|
||||
String className;
|
||||
String methodName;
|
||||
String methodDesc;
|
||||
boolean isStatic;
|
||||
ClassLoader loader;
|
||||
|
||||
static final String getFunctionIdPostfix = "_function_id";
|
||||
|
||||
public MethodId(String cls, String method, String mdesc, boolean isstatic, ClassLoader loader_) {
|
||||
className = cls;
|
||||
methodName = method;
|
||||
methodDesc = mdesc;
|
||||
isStatic = isstatic;
|
||||
loader = loader_;
|
||||
}
|
||||
|
||||
public MethodId(String encodedString, ClassLoader loader_) {
|
||||
// className + "." + methodName + "::" + methodDesc + "&&" + isStatic;
|
||||
int lastPos3 = encodedString.lastIndexOf("&&");
|
||||
int lastPos2 = encodedString.lastIndexOf("::");
|
||||
int lastPos1 = encodedString.lastIndexOf(".");
|
||||
if (lastPos1 == -1 || lastPos2 == -1 || lastPos3 == -1) {
|
||||
throw new RuntimeException("invalid given method id " + encodedString
|
||||
+ " - it must be className.methodName::methodDesc&&isStatic");
|
||||
}
|
||||
|
||||
className = encodedString.substring(0, lastPos1);
|
||||
methodName = encodedString.substring(lastPos1 + ".".length(), lastPos2);
|
||||
methodDesc = encodedString.substring(lastPos2 + "::".length(), lastPos3);
|
||||
isStatic = Boolean.parseBoolean(encodedString.substring(lastPos3 + "&&".length()));
|
||||
loader = loader_;
|
||||
}
|
||||
|
||||
public String getClassName() {
|
||||
return className;
|
||||
}
|
||||
|
||||
public String getMethodName() {
|
||||
return methodName;
|
||||
}
|
||||
|
||||
public String getMethodDesc() {
|
||||
return methodDesc;
|
||||
}
|
||||
|
||||
public ClassLoader getLoader() {
|
||||
return loader;
|
||||
}
|
||||
|
||||
public Boolean isStaticMethod() {
|
||||
return isStatic;
|
||||
}
|
||||
|
||||
public String getIdMethodName() {
|
||||
return this.methodName + getFunctionIdPostfix;
|
||||
}
|
||||
|
||||
public String getIdMethodDesc() {
|
||||
return "(L" + this.className + ";" + this.methodDesc.substring(1);
|
||||
}
|
||||
|
||||
public static String toHexHashString(byte[] id) {
|
||||
String s = "";
|
||||
String hex = "0123456789abcdef";
|
||||
assert (id.length == 20);
|
||||
for (int i = 0; i < 20; i++) {
|
||||
int val = id[i] & 0xff;
|
||||
s += hex.charAt(val >> 4);
|
||||
s += hex.charAt(val & 0xf);
|
||||
}
|
||||
return s;
|
||||
}
|
||||
|
||||
private String toHexHashString() {
|
||||
byte[] id = this.getSha1Hash();
|
||||
return toHexHashString(id);
|
||||
}
|
||||
|
||||
public Method load() {
|
||||
String loadClsName = className.replace('/', '.');
|
||||
Class<?> cls;
|
||||
try {
|
||||
RayLog.core.debug(
|
||||
"load class " + loadClsName + " from class loader " + (loader == null ? this.getClass()
|
||||
.getClassLoader() : loader)
|
||||
+ " for method " + toString() + " with ID = " + toHexHashString()
|
||||
);
|
||||
cls = Class
|
||||
.forName(loadClsName, true, loader == null ? this.getClass().getClassLoader() : loader);
|
||||
} catch (Throwable e) {
|
||||
RayLog.core.error("Cannot load class " + loadClsName, e);
|
||||
return null;
|
||||
}
|
||||
|
||||
Method[] ms = cls.getDeclaredMethods();
|
||||
ArrayList<Method> methods = new ArrayList<>();
|
||||
Type t = Type.getMethodType(this.methodDesc);
|
||||
Type[] params = t.getArgumentTypes();
|
||||
String rt = t.getReturnType().getDescriptor();
|
||||
|
||||
for (Method m : ms) {
|
||||
if (m.getName().equals(methodName)) {
|
||||
if (!Arrays.equals(params, Type.getArgumentTypes(m))) {
|
||||
continue;
|
||||
}
|
||||
|
||||
String mrt = Type.getDescriptor(m.getReturnType());
|
||||
if (!rt.equals(mrt)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
methods.add(m);
|
||||
}
|
||||
}
|
||||
|
||||
if (methods.size() != 1) {
|
||||
RayLog.core.error(
|
||||
"Load method " + toString() + " failed as there are " + methods.size() + " definitions");
|
||||
return null;
|
||||
}
|
||||
|
||||
Method m = methods.get(0);
|
||||
try {
|
||||
Field fld = cls.getField(getStaticHashValueFieldName());
|
||||
Object hashValue = fld.get(null);
|
||||
if (hashValue instanceof byte[] && Arrays.equals((byte[]) hashValue, this.getSha1Hash())) {
|
||||
RayLog.core.debug("Method " + toString() + " hash: " + toHexHashString((byte[]) hashValue));
|
||||
} else {
|
||||
if (hashValue instanceof byte[]) {
|
||||
RayLog.core.error(
|
||||
"Method " + toString() + " hash-field: " + toHexHashString((byte[]) hashValue)
|
||||
+ " vs id-hash: " + toHexHashString());
|
||||
} else {
|
||||
RayLog.core.error(
|
||||
"Method " + toString() + " hash-field: " + (hashValue != null ? hashValue.toString()
|
||||
: "<nil>") + " vs id-hash: " + toHexHashString());
|
||||
}
|
||||
}
|
||||
} catch (NoSuchFieldException | SecurityException | IllegalArgumentException | IllegalAccessException e) {
|
||||
RayLog.core.error("load method hash field failed for " + toString(), e);
|
||||
}
|
||||
return m;
|
||||
}
|
||||
|
||||
public String toEncodingString() {
|
||||
return className + "." + methodName + "::" + methodDesc + "&&" + isStatic;
|
||||
}
|
||||
|
||||
public byte[] getSha1Hash() {
|
||||
byte[] digests = DigestUtils.sha(toEncodingString());
|
||||
ByteBuffer bb = ByteBuffer.wrap(digests);
|
||||
bb.order(ByteOrder.LITTLE_ENDIAN);
|
||||
if (methodName.contains("createActorStage1")) {
|
||||
bb.putLong(Long.BYTES, 1);
|
||||
} else {
|
||||
bb.putLong(Long.BYTES, 0);
|
||||
}
|
||||
return digests;
|
||||
}
|
||||
|
||||
public String getStaticHashValueFieldName() {
|
||||
// _hashOf<init>_([Ljava/lang/String;)V
|
||||
String r = "_hashOf" + methodName + "_" + methodDesc;
|
||||
r = r.replace("<", "_")
|
||||
.replace(">", "_")
|
||||
.replace("(", "_")
|
||||
.replace("[", "_")
|
||||
.replace("/", "_")
|
||||
.replace(";", "_")
|
||||
.replace(")", "_")
|
||||
;
|
||||
// System.err.println(r);
|
||||
return r;
|
||||
}
|
||||
|
||||
@Override
|
||||
public int hashCode() {
|
||||
final int prime = 31;
|
||||
int result = 1;
|
||||
result = prime * result + className.hashCode();
|
||||
result = prime * result + methodName.hashCode();
|
||||
result = prime * result + methodDesc.hashCode();
|
||||
return result;
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean equals(Object obj) {
|
||||
if (this == obj) {
|
||||
return true;
|
||||
}
|
||||
if (obj == null) {
|
||||
return false;
|
||||
}
|
||||
if (getClass() != obj.getClass()) {
|
||||
return false;
|
||||
}
|
||||
MethodId other = (MethodId) obj;
|
||||
return className.equals(other.className)
|
||||
&& methodName.equals(other.methodName)
|
||||
&& methodDesc.equals(other.methodDesc)
|
||||
&& isStatic == other.isStatic
|
||||
;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String toString() {
|
||||
return toEncodingString();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
package org.ray.hook.runtime;
|
||||
|
||||
import java.io.File;
|
||||
import java.io.IOException;
|
||||
import java.lang.reflect.Method;
|
||||
import java.net.URL;
|
||||
import java.net.URLClassLoader;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collection;
|
||||
import java.util.Enumeration;
|
||||
import java.util.List;
|
||||
import java.util.jar.JarEntry;
|
||||
import java.util.jar.JarFile;
|
||||
import org.apache.commons.io.FileUtils;
|
||||
import org.apache.commons.io.filefilter.DirectoryFileFilter;
|
||||
import org.apache.commons.io.filefilter.RegexFileFilter;
|
||||
import org.ray.util.logger.RayLog;
|
||||
|
||||
/**
|
||||
* load and unload jars from a dir
|
||||
*/
|
||||
public class JarLoader {
|
||||
|
||||
private static Method AddUrl = InitAddUrl();
|
||||
|
||||
private static Method InitAddUrl() {
|
||||
try {
|
||||
Method m = URLClassLoader.class.getDeclaredMethod("addURL", URL.class);
|
||||
m.setAccessible(true);
|
||||
return m;
|
||||
} catch (NoSuchMethodException | SecurityException e) {
|
||||
// TODO Auto-generated catch block
|
||||
e.printStackTrace();
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
public static URLClassLoader loadJars(String dir, boolean explicitLoadForHook) {
|
||||
// get all jars
|
||||
Collection<File> jars = FileUtils.listFiles(
|
||||
new File(dir),
|
||||
new RegexFileFilter(".*\\.jar"),
|
||||
DirectoryFileFilter.DIRECTORY
|
||||
);
|
||||
return loadJar(jars, explicitLoadForHook);
|
||||
}
|
||||
|
||||
public static ClassLoader loadJars(String[] appJars, boolean explicitLoadForHook) {
|
||||
List<File> jars = new ArrayList<>();
|
||||
|
||||
for (String jar : appJars) {
|
||||
if (jar.endsWith(".jar")) {
|
||||
jars.add(new File(jar));
|
||||
} else {
|
||||
loadJarDir(jar, jars);
|
||||
}
|
||||
}
|
||||
|
||||
return loadJar(jars, explicitLoadForHook);
|
||||
}
|
||||
|
||||
public static void unloadJars(ClassLoader loader) {
|
||||
// TODO:
|
||||
}
|
||||
|
||||
private static URLClassLoader loadJar(Collection<File> appJars, boolean explicitLoadForHook) {
|
||||
List<JarFile> jars = new ArrayList<>();
|
||||
List<URL> urls = new ArrayList<>();
|
||||
|
||||
for (File appJar : appJars) {
|
||||
try {
|
||||
RayLog.core.info("load jar " + appJar.getAbsolutePath() + " in ray hook");
|
||||
JarFile jar = new JarFile(appJar.getAbsolutePath());
|
||||
jars.add(jar);
|
||||
urls.add(appJar.toURI().toURL());
|
||||
} catch (IOException e) {
|
||||
e.printStackTrace();
|
||||
System.err.println(
|
||||
"invalid app jar path: " + appJar.getAbsolutePath() + ", load failed with exception "
|
||||
+ e.getMessage());
|
||||
System.exit(1);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
URLClassLoader cl = URLClassLoader.newInstance(urls.toArray(new URL[0]));
|
||||
|
||||
if (!explicitLoadForHook) {
|
||||
return cl;
|
||||
}
|
||||
|
||||
for (JarFile jar : jars) {
|
||||
Enumeration<JarEntry> e = jar.entries();
|
||||
while (e.hasMoreElements()) {
|
||||
JarEntry je = e.nextElement();
|
||||
//System.err.println("check " + je.getName());
|
||||
if (je.isDirectory() || !je.getName().endsWith(".class")) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// -6 because of .class
|
||||
String className = je.getName().substring(0, je.getName().length() - 6);
|
||||
className = className.replace('/', '.');
|
||||
try {
|
||||
Class.forName(className, true, cl);
|
||||
//System.err.println("load class " + className + " OK");
|
||||
} catch (ClassNotFoundException e1) {
|
||||
e1.printStackTrace();
|
||||
}
|
||||
}
|
||||
|
||||
try {
|
||||
jar.close();
|
||||
} catch (IOException e1) {
|
||||
// TODO Auto-generated catch block
|
||||
e1.printStackTrace();
|
||||
}
|
||||
}
|
||||
return cl;
|
||||
}
|
||||
|
||||
private static void loadJarDir(String jarDir, List<File> jars) {
|
||||
Collection<File> files = FileUtils.listFiles(
|
||||
new File(jarDir),
|
||||
new RegexFileFilter(".*\\.jar"),
|
||||
DirectoryFileFilter.DIRECTORY
|
||||
);
|
||||
|
||||
jars.addAll(files);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
package org.ray.hook.runtime;
|
||||
|
||||
import java.util.Collections;
|
||||
import java.util.HashSet;
|
||||
import java.util.Set;
|
||||
import org.ray.hook.MethodId;
|
||||
|
||||
public class LoadedFunctions {
|
||||
|
||||
public ClassLoader loader = null;
|
||||
public final Set<MethodId> functions = Collections.synchronizedSet(new HashSet<>());
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
package org.ray.hook.runtime;
|
||||
|
||||
import java.util.Arrays;
|
||||
|
||||
public class MethodHash {
|
||||
|
||||
private final byte[] hash;
|
||||
|
||||
public MethodHash(byte[] hash) {
|
||||
this.hash = hash;
|
||||
}
|
||||
|
||||
public byte[] getHash() {
|
||||
return hash;
|
||||
}
|
||||
|
||||
@Override
|
||||
public int hashCode() {
|
||||
return Arrays.hashCode(getHash());
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean equals(Object obj) {
|
||||
if (this == obj) {
|
||||
return true;
|
||||
}
|
||||
if (obj == null) {
|
||||
return false;
|
||||
}
|
||||
if (getClass() != obj.getClass()) {
|
||||
return false;
|
||||
}
|
||||
MethodHash other = (MethodHash) obj;
|
||||
return Arrays.equals(this.getHash(), other.getHash());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
package org.ray.hook.runtime;
|
||||
|
||||
/**
|
||||
* method mode switch at runtime
|
||||
*/
|
||||
public class MethodSwitcher {
|
||||
|
||||
public static final ThreadLocal<Boolean> IsRemoteCall = new ThreadLocal<>();
|
||||
|
||||
public static final ThreadLocal<byte[]> MethodId = new ThreadLocal<>();
|
||||
|
||||
public static boolean execute(byte[] id) {
|
||||
Boolean hooking = IsRemoteCall.get();
|
||||
if (Boolean.TRUE.equals(hooking)) {
|
||||
MethodId.set(id);
|
||||
return true;
|
||||
} else {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user