content: /*
   * Copyright (C) 2014 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 dexfuzz.program;
  
  import dexfuzz.Log;
  import dexfuzz.rawdex.Instruction;
  import dexfuzz.rawdex.Opcode;
  
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.LinkedList;
  import java.util.List;
  
  /**
   * A class that represents a CodeItem in a way that is more amenable to mutation.
   */
  public class MutatableCode {
    /**
     * To ensure we update the correct CodeItem in the raw DEX file.
     */
    public int codeItemIdx;
  
    /**
     * This is an index into the Program's list of MutatableCodes.
     */
    public int mutatableCodeIdx;
  
    /**
     * Number of registers this code uses.
     */
    public short registersSize;
  
    /**
     * Number of ins this code has.
     */
    public short insSize;
  
    /**
     * Number of outs this code has.
     */
    public short outsSize;
  
    /**
     * Number of tries this code has.
     */
    public short triesSize;
  
    /**
     * CodeTranslator is responsible for creating this, and
     * converting it back to a list of Instructions.
     */
    private List<MInsn> mutatableInsns;
  
    /**
     * CodeTranslator is responsible for creating this, and
     * converting it back to the correct form for CodeItems.
     */
    public List<MTryBlock> mutatableTries;
  
    /**
     * The name of the method this code represents.
     */
    public String name;
    public String shorty;
    public boolean isStatic;
  
    /**
     * The Program that owns this MutatableCode.
     * Currently used to get the size of constant pools for
     * PoolIndexChanger/RandomInstructionGenerator
     */
    public Program program;
  
    private short originalInVReg;
    private short tempVRegsAllocated;
    private short initialTempVReg;
    private boolean vregsNeedCopying;
    private int numMoveInsnsGenerated;
  
    public MutatableCode(Program program) {
      this.program = program;
      this.mutatableInsns = new LinkedList<MInsn>();
    }
  
    /**
     * Call this to update all instructions after the provided mInsn, to have their
     * locations adjusted by the provided offset. It will also mark that they have been updated.
     */
    public void updateInstructionLocationsAfter(MInsn mInsn, int offset) {
      boolean updating = false;
      for (MInsn mInsnChecking : mutatableInsns) {
        if (updating) {
          mInsnChecking.locationUpdated = true;
          mInsnChecking.location += offset;
        } else {
          if (mInsnChecking == mInsn) {
            updating = true;
          }
        }
  
      }
    }
  
    private void recalculateLocations() {
      int loc = 0;
      for (MInsn mInsn : mutatableInsns) {
        mInsn.location = loc;
        loc += mInsn.insn.getSize();
      }
    }
  
    public List<MInsn> getInstructions() {
      return Collections.unmodifiableList(mutatableInsns);
    }
  
    public int getInstructionCount() {
      return mutatableInsns.size();
    }
  
    public int getInstructionIndex(MInsn mInsn) {
      return mutatableInsns.indexOf(mInsn);
    }
  
    public MInsn getInstructionAt(int idx) {
      return mutatableInsns.get(idx);
    }
  
    public void addInstructionToEnd(MInsn mInsn) {
      mutatableInsns.add(mInsn);
    }
  
    public void insertInstructionAfter(MInsn toBeInserted, int insertionIdx) {
      if ((insertionIdx + 1) < mutatableInsns.size()) {
        insertInstructionAt(toBeInserted, insertionIdx + 1);
      } else {
        // Appending to end.
        MInsn finalInsn = mutatableInsns.get(mutatableInsns.size() - 1);
        toBeInserted.location = finalInsn.location + finalInsn.insn.getSize();
        mutatableInsns.add(toBeInserted);
      }
    }
  
    /**
     * Has same semantics as List.add()
     */
    public void insertInstructionAt(MInsn toBeInserted, int insertionIdx) {
      MInsn currentInsn = mutatableInsns.get(insertionIdx);
      toBeInserted.location = currentInsn.location;
      mutatableInsns.add(insertionIdx , toBeInserted);
      updateInstructionLocationsAfter(toBeInserted, toBeInserted.insn.getSize());
    }
  
    /**
     * Checks if any MTryBlock's instruction refs pointed at the 'before' MInsn,
     * and points them to the 'after' MInsn, if so. 'twoWay' will check if 'after'
     * was pointed to, and point refs to the 'before' insn.
     * (one-way is used when deleting instructions,
     * two-way is used when swapping instructions.)
     */
    private void updateTryBlocksWithReplacementInsn(MInsn before, MInsn after,
        boolean twoWay) {
      if (triesSize > 0) {
        for (MTryBlock mTryBlock : mutatableTries) {
          if (mTryBlock.startInsn == before) {
            Log.debug("Try block's first instruction was updated");
            mTryBlock.startInsn = after;
          } else if (twoWay && mTryBlock.startInsn == after) {
            Log.debug("Try block's first instruction was updated");
            mTryBlock.startInsn = before;
          }
          if (mTryBlock.endInsn == before) {
            Log.debug("Try block's last instruction was updated");
            mTryBlock.endInsn = after;
          } else if (twoWay && mTryBlock.endInsn == after) {
            Log.debug("Try block's last instruction was updated");
            mTryBlock.endInsn = before;
          }
          if (mTryBlock.catchAllHandler == before) {
            Log.debug("Try block's catch-all instruction was updated");
            mTryBlock.catchAllHandler = after;
          } else if (twoWay && mTryBlock.catchAllHandler == after) {
            Log.debug("Try block's catch-all instruction was updated");
            mTryBlock.catchAllHandler = before;
          }
  
          List<Integer> matchesIndicesToChange = new ArrayList<Integer>();
          List<Integer> replacementIndicesToChange = null;
          if (twoWay) {
            replacementIndicesToChange = new ArrayList<Integer>();
          }
  
          int idx = 0;
          for (MInsn handler : mTryBlock.handlers) {
            if (handler == before) {
              matchesIndicesToChange.add(idx);
              Log.debug("Try block's handler instruction was updated");
            } else if (twoWay && handler == after) {
              replacementIndicesToChange.add(idx);
              Log.debug("Try block's handler instruction was updated");
            }
            idx++;
          }
  
          for (int idxToChange : matchesIndicesToChange) {
            mTryBlock.handlers.set(idxToChange, after);
          }
  
          if (twoWay) {
            for (int idxToChange : replacementIndicesToChange) {
              mTryBlock.handlers.set(idxToChange, before);
            }
          }
        }
      }
    }
  
    /**
     * The actual implementation of deleteInstruction called by
     * the single-argument deleteInstructions.
     */
    private void deleteInstructionFull(MInsn toBeDeleted, int toBeDeletedIdx) {
      // Make sure we update all locations afterwards first.
      updateInstructionLocationsAfter(toBeDeleted, -(toBeDeleted.insn.getSize()));
  
      // Remove it.
      mutatableInsns.remove(toBeDeletedIdx);
  
      // Update any branch instructions that branched to the instruction we just deleted!
  
      // First, pick the replacement target.
      int replacementTargetIdx = toBeDeletedIdx;
      if (replacementTargetIdx == mutatableInsns.size()) {
        replacementTargetIdx--;
      }
      MInsn replacementTarget = mutatableInsns.get(replacementTargetIdx);
  
      for (MInsn mInsn : mutatableInsns) {
        if (mInsn instanceof MBranchInsn) {
          // Check if this branch insn points at the insn we just deleted.
          MBranchInsn branchInsn = (MBranchInsn) mInsn;
          MInsn target = branchInsn.target;
          if (target == toBeDeleted) {
            Log.debug(branchInsn + " was pointing at the deleted instruction, updated.");
            branchInsn.target = replacementTarget;
          }
        } else if (mInsn instanceof MSwitchInsn) {
          // Check if any of this switch insn's targets points at the insn we just deleted.
          MSwitchInsn switchInsn = (MSwitchInsn) mInsn;
          List<Integer> indicesToChange = new ArrayList<Integer>();
          int idx = 0;
          for (MInsn target : switchInsn.targets) {
            if (target == toBeDeleted) {
              indicesToChange.add(idx);
              Log.debug(switchInsn + "[" + idx
                  + "] was pointing at the deleted instruction, updated.");
            }
            idx++;
          }
          for (int idxToChange : indicesToChange) {
            switchInsn.targets.remove(idxToChange);
            switchInsn.targets.add(idxToChange, replacementTarget);
          }
        }
      }
  
      // Now update the try blocks.
      updateTryBlocksWithReplacementInsn(toBeDeleted, replacementTarget, false);
    }
  
    /**
     * Delete the provided MInsn.
     */
    public void deleteInstruction(MInsn toBeDeleted) {
      deleteInstructionFull(toBeDeleted, mutatableInsns.indexOf(toBeDeleted));
    }
  
    /**
     * Delete the MInsn at the provided index.
     */
    public void deleteInstruction(int toBeDeletedIdx) {
      deleteInstructionFull(mutatableInsns.get(toBeDeletedIdx), toBeDeletedIdx);
    }
  
    public void swapInstructionsByIndex(int aIdx, int bIdx) {
      MInsn aInsn = mutatableInsns.get(aIdx);
      MInsn bInsn = mutatableInsns.get(bIdx);
  
      mutatableInsns.set(aIdx, bInsn);
      mutatableInsns.set(bIdx, aInsn);
  
      updateTryBlocksWithReplacementInsn(aInsn, bInsn, true);
  
      recalculateLocations();
    }
  
    /**
     * Some mutators may require the use of temporary registers. For instance,
     * to easily add in printing of values without having to look for registers
     * that aren't currently live.
     * The idea is to allocate these registers at the top of the set of registers.
     * Because this will then shift where the arguments to the method are, we then
     * change the start of the method to copy the arguments to the method
     * into the place where the rest of the method's code expects them to be.
     * Call allocateTemporaryVRegs(n), then use getTemporaryVReg(n),
     * and then make sure finishedUsingTemporaryVRegs() is called!
     */
    public void allocateTemporaryVRegs(int count) {
      if (count > tempVRegsAllocated) {
        if (tempVRegsAllocated == 0) {
          Log.info("Allocating temporary vregs for method...");
          initialTempVReg = registersSize;
          originalInVReg = (short) (registersSize - insSize);
        } else {
          Log.info("Extending allocation of temporary vregs for method...");
        }
        registersSize = (short) (initialTempVReg + count);
        if (outsSize < count) {
          outsSize = (short) count;
        }
        vregsNeedCopying = true;
        tempVRegsAllocated = (short) count;
      }
    }
  
    public int getTemporaryVReg(int number) {
      if (number >= tempVRegsAllocated) {
        Log.errorAndQuit("Not allocated enough temporary vregs!");
      }
      return initialTempVReg + number;
    }
  
    public void finishedUsingTemporaryVRegs() {
      if (tempVRegsAllocated > 0 && vregsNeedCopying) {
        // Just delete all the move instructions and generate again, if we already have some.
        while (numMoveInsnsGenerated > 0) {
          deleteInstruction(0);
          numMoveInsnsGenerated--;
        }
  
        Log.info("Moving 'in' vregs to correct locations after allocating temporary vregs");
  
        int shortyIdx = 0;
        if (isStatic) {
          shortyIdx = 1;
        }
  
        int insertionCounter = 0;
  
        // Insert copy insns that move all the in VRs down.
        for (int i = 0; i < insSize; i++) {
          MInsn moveInsn = new MInsn();
          moveInsn.insn = new Instruction();
          moveInsn.insn.vregA = originalInVReg + i;
          moveInsn.insn.vregB = originalInVReg + i + tempVRegsAllocated;
  
          char type = 'L';
          if (shortyIdx > 0) {
            type = shorty.charAt(shortyIdx);
          }
          shortyIdx++;
  
          if (type == 'L') {
            moveInsn.insn.info = Instruction.getOpcodeInfo(Opcode.MOVE_OBJECT_16);
          } else if (type == 'D' || type == 'J') {
            moveInsn.insn.info = Instruction.getOpcodeInfo(Opcode.MOVE_WIDE_16);
            i++;
          } else {
            moveInsn.insn.info = Instruction.getOpcodeInfo(Opcode.MOVE_16);
          }
  
          insertInstructionAt(moveInsn, insertionCounter);
          insertionCounter++;
          Log.info("Temp vregs creation, Added instruction " + moveInsn);
          numMoveInsnsGenerated++;
        }
  
        vregsNeedCopying = false;
      }
    }
  
    /**
     * When we insert new Field/Type/MethodIds into the DEX file, this may shunt some Ids
     * into a new position in the table. If this happens, every reference to the Ids must
     * be updated! Because CodeItems have their Instructions wrapped into a graph of MInsns
     * during mutation, they don't have a view of all their instructions during mutation,
     * and so if they are asked to update their instructions' indices into the tables, they
     * must call this method to get the actual list of instructions they currently own.
     */
    public List<Instruction> requestLatestInstructions() {
      List<Instruction> latestInsns = new ArrayList<Instruction>();
      for (MInsn mInsn : mutatableInsns) {
        latestInsns.add(mInsn.insn);
      }
      return latestInsns;
    }
  }
  