1 files changed, 259 insertions, 0 deletions
diff --git a/tools/dexfuzz/src/dexfuzz/listeners/UniqueProgramTrackerListener.java b/tools/dexfuzz/src/dexfuzz/listeners/UniqueProgramTrackerListener.java
new file mode 100644
index 0000000..affaffc
--- /dev/null
+++ b/tools/dexfuzz/src/dexfuzz/listeners/UniqueProgramTrackerListener.java
@@ -0,0 +1,259 @@
+/*
+ * 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.listeners;
+
+import dexfuzz.Log;
+import dexfuzz.Options;
+import dexfuzz.executors.Executor;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Tracks unique programs and outputs. Also saves divergent programs!
+ */
+public class UniqueProgramTrackerListener extends BaseListener {
+  /**
+   * Map of unique program MD5 sums, mapped to times seen.
+   */
+  private Map<String, Integer> uniquePrograms;
+
+  /**
+   * Map of unique program outputs (MD5'd), mapped to times seen.
+   */
+  private Map<String, Integer> uniqueOutputs;
+
+  /**
+   * Used to remember the seed used to fuzz the fuzzed file, so we can save it with this
+   * seed as a name, if we find a divergence.
+   */
+  private long currentSeed;
+
+  /**
+   * Used to remember the name of the file we've fuzzed, so we can save it if we
+   * find a divergence.
+   */
+  private String fuzzedFile;
+
+  private MessageDigest digest;
+  private String databaseFile;
+
+  /**
+   * Save the database every X number of iterations.
+   */
+  private static final int saveDatabasePeriod = 20;
+
+  public UniqueProgramTrackerListener(String databaseFile) {
+    this.databaseFile = databaseFile;
+  }
+
+  @Override
+  public void handleSeed(long seed) {
+    currentSeed = seed;
+  }
+
+  /**
+   * Given a program filename, calculate the MD5sum of
+   * this program.
+   */
+  private String getMD5SumOfProgram(String programName) {
+    byte[] buf = new byte[256];
+    try {
+      FileInputStream stream = new FileInputStream(programName);
+      boolean done = false;
+      while (!done) {
+        int bytesRead = stream.read(buf);
+        if (bytesRead == -1) {
+          done = true;
+        } else {
+          digest.update(buf);
+        }
+      }
+      stream.close();
+    } catch (FileNotFoundException e) {
+      e.printStackTrace();
+    } catch (IOException e) {
+      e.printStackTrace();
+    }
+    return new String(digest.digest());
+  }
+
+  private String getMD5SumOfOutput(String output) {
+    digest.update(output.getBytes());
+    return new String(digest.digest());
+  }
+
+  @SuppressWarnings("unchecked")
+  private void loadUniqueProgsData() {
+    File file = new File(databaseFile);
+    if (!file.exists()) {
+      uniquePrograms = new HashMap<String, Integer>();
+      uniqueOutputs = new HashMap<String, Integer>();
+      return;
+    }
+
+    try {
+      ObjectInputStream objectStream =
+          new ObjectInputStream(new FileInputStream(databaseFile));
+      uniquePrograms = (Map<String, Integer>) objectStream.readObject();
+      uniqueOutputs = (Map<String, Integer>) objectStream.readObject();
+      objectStream.close();
+    } catch (FileNotFoundException e) {
+      e.printStackTrace();
+    } catch (IOException e) {
+      e.printStackTrace();
+    } catch (ClassNotFoundException e) {
+      e.printStackTrace();
+    }
+
+  }
+
+  private void saveUniqueProgsData() {
+    // Since we could potentially stop the program while writing out this DB,
+    // copy the old file beforehand, and then delete it if we successfully wrote out the DB.
+    boolean oldWasSaved = false;
+    File file = new File(databaseFile);
+    if (file.exists()) {
+      try {
+        Process process =
+            Runtime.getRuntime().exec(String.format("cp %1$s %1$s.old", databaseFile));
+        // Shouldn't block, cp shouldn't produce output.
+        process.waitFor();
+        oldWasSaved = true;
+      } catch (IOException exception) {
+        exception.printStackTrace();
+      } catch (InterruptedException exception) {
+        exception.printStackTrace();
+      }
+    }
+
+    // Now write out the DB.
+    boolean success = false;
+    try {
+      ObjectOutputStream objectStream =
+          new ObjectOutputStream(new FileOutputStream(databaseFile));
+      objectStream.writeObject(uniquePrograms);
+      objectStream.writeObject(uniqueOutputs);
+      objectStream.close();
+      success = true;
+    } catch (FileNotFoundException e) {
+      e.printStackTrace();
+    } catch (IOException e) {
+      e.printStackTrace();
+    }
+
+    // If we get here, and we successfully wrote out the DB, delete the saved one.
+    if (oldWasSaved && success) {
+      try {
+        Process process =
+            Runtime.getRuntime().exec(String.format("rm %s.old", databaseFile));
+        // Shouldn't block, rm shouldn't produce output.
+        process.waitFor();
+      } catch (IOException exception) {
+        exception.printStackTrace();
+      } catch (InterruptedException exception) {
+        exception.printStackTrace();
+      }
+    } else if (oldWasSaved && !success) {
+      Log.error("Failed to successfully write out the unique programs DB!");
+      Log.error("Old DB should be saved in " + databaseFile + ".old");
+    }
+  }
+
+  private void addToMap(String md5sum, Map<String, Integer> map) {
+    if (map.containsKey(md5sum)) {
+      map.put(md5sum, map.get(md5sum) + 1);
+    } else {
+      map.put(md5sum, 1);
+    }
+  }
+
+  private void saveDivergentProgram() {
+    File before = new File(fuzzedFile);
+    File after = new File(String.format("divergent_programs/%d.dex", currentSeed));
+    boolean success = before.renameTo(after);
+    if (!success) {
+      Log.error("Failed to save divergent program! Does divergent_programs/ exist?");
+    }
+  }
+
+  @Override
+  public void setup() {
+    try {
+      digest = MessageDigest.getInstance("MD5");
+      loadUniqueProgsData();
+    } catch (NoSuchAlgorithmException e) {
+      e.printStackTrace();
+    }
+  }
+
+  @Override
+  public void handleIterationFinished(int iteration) {
+    if ((iteration % saveDatabasePeriod) == (saveDatabasePeriod - 1)) {
+      saveUniqueProgsData();
+    }
+  }
+
+  @Override
+  public void handleSuccessfullyFuzzedFile(String programName) {
+    String md5sum = getMD5SumOfProgram(programName);
+    addToMap(md5sum, uniquePrograms);
+
+    fuzzedFile = programName;
+  }
+
+  @Override
+  public void handleDivergences(Map<String, List<Executor>> outputMap) {
+    // Just use the first one.
+    String output = (String) outputMap.keySet().toArray()[0];
+    String md5sum = getMD5SumOfOutput(output);
+    addToMap(md5sum, uniqueOutputs);
+
+    saveDivergentProgram();
+  }
+
+  @Override
+  public void handleSuccess(Map<String, List<Executor>> outputMap) {
+    // There's only one, use it.
+    String output = (String) outputMap.keySet().toArray()[0];
+    String md5sum = getMD5SumOfOutput(output);
+    addToMap(md5sum, uniqueOutputs);
+  }
+
+  @Override
+  public void handleSummary() {
+    if (Options.reportUnique) {
+      Log.always("-- UNIQUE PROGRAM REPORT --");
+      Log.always("Unique Programs Seen: " + uniquePrograms.size());
+      Log.always("Unique Outputs Seen: " + uniqueOutputs.size());
+      Log.always("---------------------------");
+    }
+
+    saveUniqueProgsData();
+  }
+
+}