All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.github.lab515.utils.ByteShifter Maven / Gradle / Ivy

There is a newer version: 1.0.12
Show newest version
package io.github.lab515.utils;

import jdk.internal.org.objectweb.asm.*;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.concurrent.atomic.AtomicInteger;

public class ByteShifter extends ClassVisitor {
    private static final String runnerCls = ByteRunner.class.getCanonicalName().replace('.', '/');
    private static final AtomicInteger id = new AtomicInteger(0);
    static Interceptor[] allInterceptors = null;
    private String clsName;
    private int[] fits; // simple impl for now
    private boolean valid = false;
    private ByteShifter(int[] fits){
        super(Opcodes.ASM5,null);
        this.fits = fits;
    }
    public static byte[] readClassAllBytes(Class target) {
        return readClassAllBytes(target.getClassLoader(), target.getName().replace('.', '/') + ".class");
    }
    public static byte[] readClassAllBytes(ClassLoader loader, String path) {
        try(InputStream inp = loader.getResourceAsStream(path)){
            if(inp == null){
                return null;
            }
            ByteArrayOutputStream bto = new ByteArrayOutputStream();
            byte[] bytes = new byte[4096];
            int l = 0;
            while((l = inp.read(bytes,0,bytes.length)) > 0){
                bto.write(bytes,0,l);
            }
            bytes = bto.toByteArray();
            return bytes;
        }catch (IOException e){
            return null;
        }
    }

    public static byte[] doInstrument(String className, byte[] buf,  Class redefiner){
        if(buf == null && redefiner != null){
            buf = readClassAllBytes(redefiner);
        }
        if(className == null || buf == null){
            return null;
        }
        int[] fits = new int[allInterceptors.length];
        String cls = className.replace('/','.');
        boolean doInstr = false;
        for(int i = 0; i < allInterceptors.length; i++){
            fits[i] = -1;
            if(!allInterceptors[i].withClass(cls)){
                continue;
            }
            doInstr = true;
            fits[i] = 1; // as true
        }
        if(doInstr){
            return ByteShifter.instr(buf,fits);
        }else{
            return null;
        }
    }

    public static byte[] instr(byte[] source, int[] fits){
        ClassReader reader = new ClassReader(source);
        ByteShifter ia = new ByteShifter(fits);
        reader.accept(ia, 0);
        if(ia.valid){
            return ((ClassWriter)ia.cv).toByteArray();
        }else{
            return null;
        }
    }

    @Override
    public void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
        if((access & (Opcodes.ACC_INTERFACE | Opcodes.ACC_ENUM)) != 0){
            return;
        }
        cv = new ClassWriter(0);
        cv.visit(version, access, name, signature, superName, interfaces);
        clsName = name;
    }
    private static final String tps = "ZBCSIJFD";
    private static final String[] ts = new String[]{"Z","B","C","S","I","J","F","D"};
    private static final String[] trs = new String[]{"()Z","()B","()C","()S","()I","()J","()F","()D"};
    private static final String[] tpds = new String[]{"java/lang/Boolean","java/lang/Byte","java/lang/Character","java/lang/Short","java/lang/Integer","java/lang/Long","java/lang/Float","java/lang/Double"};
    private static final String[] tpdss = new String[]{"(Z)Ljava/lang/Boolean;","(B)Ljava/lang/Byte;","(C)Ljava/lang/Character;","(S)Ljava/lang/Short;","(I)Ljava/lang/Integer;","(J)Ljava/lang/Long;","(F)Ljava/lang/Float;","(D)Ljava/lang/Double;"};
    private static final String[] tpms = new String[]{"booleanValue","byteValue","charValue","shortValue","intValue","longValue","floatValue","doubleValue"};
    private static final int[] tpis = new int[]{Opcodes.IRETURN, Opcodes.IRETURN,Opcodes.IRETURN,Opcodes.IRETURN,Opcodes.IRETURN,Opcodes.LRETURN,Opcodes.FRETURN,Opcodes.DRETURN};
    private static final int[] tpils = new int[]{Opcodes.ILOAD, Opcodes.ILOAD,Opcodes.ILOAD,Opcodes.ILOAD,Opcodes.ILOAD,Opcodes.LLOAD,Opcodes.FLOAD,Opcodes.DLOAD};

    // for now, simply instrument method
    @Override
    public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
        if(cv == null){
            return null;
        }
        MethodVisitor mv = cv.visitMethod(access, name, desc, signature, exceptions);
        if(name.startsWith("<")){
            return mv;
        }
        // produce method is the target
        return new MethodVisitor(Opcodes.ASM5, mv) {
            int stackSize = 0;
            StringBuilder annos = new StringBuilder();
            @Override
            public void visitCode() {
                mv.visitCode();
                int[] flags = new int[allInterceptors.length];
                if(flags.length != fits.length){
                    return;
                }
                int insts = 0;
                String stdCls = clsName.replace('/','.');
                String anns = annos.toString();
                for(int i = 0; i < flags.length;i++){
                    if(fits[i] < 0){
                        flags[i] = -1;
                    }else if((flags[i] = allInterceptors[i].withMethod(stdCls, name, desc, anns)) >= 0){
                        insts++;
                    }
                }
                if(insts < 1){
                    return;
                }
                valid = true;
                Label skip = new Label();
                mv.visitLdcInsn(insts * 2); // no need to use bipush or sipush
                mv.visitIntInsn(Opcodes.NEWARRAY, Opcodes.T_INT);
                insts = 0;
                for(int i = 0; i < flags.length;i++){
                    if(flags[i] >= 0){
                        mv.visitInsn(Opcodes.DUP);
                        mv.visitLdcInsn(insts++);
                        mv.visitLdcInsn(i);
                        mv.visitInsn(Opcodes.IASTORE);
                        mv.visitInsn(Opcodes.DUP);
                        mv.visitLdcInsn(insts++);
                        mv.visitLdcInsn(flags[i]);
                        mv.visitInsn(Opcodes.IASTORE);
                    }
                } // stack max 4, current 1
                mv.visitLdcInsn(Type.getObjectType(clsName));
                mv.visitLdcInsn(name);
                mv.visitLdcInsn(desc);
                mv.visitLdcInsn(id.getAndIncrement());
                mv.visitMethodInsn(Opcodes.INVOKESTATIC, runnerCls, "skip",
                        "([ILjava/lang/Class;Ljava/lang/String;Ljava/lang/String;I)Z",false); // max 5
                mv.visitJumpInsn(Opcodes.IFNE, skip);
                // process parameters
                int start = (Opcodes.ACC_STATIC & access) != 0 ? 0 : 1;
                if (start > 0 && !name.startsWith("<")) { // not a constructor
                    mv.visitVarInsn(Opcodes.ALOAD, 0); // the this pointer or null for
                }else {
                    mv.visitInsn(Opcodes.ACONST_NULL);
                }
                int paraLen = 0;
                char a;
                for (int i = 1; i < desc.length(); i++) {
                    a = desc.charAt(i);
                    if(a == ')') {
                        break;
                    }else if(a == '['){
                        continue;
                    }
                    if(a == 'L') {
                        i = desc.indexOf(';',i);
                    }
                    paraLen++;
                }
                mv.visitIntInsn(Opcodes.BIPUSH, paraLen);
                mv.visitTypeInsn(Opcodes.ANEWARRAY, "java/lang/Object");
                // fill in paramters
                int array = 0;
                paraLen = 0;
                boolean stackAware = false;
                boolean retPhase = false;
                for (int i = 1; i < desc.length(); i++) {
                    a = desc.charAt(i);
                    if(a == ')') {
                        mv.visitMethodInsn(Opcodes.INVOKESTATIC, runnerCls, "handle",
                                "(Ljava/lang/Object;[Ljava/lang/Object;)Ljava/lang/Object;",false);
                        retPhase = true;
                        continue;
                    }else if(a == '[') {
                        array++;
                        continue;
                    }
                    if(!retPhase) {
                        mv.visitInsn(Opcodes.DUP);
                        mv.visitIntInsn(Opcodes.BIPUSH, paraLen++);
                    }
                    if(array > 0 || a == 'L') {
                        if(retPhase){
                            if(!desc.endsWith(")Ljava/lang/Object;")){
                                mv.visitTypeInsn(Opcodes.CHECKCAST, array > 0 ? desc.substring( i - array) : desc.substring(i+1, desc.length()-1));
                            }
                            mv.visitInsn(Opcodes.ARETURN);
                        }else {
                            mv.visitVarInsn(Opcodes.ALOAD, start++);
                        }
                    }else {
                        int tp = tps.indexOf(a);
                        if(tp < 0){
                            if(retPhase){ // void
                                mv.visitInsn(Opcodes.POP);
                                mv.visitInsn(Opcodes.RETURN);
                            }else { // not likely
                                mv.visitInsn(Opcodes.ACONST_NULL);
                                start++;
                            }
                        }else{
                            if(retPhase){
                                mv.visitTypeInsn(Opcodes.CHECKCAST, tpds[tp]);
                                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, tpds[tp], tpms[tp],
                                        trs[tp], false);
                                mv.visitInsn(tpis[tp]);
                            }else {
                                mv.visitVarInsn(tpils[tp], start++);
                                if(a == 'J'||a=='D'){
                                    stackAware = true;
                                    start++;
                                }
                                mv.visitMethodInsn(Opcodes.INVOKESTATIC, tpds[tp], "valueOf",
                                        tpdss[tp], false);
                            }
                        }
                    }
                    if(!retPhase) {
                        mv.visitInsn(Opcodes.AASTORE);
                    }
                    array = 0;
                    if(a == 'L') {
                        i = desc.indexOf(';',i);
                    }
                }
                mv.visitLabel(skip);
                mv.visitFrame(Opcodes.F_SAME, 0, null, 0, null);
                stackSize = paraLen > 0 && stackAware ? 6 : 5; // long. double set in array require 4
                // if(agent.run()){
                //    return agent.incpt(cls,name, desc, args);
                // }


            }

            @Override
            public void visitMaxs(int maxStack, int maxLocals) {
                if (maxStack < stackSize) {
                    maxStack = stackSize;
                }
                if(mv != null) {
                    mv.visitMaxs(maxStack, maxLocals);
                }
            }
            @Override
            public void visitMethodInsn(int opcode, String owner, String name, String desc, boolean itf) {
                if(mv == null){
                    return;
                }
                if(owner.equals(runnerCls)){
                    // stopping
                    cv = null;
                    mv = null;
                    return;
                }
                mv.visitMethodInsn(opcode,owner,name,desc,itf);
            }

            @Override
            public AnnotationVisitor visitAnnotation(String desc, boolean visible){
                annos.append(desc);
                return mv.visitAnnotation(desc,visible);
            }
        };
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy