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.rawdex;
  
  import dexfuzz.Log;
  
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  
  /**
   * This class allows the recording of both Offsettable items (that is, items that can be
   * referred to by an offset somewhere else in the file - RawDexObjects) and Offsets.
   * The idea in a nutshell is that for every Offsettable item we read, we remember
   * its original position in the file using a map, and the order in which the Offsettables were
   * written out. We also remember every Offset we read in, and its value. Then, after reading
   * the whole file, we use the map to find the Offsettable it pointed at.
   * Then, as we write out the file, for every Offsettable we write out, we record its new position,
   * using the order we collected earlier. For every Offset we write out, we look at its Offsettable
   * to see where it was written. If it hasn't been written yet, then we write out a blank value
   * for the time being, remember where that blank value was written, and put the Offset into a
   * table for patching once everything has been written out.
   * There are some variables (index_after_map_list, restore_point) used for remembering certain
   * points to jump forward and back to, because we cannot read and write the file out in exactly
   * the same order.
   * TODO: Perhaps it makes more sense to just reorder the offsettable_table once it's been read,
   * in preparation for the order in which the file is written out?
   * Finally, we provide methods for adding new Offsettable items into the right place in the order
   * table.
   */
  public class OffsetTracker {
    /**
     * A map from the original offset in the input DEX file to
     * the Offsettable it points to. (That Offsettable will contain
     * the actual item, and later on the new offset for the item when
     * the item is written out.
     */
    private Map<Integer, Offsettable> offsettableMap;
  
    /**
     * A table of all Offsettables. We need to ensure we write out
     * all items in the same order we read them in, to make sure we update
     * the Offsettable.new_position field with the correct value wrt to
     * the original_position field.
     */
    private List<Offsettable> offsettableTable;
  
    /**
     * A table of all offsets that is populated as we read in the DEX file.
     * As the end, we find the correct Offsettable for the Offset in the above
     * map, and associate them.
     */
    private List<Offset> needsAssociationTable;
  
    /**
     * A table of all offsets that we could not write out an updated offset for
     * as we write out a DEX file. Will be checked after writing is complete,
     * to allow specific patching of each offset's location as at that point
     * all Offsettables will have been updated with their new position.
     */
    private List<Offset> needsUpdateTable;
  
    /**
     * Tracks how far we are through the offsettable_table as we write out the file.
     */
    private int offsettableTableIdx;
  
    /**
     * Because we write in a slightly different order to how we read
     * (why? to read, we read the header, then the map list, and then use the map
     *  list to read everything else.
     *  when we write, we write the header, and then we cannot write the map list
     *  because we don't know where it will go yet, so we write everything else first)
     * so: we remember the position in the offsettable_table after we read the map list,
     * so we can skip there after we write out the header.
     */
    private int indexAfterMapList;
  
    /**
     * Related to index_after_map_list, this is the index we save when we're jumping back to
     * write the map list.
     */
    private int restorePoint;
  
    /**
     * Create a new OffsetTracker. Should persist between parsing a DEX file, and outputting
     * the mutated DEX file.
     */
    public OffsetTracker() {
      offsettableMap = new HashMap<Integer,Offsettable>();
      offsettableTable = new ArrayList<Offsettable>();
      needsAssociationTable = new ArrayList<Offset>();
      needsUpdateTable = new ArrayList<Offset>();
    }
  
    /**
     * Lookup an Item by the offset it had in the input DEX file.
     * @param offset The offset in the input DEX file.
     * @return The corresponding Item.
     */
    public RawDexObject getItemByOffset(int offset) {
      return offsettableMap.get(offset).getItem();
    }
  
    /**
     * As Items are read in, they call this function once they have word-aligned the file pointer,
     * to record their position and themselves into an Offsettable object, that will be tracked.
     * @param file Used for recording position into the new Offsettable.
     * @param item Used for recording the relevant Item into the new Offsettable.
     */
    public void getNewOffsettable(DexRandomAccessFile file, RawDexObject item) throws IOException {
      Offsettable offsettable = new Offsettable(item, false);
      offsettable.setOriginalPosition((int) file.getFilePointer());
      offsettableMap.put(offsettable.getOriginalPosition(), offsettable);
      offsettableTable.add(offsettable);
    }
  
    /**
     * As Items read in Offsets, they call this function with the offset they originally
     * read from the file, to allow later association with an Offsettable.
     * @param originalOffset The original offset read from the input DEX file.
     * @return An Offset that will later be associated with an Offsettable.
     */
    public Offset getNewOffset(int originalOffset) throws IOException {
      Offset offset = new Offset(false);
      offset.setOriginalOffset(originalOffset);
      needsAssociationTable.add(offset);
      return offset;
    }
  
    /**
     * Only MapItem should call this method, when the MapItem that points to the header
     * is read.
     */
    public Offset getNewHeaderOffset(int originalOffset) throws IOException {
      Offset offset = new Offset(true);
      offset.setOriginalOffset(originalOffset);
      needsAssociationTable.add(offset);
      return offset;
    }
  
    /**
     * Call this after reading, to associate Offsets with Offsettables.
     */
    public void associateOffsets() {
      for (Offset offset : needsAssociationTable) {
        if (offset.getOriginalOffset() == 0 && !(offset.pointsAtHeader())) {
          offset.setPointsAtNull();
        } else {
          offset.pointTo(offsettableMap.get(offset.getOriginalOffset()));
          if (!offset.pointsToSomething()) {
            Log.error(String.format("Couldn't find original offset 0x%x!",
                offset.getOriginalOffset()));
          }
        }
      }
      needsAssociationTable.clear();
    }
  
    /**
     * As Items are written out into the output DEX file, this function is called
     * to update the next Offsettable with the file pointer's current position.
     * This should allow the tracking of new offset locations.
     * This also requires that reading and writing of all items happens in the same order
     * (with the exception of the map list, see above)
     * @param file Used for recording the new position.
     */
    public void updatePositionOfNextOffsettable(DexRandomAccessFile file) throws IOException {
      if (offsettableTableIdx == offsettableTable.size()) {
        Log.errorAndQuit("Not all created Offsettable items have been added to the "
            + "Offsettable Table!");
      }
      Offsettable offsettable = offsettableTable.get(offsettableTableIdx);
      offsettable.setNewPosition((int) file.getFilePointer());
      offsettableTableIdx++;
    }
  
    /**
     * As Items are written out, any writing out of an offset must call this function, passing
     * in the relevant offset. This function will write out the offset, if the associated
     * Offsettable has been updated with its new position, or else will write out a null value, and
     * the Offset will be stored for writing after all Items have been written, and all
     * Offsettables MUST have been updated.
     * @param offset The offset received from getNewOffset().
     * @param file Used for writing out to the file.
     * @param useUleb128 Whether or not the offset should be written in UINT or ULEB128 form.
     */
    public void tryToWriteOffset(Offset offset, DexRandomAccessFile file, boolean useUleb128)
        throws IOException {
      if (!offset.isNewOffset() && (!offset.pointsToSomething())) {
        if (useUleb128) {
          file.writeUleb128(0);
        } else {
          file.writeUInt(0);
        }
        return;
      }
  
      if (offset.readyForWriting()) {
        if (useUleb128) {
          file.writeUleb128(offset.getNewPositionOfItem());
        } else {
          file.writeUInt(offset.getNewPositionOfItem());
        }
      } else {
        offset.setOutputLocation((int) file.getFilePointer());
        if (useUleb128) {
          file.writeLargestUleb128(offset.getOriginalOffset());
          offset.setUsesUleb128();
        } else {
          file.writeUInt(offset.getOriginalOffset());
        }
        needsUpdateTable.add(offset);
      }
    }
  
    /**
     * This is called after all writing has finished, to write out any Offsets
     * that could not be written out during the original writing phase, because their
     * associated Offsettables hadn't had their new positions calculated yet.
     * @param file Used for writing out to the file.
     */
    public void updateOffsets(DexRandomAccessFile file) throws IOException {
      if (offsettableTableIdx != offsettableTable.size()) {
        Log.errorAndQuit("Being asked to update dangling offsets but the "
            + "correct number of offsettables has not been written out!");
      }
      for (Offset offset : needsUpdateTable) {
        file.seek(offset.getOutputLocation());
        if (offset.usesUleb128()) {
          file.writeLargestUleb128(offset.getNewPositionOfItem());
        } else {
          file.writeUInt(offset.getNewPositionOfItem());
        }
      }
      needsUpdateTable.clear();
    }
  
    /**
     * Called after writing out the header, to skip to after the map list.
     */
    public void skipToAfterMapList() {
      offsettableTableIdx = indexAfterMapList;
    }
  
    /**
     * Called once the map list needs to be written out, to set the
     * offsettable table index back to the right place.
     */
    public void goBackToMapList() {
      restorePoint = offsettableTableIdx;
      offsettableTableIdx = (indexAfterMapList - 1);
    }
  
    /**
     * Called once the map list has been written out, to set the
     * offsettable table index back to where it was before.
     */
    public void goBackToPreviousPoint() {
      if (offsettableTableIdx != indexAfterMapList) {
        Log.errorAndQuit("Being asked to go to the point before the MapList was written out,"
            + " but we're not in the right place.");
      }
      offsettableTableIdx = restorePoint;
    }
  
    /**
     * Called after reading in the map list, to remember the point to be skipped
     * to later.
     */
    public void rememberPointAfterMapList() {
      indexAfterMapList = offsettableTable.size();
    }
  
    private void updateHeaderOffsetIfValid(Offset offset, Offsettable previousFirst,
        Offsettable newFirst, String offsetName) {
      if (offset.pointsToThisOffsettable(previousFirst)) {
        offset.pointToNew(newFirst);
      } else {
        Log.errorAndQuit("Header " + offsetName + " offset not pointing at first element?");
      }
    }
  
    private void addTypeListsToMapFile(RawDexFile rawDexFile, Offsettable typeListOffsettable) {
      // Create a MapItem for the TypeLists
      MapItem typeListMapItem = new MapItem();
      typeListMapItem.offset = new Offset(false);
      typeListMapItem.offset.pointToNew(typeListOffsettable);
      typeListMapItem.type = MapItem.TYPE_TYPE_LIST;
      typeListMapItem.size = 1;
  
      // Insert into the MapList.
      // (first, find the MapItem that points to StringDataItems...)
      int idx = 0;
      for (MapItem mapItem : rawDexFile.mapList.mapItems) {
        if (mapItem.type == MapItem.TYPE_STRING_DATA_ITEM) {
          break;
        }
        idx++;
      }
      // (now insert the TypeList MapItem just before the StringDataItem one...)
      rawDexFile.mapList.mapItems.add(idx, typeListMapItem);
    }
  
    private void addFieldIdsToHeaderAndMapFile(RawDexFile rawDexFile,
        Offsettable fieldOffsettable) {
      // Add the field IDs to the header.
      rawDexFile.header.fieldIdsOff.unsetNullAndPointTo(fieldOffsettable);
      rawDexFile.header.fieldIdsSize = 1;
  
      // Create a MapItem for the field IDs.
      MapItem fieldMapItem = new MapItem();
      fieldMapItem.offset = new Offset(false);
      fieldMapItem.offset.pointToNew(fieldOffsettable);
      fieldMapItem.type = MapItem.TYPE_FIELD_ID_ITEM;
      fieldMapItem.size = 1;
  
      // Insert into the MapList.
      // (first, find the MapItem that points to MethodIdItems...)
      int idx = 0;
      for (MapItem mapItem : rawDexFile.mapList.mapItems) {
        if (mapItem.type == MapItem.TYPE_METHOD_ID_ITEM) {
          break;
        }
        idx++;
      }
      // (now insert the FieldIdItem MapItem just before the MethodIdItem one...)
      rawDexFile.mapList.mapItems.add(idx, fieldMapItem);
    }
  
  
    private void updateOffsetsInHeaderAndMapFile(RawDexFile rawDexFile,
        Offsettable newFirstOffsettable) {
      Offsettable prevFirstOffsettable = null;
      for (int i = 0; i < offsettableTable.size(); i++) {
        if (offsettableTable.get(i) == newFirstOffsettable) {
          prevFirstOffsettable = offsettableTable.get(i + 1);
          break;
        }
      }
      if (prevFirstOffsettable == null) {
        Log.errorAndQuit("When calling updateMapListOffsets, could not find new "
            + "first offsettable?");
      }
  
      // Based on the type of the item we just added, check the relevant Offset in the header
      // and if it pointed at the prev_first_offsettable, make it point at the new one.
      // NB: if it isn't pointing at the prev one, something is wrong.
      HeaderItem header = rawDexFile.header;
      if (newFirstOffsettable.getItem() instanceof StringIdItem) {
        updateHeaderOffsetIfValid(header.stringIdsOff, prevFirstOffsettable,
            newFirstOffsettable, "StringID");
      } else if (newFirstOffsettable.getItem() instanceof TypeIdItem) {
        updateHeaderOffsetIfValid(header.typeIdsOff, prevFirstOffsettable,
            newFirstOffsettable, "TypeID");
      } else if (newFirstOffsettable.getItem() instanceof ProtoIdItem) {
        updateHeaderOffsetIfValid(header.protoIdsOff, prevFirstOffsettable,
            newFirstOffsettable, "ProtoID");
      } else if (newFirstOffsettable.getItem() instanceof FieldIdItem) {
        updateHeaderOffsetIfValid(header.fieldIdsOff, prevFirstOffsettable,
            newFirstOffsettable, "FieldID");
      } else if (newFirstOffsettable.getItem() instanceof MethodIdItem) {
        updateHeaderOffsetIfValid(header.methodIdsOff, prevFirstOffsettable,
            newFirstOffsettable, "MethodID");
      } else if (newFirstOffsettable.getItem() instanceof ClassDefItem) {
        updateHeaderOffsetIfValid(header.classDefsOff, prevFirstOffsettable,
            newFirstOffsettable, "ClassDef");
      }
  
      // Now iterate through the MapList's MapItems, and see if their Offsets pointed at the
      // prev_first_offsettable, and if so, make them now point at the new_first_offsettable.
      for (MapItem mapItem : rawDexFile.mapList.mapItems) {
        if (mapItem.offset.pointsToThisOffsettable(prevFirstOffsettable)) {
          Log.info("Updating offset in MapItem (type: " + mapItem.type + ") after "
              + "we called insertNewOffsettableAsFirstOfType()");
          mapItem.offset.pointToNew(newFirstOffsettable);
        }
      }
    }
  
    private void insertOffsettableAt(int idx, Offsettable offsettable) {
      offsettableTable.add(idx, offsettable);
      if (indexAfterMapList > idx) {
        indexAfterMapList++;
      }
      if (restorePoint > idx) {
        restorePoint++;
      }
    }
  
    /**
     * If we're creating our first TypeList, then IdCreator has to call this method to
     * ensure it gets put into the correct place in the offsettable table.
     * This assumes TypeLists always come before StringDatas.
     */
    public Offsettable insertNewOffsettableAsFirstEverTypeList(RawDexObject item,
        RawDexFile rawDexFile) {
      // We find the first StringDataItem, the type lists will come before this.
      Log.info("Calling insertNewOffsettableAsFirstEverTypeList()");
      for (int i = 0; i < offsettableTable.size(); i++) {
        if (offsettableTable.get(i).getItem() instanceof StringDataItem) {
          Offsettable offsettable = new Offsettable(item, true);
          insertOffsettableAt(i, offsettable);
          addTypeListsToMapFile(rawDexFile, offsettable);
          return offsettable;
        }
      }
      Log.errorAndQuit("Could not find any StringDataItems to insert the type list before.");
      return null;
    }
  
    /**
     * If we're creating our first FieldId, then IdCreator has to call this method to
     * ensure it gets put into the correct place in the offsettable table.
     * This assumes FieldIds always come before MethodIds.
     */
    public Offsettable insertNewOffsettableAsFirstEverField(RawDexObject item,
        RawDexFile rawDexFile) {
      // We find the first MethodIdItem, the fields will come before this.
      Log.info("Calling insertNewOffsettableAsFirstEverField()");
      for (int i = 0; i < offsettableTable.size(); i++) {
        if (offsettableTable.get(i).getItem() instanceof MethodIdItem) {
          Offsettable offsettable = new Offsettable(item, true);
          insertOffsettableAt(i, offsettable);
          addFieldIdsToHeaderAndMapFile(rawDexFile, offsettable);
          return offsettable;
        }
      }
      Log.errorAndQuit("Could not find any MethodIdItems to insert the field before.");
      return null;
    }
  
    /**
     * If we're creating a new Item (such as FieldId, MethodId) that is placed into the
     * first position of the relevant ID table, then IdCreator has to call this method to
     * ensure it gets put into the correct place in the offsettable table.
     */
    public Offsettable insertNewOffsettableAsFirstOfType(RawDexObject item,
        RawDexFile rawDexFile) {
      Log.debug("Calling insertNewOffsettableAsFirstOfType()");
      int index = getOffsettableIndexForFirstItemType(item);
      if (index == -1) {
        Log.errorAndQuit("Could not find any object of class: " + item.getClass());
      }
      Offsettable offsettable = new Offsettable(item, true);
      insertOffsettableAt(index, offsettable);
      updateOffsetsInHeaderAndMapFile(rawDexFile, offsettable);
      return offsettable;
    }
  
    /**
     * IdCreator has to call this method when it creates a new IdItem, to make sure it
     * gets put into the correct place in the offsettable table. IdCreator should
     * provide the IdItem that should come before this new IdItem.
     */
    public Offsettable insertNewOffsettableAfter(RawDexObject item, RawDexObject itemBefore) {
      Log.debug("Calling insertNewOffsettableAfter()");
      int index = getOffsettableIndexForItem(itemBefore);
      if (index == -1) {
        Log.errorAndQuit("Did not find specified 'after' object in offsettable table.");
      }
      Offsettable offsettable = new Offsettable(item, true);
      insertOffsettableAt(index + 1, offsettable);
      return offsettable;
    }
  
    private int getOffsettableIndexForFirstItemType(RawDexObject item) {
      Class<?> itemClass = item.getClass();
      for (int i = 0; i < offsettableTable.size(); i++) {
        if (offsettableTable.get(i).getItem().getClass().equals(itemClass)) {
          return i;
        }
      }
      return -1;
    }
  
    private int getOffsettableIndexForItem(RawDexObject item) {
      for (int i = 0; i < offsettableTable.size(); i++) {
        if (offsettableTable.get(i).getItem() == item) {
          return i;
        }
      }
      return -1;
    }
  
    /**
     * Given a RawDexObject, get the Offsettable that contains it.
     */
    public Offsettable getOffsettableForItem(RawDexObject item) {
      for (int i = 0; i < offsettableTable.size(); i++) {
        if (offsettableTable.get(i).getItem() == item) {
          return offsettableTable.get(i);
        }
      }
      return null;
    }
  }
  