tools/locked_region_code_injection/src/lockedregioncodeinjection/LockFindingClassVisitor.java - platform/frameworks/base - Git at Google

 /*
  * Copyright (C) 2017 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
  * in compliance with the License. You may obtain a copy of the License at
  *
  * http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software distributed under the License
  * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  * or implied. See the License for the specific language governing permissions and limitations under
  * the License.
  */
 package lockedregioncodeinjection;

 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.LinkedList;
 import java.util.List;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.commons.TryCatchBlockSorter;
 import org.objectweb.asm.tree.AbstractInsnNode;
 import org.objectweb.asm.tree.InsnList;
 import org.objectweb.asm.tree.LabelNode;
 import org.objectweb.asm.tree.MethodInsnNode;
 import org.objectweb.asm.tree.MethodNode;
 import org.objectweb.asm.tree.TryCatchBlockNode;
 import org.objectweb.asm.tree.analysis.Analyzer;
 import org.objectweb.asm.tree.analysis.AnalyzerException;
 import org.objectweb.asm.tree.analysis.BasicValue;
 import org.objectweb.asm.tree.analysis.Frame;

 import static com.google.common.base.Preconditions.checkElementIndex;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;

 /**
  * This visitor does two things:
  *
  * 1. Finds all the MONITOR_ENTER / MONITOR_EXIT in the byte code and insert the corresponding pre
  * and post methods calls should it matches one of the given target type in the Configuration.
  *
  * 2. Find all methods that are synchronized and insert pre method calls in the beginning and post
  * method calls just before all return instructions.
  */
 class LockFindingClassVisitor extends ClassVisitor {
     private String className = null;
     private final List<LockTarget> targets;

     public LockFindingClassVisitor(List<LockTarget> targets, ClassVisitor chain) {
         super(Utils.ASM_VERSION, chain);
         this.targets = targets;
     }

     @Override
     public MethodVisitor visitMethod(int access, String name, String desc, String signature,
             String[] exceptions) {
         assert this.className != null;
         MethodNode mn = new TryCatchBlockSorter(null, access, name, desc, signature, exceptions);
         MethodVisitor chain = super.visitMethod(access, name, desc, signature, exceptions);
         return new LockFindingMethodVisitor(this.className, mn, chain);
     }

     @Override
     public void visit(int version, int access, String name, String signature, String superName,
             String[] interfaces) {
         this.className = name;
         super.visit(version, access, name, signature, superName, interfaces);
     }

     class LockFindingMethodVisitor extends MethodVisitor {
         private String owner;
         private MethodVisitor chain;

         public LockFindingMethodVisitor(String owner, MethodNode mn, MethodVisitor chain) {
             super(Utils.ASM_VERSION, mn);
             assert owner != null;
             this.owner = owner;
             this.chain = chain;
         }

         @SuppressWarnings("unchecked")
         @Override
         public void visitEnd() {
             MethodNode mn = (MethodNode) mv;

             Analyzer a = new Analyzer(new LockTargetStateAnalysis(targets));

             LockTarget ownerMonitor = null;
             if ((mn.access & Opcodes.ACC_SYNCHRONIZED) != 0) {
                 for (LockTarget t : targets) {
                     if (t.getTargetDesc().equals("L" + owner + ";")) {
                         ownerMonitor = t;
                     }
                 }
             }

             try {
                 a.analyze(owner, mn);
             } catch (AnalyzerException e) {
                 throw new RuntimeException("Locked region code injection: " + e.getMessage(), e);
             }
             InsnList instructions = mn.instructions;

             Frame[] frames = a.getFrames();
             List<Frame> frameMap = new LinkedList<>();
             frameMap.addAll(Arrays.asList(frames));

             List<List<TryCatchBlockNode>> handlersMap = new LinkedList<>();

             for (int i = 0; i < instructions.size(); i++) {
                 handlersMap.add(a.getHandlers(i));
             }

             if (ownerMonitor != null) {
                 AbstractInsnNode s = instructions.getFirst();
                 MethodInsnNode call = new MethodInsnNode(Opcodes.INVOKESTATIC,
                         ownerMonitor.getPreOwner(), ownerMonitor.getPreMethod(), "()V", false);
                 insertMethodCallBefore(mn, frameMap, handlersMap, s, 0, call);
             }

             for (int i = 0; i < instructions.size(); i++) {
                 AbstractInsnNode s = instructions.get(i);

                 if (s.getOpcode() == Opcodes.MONITORENTER) {
                     Frame f = frameMap.get(i);
                     BasicValue operand = (BasicValue) f.getStack(f.getStackSize() - 1);
                     if (operand instanceof LockTargetState) {
                         LockTargetState state = (LockTargetState) operand;
                         for (int j = 0; j < state.getTargets().size(); j++) {
                             LockTarget target = state.getTargets().get(j);
                             MethodInsnNode call = new MethodInsnNode(Opcodes.INVOKESTATIC,
                                     target.getPreOwner(), target.getPreMethod(), "()V", false);
                             insertMethodCallAfter(mn, frameMap, handlersMap, s, i, call);
                         }
                     }
                 }

                 if (s.getOpcode() == Opcodes.MONITOREXIT) {
                     Frame f = frameMap.get(i);
                     BasicValue operand = (BasicValue) f.getStack(f.getStackSize() - 1);
                     if (operand instanceof LockTargetState) {
                         LockTargetState state = (LockTargetState) operand;
                         for (int j = 0; j < state.getTargets().size(); j++) {
                             // The instruction after a monitor_exit should be a label for the end of the implicit
                             // catch block that surrounds the synchronized block to call monitor_exit when an exception
                             // occurs.
                             checkState(instructions.get(i + 1).getType() == AbstractInsnNode.LABEL,
                                 "Expected to find label after monitor exit");

                             int labelIndex = i + 1;
                             checkElementIndex(labelIndex, instructions.size());

                             LabelNode label = (LabelNode)instructions.get(labelIndex);

                             checkNotNull(handlersMap.get(i));
                             checkElementIndex(0, handlersMap.get(i).size());
                             checkState(handlersMap.get(i).get(0).end == label,
                                 "Expected label to be the end of monitor exit's try block");

                             LockTarget target = state.getTargets().get(j);
                             MethodInsnNode call = new MethodInsnNode(Opcodes.INVOKESTATIC,
                                     target.getPostOwner(), target.getPostMethod(), "()V", false);
                             insertMethodCallAfter(mn, frameMap, handlersMap, label, labelIndex, call);
                         }
                     }
                 }

                 if (ownerMonitor != null && (s.getOpcode() == Opcodes.RETURN
                         || s.getOpcode() == Opcodes.ARETURN || s.getOpcode() == Opcodes.DRETURN
                         || s.getOpcode() == Opcodes.FRETURN || s.getOpcode() == Opcodes.IRETURN)) {
                     MethodInsnNode call =
                             new MethodInsnNode(Opcodes.INVOKESTATIC, ownerMonitor.getPostOwner(),
                                     ownerMonitor.getPostMethod(), "()V", false);
                     insertMethodCallBefore(mn, frameMap, handlersMap, s, i, call);
                     i++; // Skip ahead. Otherwise, we will revisit this instruction again.
                 }
             }
             super.visitEnd();
             mn.accept(chain);
         }
     }

     public static void insertMethodCallBefore(MethodNode mn, List<Frame> frameMap,
             List<List<TryCatchBlockNode>> handlersMap, AbstractInsnNode node, int index,
             MethodInsnNode call) {
         List<TryCatchBlockNode> handlers = handlersMap.get(index);
         InsnList instructions = mn.instructions;
         LabelNode end = new LabelNode();
         instructions.insert(node, end);
         frameMap.add(index, null);
         handlersMap.add(index, null);
         instructions.insertBefore(node, call);
         frameMap.add(index, null);
         handlersMap.add(index, null);

         LabelNode start = new LabelNode();
         instructions.insert(node, start);
         frameMap.add(index, null);
         handlersMap.add(index, null);
         updateCatchHandler(mn, handlers, start, end, handlersMap);
     }

     public static void insertMethodCallAfter(MethodNode mn, List<Frame> frameMap,
             List<List<TryCatchBlockNode>> handlersMap, AbstractInsnNode node, int index,
             MethodInsnNode call) {
         List<TryCatchBlockNode> handlers = handlersMap.get(index + 1);
         InsnList instructions = mn.instructions;

         LabelNode end = new LabelNode();
         instructions.insert(node, end);
         frameMap.add(index + 1, null);
         handlersMap.add(index + 1, null);

         instructions.insert(node, call);
         frameMap.add(index + 1, null);
         handlersMap.add(index + 1, null);

         LabelNode start = new LabelNode();
         instructions.insert(node, start);
         frameMap.add(index + 1, null);
         handlersMap.add(index + 1, null);

         updateCatchHandler(mn, handlers, start, end, handlersMap);
     }

     @SuppressWarnings("unchecked")
     public static void updateCatchHandler(MethodNode mn, List<TryCatchBlockNode> handlers,
             LabelNode start, LabelNode end, List<List<TryCatchBlockNode>> handlersMap) {
         if (handlers == null || handlers.size() == 0) {
             return;
         }

         InsnList instructions = mn.instructions;
         List<TryCatchBlockNode> newNodes = new ArrayList<>(handlers.size());
         for (TryCatchBlockNode handler : handlers) {
             if (!(instructions.indexOf(handler.start) <= instructions.indexOf(start)
                     && instructions.indexOf(end) <= instructions.indexOf(handler.end))) {
                 TryCatchBlockNode newNode =
                         new TryCatchBlockNode(start, end, handler.handler, handler.type);
                 newNodes.add(newNode);
                 for (int i = instructions.indexOf(start); i <= instructions.indexOf(end); i++) {
                     if (handlersMap.get(i) == null) {
                         handlersMap.set(i, new ArrayList<>());
                     }
                     handlersMap.get(i).add(newNode);
                 }
             } else {
                 for (int i = instructions.indexOf(start); i <= instructions.indexOf(end); i++) {
                     if (handlersMap.get(i) == null) {
                         handlersMap.set(i, new ArrayList<>());
                     }
                     handlersMap.get(i).add(handler);
                 }
             }
         }
         mn.tryCatchBlocks.addAll(0, newNodes);
     }
 }
	/*
	* Copyright (C) 2017 The Android Open Source Project
	*
	* Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
	* in compliance with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software distributed under the License
	* is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
	* or implied. See the License for the specific language governing permissions and limitations under
	* the License.
	*/
	package lockedregioncodeinjection;

	import java.util.ArrayList;
	import java.util.Arrays;
	import java.util.LinkedList;
	import java.util.List;
	import org.objectweb.asm.ClassVisitor;
	import org.objectweb.asm.MethodVisitor;
	import org.objectweb.asm.Opcodes;
	import org.objectweb.asm.commons.TryCatchBlockSorter;
	import org.objectweb.asm.tree.AbstractInsnNode;
	import org.objectweb.asm.tree.InsnList;
	import org.objectweb.asm.tree.LabelNode;
	import org.objectweb.asm.tree.MethodInsnNode;
	import org.objectweb.asm.tree.MethodNode;
	import org.objectweb.asm.tree.TryCatchBlockNode;
	import org.objectweb.asm.tree.analysis.Analyzer;
	import org.objectweb.asm.tree.analysis.AnalyzerException;
	import org.objectweb.asm.tree.analysis.BasicValue;
	import org.objectweb.asm.tree.analysis.Frame;

	import static com.google.common.base.Preconditions.checkElementIndex;
	import static com.google.common.base.Preconditions.checkNotNull;
	import static com.google.common.base.Preconditions.checkState;

	/**
	* This visitor does two things:
	*
	* 1. Finds all the MONITOR_ENTER / MONITOR_EXIT in the byte code and insert the corresponding pre
	* and post methods calls should it matches one of the given target type in the Configuration.
	*
	* 2. Find all methods that are synchronized and insert pre method calls in the beginning and post
	* method calls just before all return instructions.
	*/
	class LockFindingClassVisitor extends ClassVisitor {
	private String className = null;
	private final List<LockTarget> targets;

	public LockFindingClassVisitor(List<LockTarget> targets, ClassVisitor chain) {
	super(Utils.ASM_VERSION, chain);
	this.targets = targets;
	}

	@Override
	public MethodVisitor visitMethod(int access, String name, String desc, String signature,
	String[] exceptions) {
	assert this.className != null;
	MethodNode mn = new TryCatchBlockSorter(null, access, name, desc, signature, exceptions);
	MethodVisitor chain = super.visitMethod(access, name, desc, signature, exceptions);
	return new LockFindingMethodVisitor(this.className, mn, chain);
	}

	@Override
	public void visit(int version, int access, String name, String signature, String superName,
	String[] interfaces) {
	this.className = name;
	super.visit(version, access, name, signature, superName, interfaces);
	}

	class LockFindingMethodVisitor extends MethodVisitor {
	private String owner;
	private MethodVisitor chain;

	public LockFindingMethodVisitor(String owner, MethodNode mn, MethodVisitor chain) {
	super(Utils.ASM_VERSION, mn);
	assert owner != null;
	this.owner = owner;
	this.chain = chain;
	}

	@SuppressWarnings("unchecked")
	@Override
	public void visitEnd() {
	MethodNode mn = (MethodNode) mv;

	Analyzer a = new Analyzer(new LockTargetStateAnalysis(targets));

	LockTarget ownerMonitor = null;
	if ((mn.access & Opcodes.ACC_SYNCHRONIZED) != 0) {
	for (LockTarget t : targets) {
	if (t.getTargetDesc().equals("L" + owner + ";")) {
	ownerMonitor = t;
	}
	}
	}

	try {
	a.analyze(owner, mn);
	} catch (AnalyzerException e) {
	throw new RuntimeException("Locked region code injection: " + e.getMessage(), e);
	}
	InsnList instructions = mn.instructions;

	Frame[] frames = a.getFrames();
	List<Frame> frameMap = new LinkedList<>();
	frameMap.addAll(Arrays.asList(frames));

	List<List<TryCatchBlockNode>> handlersMap = new LinkedList<>();

	for (int i = 0; i < instructions.size(); i++) {
	handlersMap.add(a.getHandlers(i));
	}

	if (ownerMonitor != null) {
	AbstractInsnNode s = instructions.getFirst();
	MethodInsnNode call = new MethodInsnNode(Opcodes.INVOKESTATIC,
	ownerMonitor.getPreOwner(), ownerMonitor.getPreMethod(), "()V", false);
	insertMethodCallBefore(mn, frameMap, handlersMap, s, 0, call);
	}

	for (int i = 0; i < instructions.size(); i++) {
	AbstractInsnNode s = instructions.get(i);

	if (s.getOpcode() == Opcodes.MONITORENTER) {
	Frame f = frameMap.get(i);
	BasicValue operand = (BasicValue) f.getStack(f.getStackSize() - 1);
	if (operand instanceof LockTargetState) {
	LockTargetState state = (LockTargetState) operand;
	for (int j = 0; j < state.getTargets().size(); j++) {
	LockTarget target = state.getTargets().get(j);
	MethodInsnNode call = new MethodInsnNode(Opcodes.INVOKESTATIC,
	target.getPreOwner(), target.getPreMethod(), "()V", false);
	insertMethodCallAfter(mn, frameMap, handlersMap, s, i, call);
	}
	}
	}

	if (s.getOpcode() == Opcodes.MONITOREXIT) {
	Frame f = frameMap.get(i);
	BasicValue operand = (BasicValue) f.getStack(f.getStackSize() - 1);
	if (operand instanceof LockTargetState) {
	LockTargetState state = (LockTargetState) operand;
	for (int j = 0; j < state.getTargets().size(); j++) {
	// The instruction after a monitor_exit should be a label for the end of the implicit
	// catch block that surrounds the synchronized block to call monitor_exit when an exception
	// occurs.
	checkState(instructions.get(i + 1).getType() == AbstractInsnNode.LABEL,
	"Expected to find label after monitor exit");

	int labelIndex = i + 1;
	checkElementIndex(labelIndex, instructions.size());

	LabelNode label = (LabelNode)instructions.get(labelIndex);

	checkNotNull(handlersMap.get(i));
	checkElementIndex(0, handlersMap.get(i).size());
	checkState(handlersMap.get(i).get(0).end == label,
	"Expected label to be the end of monitor exit's try block");

	LockTarget target = state.getTargets().get(j);
	MethodInsnNode call = new MethodInsnNode(Opcodes.INVOKESTATIC,
	target.getPostOwner(), target.getPostMethod(), "()V", false);
	insertMethodCallAfter(mn, frameMap, handlersMap, label, labelIndex, call);
	}
	}
	}

	if (ownerMonitor != null && (s.getOpcode() == Opcodes.RETURN
	\|\| s.getOpcode() == Opcodes.ARETURN \|\| s.getOpcode() == Opcodes.DRETURN
	\|\| s.getOpcode() == Opcodes.FRETURN \|\| s.getOpcode() == Opcodes.IRETURN)) {
	MethodInsnNode call =
	new MethodInsnNode(Opcodes.INVOKESTATIC, ownerMonitor.getPostOwner(),
	ownerMonitor.getPostMethod(), "()V", false);
	insertMethodCallBefore(mn, frameMap, handlersMap, s, i, call);
	i++; // Skip ahead. Otherwise, we will revisit this instruction again.
	}
	}
	super.visitEnd();
	mn.accept(chain);
	}
	}

	public static void insertMethodCallBefore(MethodNode mn, List<Frame> frameMap,
	List<List<TryCatchBlockNode>> handlersMap, AbstractInsnNode node, int index,
	MethodInsnNode call) {
	List<TryCatchBlockNode> handlers = handlersMap.get(index);
	InsnList instructions = mn.instructions;
	LabelNode end = new LabelNode();
	instructions.insert(node, end);
	frameMap.add(index, null);
	handlersMap.add(index, null);
	instructions.insertBefore(node, call);
	frameMap.add(index, null);
	handlersMap.add(index, null);

	LabelNode start = new LabelNode();
	instructions.insert(node, start);
	frameMap.add(index, null);
	handlersMap.add(index, null);
	updateCatchHandler(mn, handlers, start, end, handlersMap);
	}

	public static void insertMethodCallAfter(MethodNode mn, List<Frame> frameMap,
	List<List<TryCatchBlockNode>> handlersMap, AbstractInsnNode node, int index,
	MethodInsnNode call) {
	List<TryCatchBlockNode> handlers = handlersMap.get(index + 1);
	InsnList instructions = mn.instructions;

	LabelNode end = new LabelNode();
	instructions.insert(node, end);
	frameMap.add(index + 1, null);
	handlersMap.add(index + 1, null);

	instructions.insert(node, call);
	frameMap.add(index + 1, null);
	handlersMap.add(index + 1, null);

	LabelNode start = new LabelNode();
	instructions.insert(node, start);
	frameMap.add(index + 1, null);
	handlersMap.add(index + 1, null);

	updateCatchHandler(mn, handlers, start, end, handlersMap);
	}

	@SuppressWarnings("unchecked")
	public static void updateCatchHandler(MethodNode mn, List<TryCatchBlockNode> handlers,
	LabelNode start, LabelNode end, List<List<TryCatchBlockNode>> handlersMap) {
	if (handlers == null \|\| handlers.size() == 0) {
	return;
	}

	InsnList instructions = mn.instructions;
	List<TryCatchBlockNode> newNodes = new ArrayList<>(handlers.size());
	for (TryCatchBlockNode handler : handlers) {
	if (!(instructions.indexOf(handler.start) <= instructions.indexOf(start)
	&& instructions.indexOf(end) <= instructions.indexOf(handler.end))) {
	TryCatchBlockNode newNode =
	new TryCatchBlockNode(start, end, handler.handler, handler.type);
	newNodes.add(newNode);
	for (int i = instructions.indexOf(start); i <= instructions.indexOf(end); i++) {
	if (handlersMap.get(i) == null) {
	handlersMap.set(i, new ArrayList<>());
	}
	handlersMap.get(i).add(newNode);
	}
	} else {
	for (int i = instructions.indexOf(start); i <= instructions.indexOf(end); i++) {
	if (handlersMap.get(i) == null) {
	handlersMap.set(i, new ArrayList<>());
	}
	handlersMap.get(i).add(handler);
	}
	}
	}
	mn.tryCatchBlocks.addAll(0, newNodes);
	}
	}