/*

 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.

 *

 * This code is free software; you can redistribute it and/or modify it

 * under the terms of the GNU General Public License version 2 only, as

 * published by the Free Software Foundation.

 *

 * This code is distributed in the hope that it will be useful, but WITHOUT

 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or

 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License

 * version 2 for more details (a copy is included in the LICENSE file that

 * accompanied this code).

 *

 * You should have received a copy of the GNU General Public License version

 * 2 along with this work; if not, write to the Free Software Foundation,

 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

 *

 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA

 * or visit www.oracle.com if you need additional information or have any

 * questions.

 */

package ilib;

import java.io.IOException;

import java.io.File;

import java.io.FileOutputStream;

import java.io.DataOutputStream;

import java.util.List;

import java.util.Iterator;

import java.util.ArrayList;

public class Inject implements RuntimeConstants {

    public static byte[] instrumentation(Options opt,

                                         ClassLoader loader,

                                         String className,

                                         byte[] classfileBuffer) {

        ClassReaderWriter c = new ClassReaderWriter(classfileBuffer);

        (new Inject(className, c, loader == null, opt)).doit();

        return c.result();

    }

    static boolean verbose = false;

    final String className;

    final ClassReaderWriter c;

    final boolean isSystem;

    final Options options;

    int constantPoolCount;

    int methodsCount;

    int methodsCountPos;

    int profiler;

    int wrappedTrackerIndex = 0;

    int thisClassIndex = 0;

    TrackerInjector callInjector;

    TrackerInjector allocInjector;

    TrackerInjector defaultInjector;

    static interface TrackerInjector extends Injector {

        void reinit(int tracker);

        int stackSize(int currentSize);

    }

    static class SimpleInjector implements TrackerInjector {

        byte[] injection;

        public int stackSize(int currentSize) {

            return currentSize;

        }

        public void reinit(int tracker) {

            injection = new byte[3];

            injection[0] = (byte)opc_invokestatic;

            injection[1] = (byte)(tracker >> 8);

            injection[2] = (byte)tracker;

        }

        public byte[] bytecodes(String className, String methodName, int location) {

            return injection;

        }

    }

    static class ObjectInjector implements TrackerInjector {

        byte[] injection;

        public int stackSize(int currentSize) {

            return currentSize + 1;

        }

        public void reinit(int tracker) {

            injection = new byte[4];

            injection[0] = (byte)opc_dup;

            injection[1] = (byte)opc_invokestatic;

            injection[2] = (byte)(tracker >> 8);

            injection[3] = (byte)tracker;

        }

        public byte[] bytecodes(String className, String methodName, int location) {

            return injection;

        }

    }

    class IndexedInjector implements TrackerInjector {

        int counter = 0;

        int tracker;

        List<Info> infoList = new ArrayList<>();

        public int stackSize(int currentSize) {

            return currentSize + 1;

        }

        public void reinit(int tracker) {

            this.tracker = tracker;

        }

        void dump(File outDir, String filename) throws IOException {

            try (FileOutputStream fileOut =

                     new FileOutputStream(new File(outDir, filename));

                 DataOutputStream dataOut = new DataOutputStream(fileOut))

            {

                String currentClassName = null;

                dataOut.writeInt(infoList.size());

                for (Iterator<Info> it = infoList.iterator(); it.hasNext(); ) {

                    Info info = it.next();

                    if (!info.className.equals(currentClassName)) {

                        dataOut.writeInt(123456); // class name marker

                        currentClassName = info.className;

                        dataOut.writeUTF(currentClassName);

                    }

                    dataOut.writeInt(info.location);

                    dataOut.writeUTF(info.methodName);

                }

            }

        }

        public byte[] bytecodes(String className, String methodName, int location) {

            byte[] injection = new byte[6];

            int injectedIndex = options.fixedIndex != 0? options.fixedIndex : ++counter;

            infoList.add(new Info(counter, className, methodName, location));

            injection[0] = (byte)opc_sipush;

            injection[1] = (byte)(injectedIndex >> 8);

            injection[2] = (byte)injectedIndex;

            injection[3] = (byte)opc_invokestatic;

            injection[4] = (byte)(tracker >> 8);

            injection[5] = (byte)tracker;

            return injection;

        }

    }

    Inject(String className, ClassReaderWriter c, boolean isSystem, Options options) {

        this.className = className;

        this.c = c;

        this.isSystem = isSystem;

        this.options = options;

    }

    void doit() {

        int i;

        c.copy(4 + 2 + 2); // magic min/maj version

        int constantPoolCountPos = c.generatedPosition();

        constantPoolCount = c.copyU2();

        // copy old constant pool

        c.copyConstantPool(constantPoolCount);

        if (verbose) {

            System.out.println("ConstantPool expanded from: " +

                               constantPoolCount);

        }

        profiler = addClassToConstantPool(options.trackerClassName);

        if (options.shouldInstrumentNew || options.shouldInstrumentObjectInit) {

            if (options.shouldInstrumentIndexed) {

                if (allocInjector == null) {

                    // first time - create it

                    allocInjector = new IndexedInjector();

                }

                int allocTracker = addMethodToConstantPool(profiler,

                                                           options.allocTrackerMethodName,

                                                           "(I)V");

                allocInjector.reinit(allocTracker);

            } else if (options.shouldInstrumentObject) {

                if (allocInjector == null) {

                    // first time - create it

                    allocInjector = new ObjectInjector();

                }

                int allocTracker = addMethodToConstantPool(profiler,

                                                           options.allocTrackerMethodName,

                                                           "(Ljava/lang/Object;)V");

                allocInjector.reinit(allocTracker);

            } else {

                if (allocInjector == null) {

                    // first time - create it

                    allocInjector = new SimpleInjector();

                }

                int allocTracker = addMethodToConstantPool(profiler,

                                                           options.allocTrackerMethodName,

                                                           "()V");

                allocInjector.reinit(allocTracker);

            }

            defaultInjector = allocInjector;

        }

        if (options.shouldInstrumentCall) {

            if (options.shouldInstrumentIndexed) {

                if (callInjector == null) {

                    // first time - create it

                    callInjector = new IndexedInjector();

                }

                int callTracker = addMethodToConstantPool(profiler,

                                                          options.callTrackerMethodName,

                                                          "(I)V");

                callInjector.reinit(callTracker);

            } else {

                if (callInjector == null) {

                    // first time - create it

                    callInjector = new SimpleInjector();

                }

                int callTracker = addMethodToConstantPool(profiler,

                                                          options.callTrackerMethodName,

                                                          "()V");

                callInjector.reinit(callTracker);

            }

            defaultInjector = callInjector;

        }

        if (verbose) {

            System.out.println("To: " + constantPoolCount);

        }

        c.setSection(1);

        c.copy(2 + 2 + 2);  // access, this, super

        int interfaceCount = c.copyU2();

        if (verbose) {

            System.out.println("interfaceCount: " + interfaceCount);

        }

        c.copy(interfaceCount * 2);

        copyFields(); // fields

        copyMethods(); // methods

        int attrCountPos = c.generatedPosition();

        int attrCount = c.copyU2();

        if (verbose) {

            System.out.println("class attrCount: " + attrCount);

        }

        // copy the class attributes

        copyAttrs(attrCount);

        c.randomAccessWriteU2(constantPoolCountPos, constantPoolCount);

    }

    void copyFields() {

        int count = c.copyU2();

        if (verbose) {

            System.out.println("fields count: " + count);

        }

        for (int i = 0; i < count; ++i) {

            c.copy(6); // access, name, descriptor

            int attrCount = c.copyU2();

            if (verbose) {

                System.out.println("field attr count: " + attrCount);

            }

            copyAttrs(attrCount);

        }

    }

    void copyMethods() {

        methodsCountPos = c.generatedPosition();

        methodsCount = c.copyU2();

        int initialMethodsCount = methodsCount;

        if (verbose) {

            System.out.println("methods count: " + methodsCount);

        }

        for (int i = 0; i < initialMethodsCount; ++i) {

            copyMethod();

        }

    }

    void copyMethod() {

        int accessFlags = c.copyU2();// access flags

        if (options.shouldInstrumentNativeMethods && (accessFlags & ACC_NATIVE) != 0) {

            wrapNativeMethod(accessFlags);

            return;

        }

        int nameIndex = c.copyU2();  // name

        String methodName = c.constantPoolString(nameIndex);

        c.copyU2();                  // descriptor

        int attrCount = c.copyU2();  // attribute count

        if (verbose) {

            System.out.println("methods attr count: " + attrCount);

        }

        for (int i = 0; i < attrCount; ++i) {

            copyAttrForMethod(methodName, accessFlags);

        }

    }

    void wrapNativeMethod(int accessFlags) {

        // first, copy the native method with the name changed

        // accessFlags have already been copied

        int nameIndex = c.readU2();        // name

        String methodName = c.constantPoolString(nameIndex);

        String wrappedMethodName = options.wrappedPrefix + methodName;

        int wrappedNameIndex = writeCPEntryUtf8(wrappedMethodName);

        c.writeU2(wrappedNameIndex);       // change to the wrapped name

        int descriptorIndex = c.copyU2();  // descriptor index

        int attrCount = c.copyU2();        // attribute count

        // need to replicate these attributes (esp Exceptions) in wrapper

        // so mark this location so we can rewind

        c.markLocalPositionStart();

        for (int i = 0; i < attrCount; ++i) {

            copyAttrForMethod(methodName, accessFlags);

        }

        if (true) {

            System.err.println("   wrapped: " + methodName);

        }

        // now write the wrapper method

        c.writeU2(accessFlags & ~ACC_NATIVE);

        c.writeU2(nameIndex);           // original unwrapped name

        c.writeU2(descriptorIndex);     // descriptor is the same

        c.writeU2(attrCount + 1);       // wrapped plus a code attribute

        // rewind to wrapped attributes

        c.rewind();

        for (int i = 0; i < attrCount; ++i) {

            copyAttrForMethod(methodName, accessFlags);

        }

        // generate a Code attribute for the wrapper method

        int wrappedIndex = addMethodToConstantPool(getThisClassIndex(),

                                                   wrappedNameIndex,

                                                   descriptorIndex);

        String descriptor = c.constantPoolString(descriptorIndex);

        createWrapperCodeAttr(nameIndex, accessFlags, descriptor, wrappedIndex);

        // increment method count

        c.randomAccessWriteU2(methodsCountPos, ++methodsCount);

    }

    void copyAttrs(int attrCount) {

        for (int i = 0; i < attrCount; ++i) {

            copyAttr();

        }

    }

    void copyAttr() {

        c.copy(2);             // name

        int len = c.copyU4();  // attr len

        if (verbose) {

            System.out.println("attr len: " + len);

        }

        c.copy(len);           // attribute info

    }

    void copyAttrForMethod(String methodName, int accessFlags) {

        int nameIndex = c.copyU2();   // name

        // check for Code attr

        if (nameIndex == c.codeAttributeIndex) {

            try {

                copyCodeAttr(methodName);

            } catch (IOException exc) {

                System.err.println("Code Exception - " + exc);

                System.exit(1);

            }

        } else {

            int len = c.copyU4();     // attr len

            if (verbose) {

                System.out.println("method attr len: " + len);

            }

            c.copy(len);              // attribute info

        }

    }

    void copyAttrForCode(InjectBytecodes ib) throws IOException {

        int nameIndex = c.copyU2();   // name

        // check for Code attr

        if (nameIndex == c.lineNumberAttributeIndex) {

            ib.copyLineNumberAttr();

        } else if (nameIndex == c.localVarAttributeIndex) {

            ib.copyLocalVarAttr();

        } else {

            int len = c.copyU4();     // attr len

            if (verbose) {

                System.out.println("code attr len: " + len);

            }

            c.copy(len);              // attribute info

        }

    }

    void copyCodeAttr(String methodName) throws IOException {

        if (verbose) {

            System.out.println("Code attr found");

        }

        int attrLengthPos = c.generatedPosition();

        int attrLength = c.copyU4();        // attr len

        int maxStack = c.readU2();          // max stack

        c.writeU2(defaultInjector == null? maxStack :

                  defaultInjector.stackSize(maxStack));  // big enough for injected code

        c.copyU2();                         // max locals

        int codeLengthPos = c.generatedPosition();

        int codeLength = c.copyU4();        // code length

        if (options.targetMethod != null && !options.targetMethod.equals(methodName)) {

            c.copy(attrLength - 8); // copy remainder minus already copied

            return;

        }

        if (isSystem) {

            if (codeLength == 1 && methodName.equals("finalize")) {

                if (verbose) {

                    System.out.println("empty system finalizer not instrumented");

                }

                c.copy(attrLength - 8); // copy remainder minus already copied

                return;

            }

            if (codeLength == 1 && methodName.equals("<init>")) {

                if (verbose) {

                    System.out.println("empty system constructor not instrumented");

                }

                if (!options.shouldInstrumentObjectInit) {

                    c.copy(attrLength - 8); // copy remainder minus already copied

                    return;

                }

            }

            if (methodName.equals("<clinit>")) {

                if (verbose) {

                    System.out.println("system class initializer not instrumented");

                }

                c.copy(attrLength - 8); // copy remainder minus already copied

                return;

            }

        }

        if (options.shouldInstrumentObjectInit

            && (!className.equals("java/lang/Object")

                || !methodName.equals("<init>"))) {

            c.copy(attrLength - 8); // copy remainder minus already copied

            return;

        }

        InjectBytecodes ib = new InjectBytecodes(c, codeLength, className, methodName);

        if (options.shouldInstrumentNew) {

            ib.injectAfter(opc_new, allocInjector);

            ib.injectAfter(opc_newarray, allocInjector);

            ib.injectAfter(opc_anewarray, allocInjector);

            ib.injectAfter(opc_multianewarray, allocInjector);

        }

        if (options.shouldInstrumentCall) {

            ib.inject(0, callInjector.bytecodes(className, methodName, 0));

        }

        if (options.shouldInstrumentObjectInit) {

            ib.inject(0, allocInjector.bytecodes(className, methodName, 0));

        }

        ib.adjustOffsets();

        // fix up code length

        int newCodeLength = c.generatedPosition() - (codeLengthPos + 4);

        c.randomAccessWriteU4(codeLengthPos, newCodeLength);

        if (verbose) {

            System.out.println("code length old: " + codeLength +

                               ", new: " + newCodeLength);

        }

        ib.copyExceptionTable();

        int attrCount = c.copyU2();

        for (int i = 0; i < attrCount; ++i) {

            copyAttrForCode(ib);

        }

        // fix up attr length

        int newAttrLength = c.generatedPosition() - (attrLengthPos + 4);

        c.randomAccessWriteU4(attrLengthPos, newAttrLength);

        if (verbose) {

            System.out.println("attr length old: " + attrLength +

                               ", new: " + newAttrLength);

        }

    }

    int nextDescriptorIndex(String descriptor, int index) {

        switch (descriptor.charAt(index)) {

        case 'B': // byte

        case 'C': // char

        case 'I': // int

        case 'S': // short

        case 'Z': // boolean

        case 'F': // float

        case 'D': // double

        case 'J': // long

            return index + 1;

        case 'L': // object

            int i = index + 1;

            while (descriptor.charAt(i) != ';') {

                ++i;

            }

            return i + 1;

        case '[': // array

            return nextDescriptorIndex(descriptor, index + 1);

        }

        throw new InternalError("should not reach here");

    }

    int getWrappedTrackerIndex() {

        if (wrappedTrackerIndex == 0) {

            wrappedTrackerIndex = addMethodToConstantPool(profiler,

                                                          options.wrappedTrackerMethodName,

                                                          "(Ljava/lang/String;I)V");

        }

        return wrappedTrackerIndex;

    }

    int getThisClassIndex() {

        if (thisClassIndex == 0) {

            thisClassIndex = addClassToConstantPool(className);

        }

        return thisClassIndex;

    }

    int computeMaxLocals(String descriptor, int accessFlags) {

        int index = 1;

        int slot = 0;