[Zucchini] Introduce Imposed Ensemble Matcher.

Previously Zucchini-gen uses built-in heuristics to perform element
matching for ensemble patch generation. This CL adds an option
(accessible via the -impose parameter) to specify elements in "old"
and "new", and how they match. This allows the default heuristics to be
overridden, and enables external applications (who perhaps have better
ideas of element matching, e.g., have access to archiving programs) to
better use Zucchini to patch archives.

Zucchini-match is updated to prints the -impose command line to repeat
its results. Also, ElementMatch::ToString() is added.

Change-Id: I541b64722904c2fcd19ed75246d87e0268fbf86c
Reviewed-on: https://chromium-review.googlesource.com/1027191
Reviewed-by: Samuel Huang <huangs@chromium.org>
Reviewed-by: Greg Thompson <grt@chromium.org>
Commit-Queue: Samuel Huang <huangs@chromium.org>
Cr-Commit-Position: refs/heads/master@{#554909}
NOKEYCHECK=True
GitOrigin-RevId: 73a64ffde3f3b64df576aa1f2b5baebf7ec964ba

13 files changed, 569 insertions, 31 deletions

diff --git a/BUILD.gn b/BUILD.gn
index 1204cd4..800216e 100644
--- a/BUILD.gn
+++ b/BUILD.gn
@@ -42,6 +42,8 @@ static_library("zucchini_lib") {
     "image_index.cc",
     "image_index.h",
     "image_utils.h",
+    "imposed_ensemble_matcher.cc",
+    "imposed_ensemble_matcher.h",
     "io_utils.cc",
     "io_utils.h",
     "patch_reader.cc",
@@ -139,6 +141,7 @@ test("zucchini_unittests") {
     "equivalence_map_unittest.cc",
     "image_index_unittest.cc",
     "image_utils_unittest.cc",
+    "imposed_ensemble_matcher_unittest.cc",
     "io_utils_unittest.cc",
     "mapped_file_unittest.cc",
     "patch_read_write_unittest.cc",
diff --git a/image_utils.h b/image_utils.h
index 0f82140..de8cfee 100644
--- a/image_utils.h
+++ b/image_utils.h
@@ -8,9 +8,13 @@
 #include <stddef.h>
 #include <stdint.h>
 
+#include <string>
+
+#include "base/format_macros.h"
 #include "base/macros.h"
 #include "base/numerics/safe_conversions.h"
 #include "base/optional.h"
+#include "base/strings/stringprintf.h"
 #include "components/zucchini/buffer_view.h"
 #include "components/zucchini/typed_value.h"
 
@@ -209,6 +213,15 @@ struct ElementMatch {
   bool IsValid() const { return old_element.exe_type == new_element.exe_type; }
   ExecutableType exe_type() const { return old_element.exe_type; }
 
+  // Represents match as "#+#=#+#", where "#" denotes the integers:
+  //   [offset in "old", size in "old", offset in "new", size in "new"].
+  // Note that element type is omitted.
+  std::string ToString() const {
+    return base::StringPrintf("%" PRIuS "+%" PRIuS "=%" PRIuS "+%" PRIuS "",
+                              old_element.offset, old_element.size,
+                              new_element.offset, new_element.size);
+  }
+
   Element old_element;
   Element new_element;
 };
diff --git a/image_utils_unittest.cc b/image_utils_unittest.cc
index f483625..81695e9 100644
--- a/image_utils_unittest.cc
+++ b/image_utils_unittest.cc
@@ -21,4 +21,14 @@ TEST(ImageUtilsTest, CastExecutableTypeToString) {
   EXPECT_EQ("DEX ", CastExecutableTypeToString(kExeTypeDex));
 }
 
+TEST(ImageUtilsTest, ElementMatchToString) {
+  constexpr ExecutableType kAnyType = kExeTypeWin32X86;
+  EXPECT_EQ("1+2=3+4",
+            (ElementMatch{{{1, 2}, kAnyType}, {{3, 4}, kAnyType}}).ToString());
+  EXPECT_EQ(
+      "1000000000+1=0+1000000000",
+      (ElementMatch{{{1000000000, 1}, kAnyType}, {{0, 1000000000}, kAnyType}})
+          .ToString());
+}
+
 }  // namespace zucchini
diff --git a/imposed_ensemble_matcher.cc b/imposed_ensemble_matcher.cc
new file mode 100644
index 0000000..e735bc4
--- /dev/null
+++ b/imposed_ensemble_matcher.cc
@@ -0,0 +1,143 @@
+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "components/zucchini/imposed_ensemble_matcher.h"
+
+#include <algorithm>
+#include <sstream>
+#include <utility>
+
+#include "base/bind.h"
+#include "base/logging.h"
+#include "components/zucchini/io_utils.h"
+
+namespace zucchini {
+
+/******** ImposedMatchParser ********/
+
+ImposedMatchParser::ImposedMatchParser() = default;
+
+ImposedMatchParser::~ImposedMatchParser() = default;
+
+ImposedMatchParser::Status ImposedMatchParser::Parse(
+    std::string imposed_matches,
+    ConstBufferView old_image,
+    ConstBufferView new_image,
+    ElementDetector&& detector) {
+  CHECK(matches_.empty());
+  CHECK(bad_matches_.empty());
+
+  // Parse |imposed_matches| and check bounds.
+  std::istringstream iss(std::move(imposed_matches));
+  bool first = true;
+  iss.peek();  // Makes empty |iss| realize EOF is reached.
+  while (iss && !iss.eof()) {
+    // Eat delimiter.
+    if (first) {
+      first = false;
+    } else if (!(iss >> EatChar(','))) {
+      return kInvalidDelimiter;
+    }
+    // Extract parameters for one imposed match.
+    offset_t old_offset = 0U;
+    size_t old_size = 0U;
+    offset_t new_offset = 0U;
+    size_t new_size = 0U;
+    if (!(iss >> StrictUInt<offset_t>(old_offset) >> EatChar('+') >>
+          StrictUInt<size_t>(old_size) >> EatChar('=') >>
+          StrictUInt<offset_t>(new_offset) >> EatChar('+') >>
+          StrictUInt<size_t>(new_size))) {
+      return kParseError;
+    }
+    // Check bounds.
+    if (old_size == 0 || new_size == 0 ||
+        !old_image.covers({old_offset, old_size}) ||
+        !new_image.covers({new_offset, new_size})) {
+      return kOutOfBound;
+    }
+    matches_.push_back(
+        {{{old_offset, old_size}, kExeTypeUnknown},    // Assign type later.
+         {{new_offset, new_size}, kExeTypeUnknown}});  // Assign type later.
+  }
+  // Sort matches by "new" file offsets. This helps with overlap checks.
+  std::sort(matches_.begin(), matches_.end(),
+            [](const ElementMatch& match_a, const ElementMatch& match_b) {
+              return match_a.new_element.offset < match_b.new_element.offset;
+            });
+
+  // Check for overlaps in "new" file.
+  if (std::adjacent_find(
+          matches_.begin(), matches_.end(),
+          [](const ElementMatch& match1, const ElementMatch& match2) {
+            return match1.new_element.hi() > match2.new_element.lo();
+          }) != matches_.end()) {
+    return kOverlapInNew;
+  }
+
+  // Compute types and verify consistency. Remove identical matches and matches
+  // where any sub-image has an unknown type.
+  size_t write_idx = 0;
+  for (size_t read_idx = 0; read_idx < matches_.size(); ++read_idx) {
+    ConstBufferView old_sub_image(
+        old_image[matches_[read_idx].old_element.region()]);
+    ConstBufferView new_sub_image(
+        new_image[matches_[read_idx].new_element.region()]);
+    // Remove identical match.
+    if (old_sub_image.equals(new_sub_image)) {
+      ++num_identical_;
+      continue;
+    }
+    // Check executable types of sub-images.
+    base::Optional<Element> old_element = detector.Run(old_sub_image);
+    base::Optional<Element> new_element = detector.Run(new_sub_image);
+    if (!old_element || !new_element) {
+      // Skip unknown types, including those mixed with known types.
+      bad_matches_.push_back(matches_[read_idx]);
+      continue;
+    } else if (old_element->exe_type != new_element->exe_type) {
+      // Error if types are known, but inconsistent.
+      return kTypeMismatch;
+    }
+
+    // Keep match and remove gaps.
+    matches_[read_idx].old_element.exe_type = old_element->exe_type;
+    matches_[read_idx].new_element.exe_type = new_element->exe_type;
+    if (write_idx < read_idx)
+      matches_[write_idx] = matches_[read_idx];
+    ++write_idx;
+  }
+  matches_.resize(write_idx);
+  return kSuccess;
+}
+
+/******** ImposedEnsembleMatcher ********/
+
+ImposedEnsembleMatcher::ImposedEnsembleMatcher(
+    const std::string& imposed_matches)
+    : imposed_matches_(imposed_matches) {}
+
+ImposedEnsembleMatcher::~ImposedEnsembleMatcher() = default;
+
+bool ImposedEnsembleMatcher::RunMatch(ConstBufferView old_image,
+                                      ConstBufferView new_image) {
+  DCHECK(matches_.empty());
+  LOG(INFO) << "Start matching.";
+  ImposedMatchParser parser;
+  ImposedMatchParser::Status status =
+      parser.Parse(std::move(imposed_matches_), old_image, new_image,
+                   base::BindRepeating(DetectElementFromDisassembler));
+  // Print all warnings first.
+  for (const ElementMatch& bad_match : *parser.mutable_bad_matches())
+    LOG(WARNING) << "Skipped match with unknown type: " << bad_match.ToString();
+  if (status != ImposedMatchParser::kSuccess) {
+    LOG(ERROR) << "Imposed match failed with error code " << status << ".";
+    return false;
+  }
+  num_identical_ = parser.num_identical();
+  matches_ = std::move(*parser.mutable_matches());
+  Trim();
+  return true;
+}
+
+}  // namespace zucchini
diff --git a/imposed_ensemble_matcher.h b/imposed_ensemble_matcher.h
new file mode 100644
index 0000000..4dfc38e
--- /dev/null
+++ b/imposed_ensemble_matcher.h
@@ -0,0 +1,83 @@
+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef COMPONENTS_ZUCCHINI_IMPOSED_ENSEMBLE_MATCHER_H_
+#define COMPONENTS_ZUCCHINI_IMPOSED_ENSEMBLE_MATCHER_H_
+
+#include <stddef.h>
+
+#include <string>
+#include <vector>
+
+#include "base/macros.h"
+#include "components/zucchini/buffer_view.h"
+#include "components/zucchini/element_detection.h"
+#include "components/zucchini/ensemble_matcher.h"
+
+namespace zucchini {
+
+// A class to parse imposed match format, which is either an empty string (no
+// imposed patch), or a string formatted as:
+//   "#+#=#+#,#+#=#+#,..."  (e.g., "1+2=3+4", "1+2=3+4,5+6=7+8"),
+// where "#+#=#+#" encodes a match as 4 unsigned integers:
+//   [offset in "old", size in "old", offset in "new", size in "new"].
+class ImposedMatchParser {
+ public:
+  enum Status {
+    kSuccess,
+    kInvalidDelimiter,
+    kParseError,
+    kOutOfBound,
+    kOverlapInNew,
+    kTypeMismatch,
+  };
+
+  ImposedMatchParser();
+  ~ImposedMatchParser();
+
+  // Parses |imposed_matches| and writes the results to member variables.
+  // |old_image| and |new_image| are used for validation. Returns a Status value
+  // to signal success or various error modes. |detector| is used to validate
+  // Element types for matched pairs. This should only be called once for each
+  // instance.
+  Status Parse(std::string imposed_matches,
+               ConstBufferView old_image,
+               ConstBufferView new_image,
+               ElementDetector&& detector);
+
+  size_t num_identical() const { return num_identical_; }
+  std::vector<ElementMatch>* mutable_matches() { return &matches_; }
+  std::vector<ElementMatch>* mutable_bad_matches() { return &bad_matches_; }
+
+ private:
+  size_t num_identical_ = 0;
+  std::vector<ElementMatch> matches_;
+  // Stores "forgiven" bad matches, so the caller can impose matches for
+  // unsupported image types (which will simply be ignored). Note that imposing
+  // matches for known but incompatible image types would result in error.
+  std::vector<ElementMatch> bad_matches_;
+
+  DISALLOW_COPY_AND_ASSIGN(ImposedMatchParser);
+};
+
+// An ensemble matcher that parses a format string that describes matches.
+class ImposedEnsembleMatcher : public EnsembleMatcher {
+ public:
+  // |imposed_matches| specifies imposed maches, using a format described below.
+  // Validation is performed in RunMatch().
+  explicit ImposedEnsembleMatcher(const std::string& imposed_matches);
+  ~ImposedEnsembleMatcher() override;
+
+  // EnsembleMatcher:
+  bool RunMatch(ConstBufferView old_image, ConstBufferView new_image) override;
+
+ private:
+  const std::string imposed_matches_;
+
+  DISALLOW_COPY_AND_ASSIGN(ImposedEnsembleMatcher);
+};
+
+}  // namespace zucchini
+
+#endif  // COMPONENTS_ZUCCHINI_IMPOSED_ENSEMBLE_MATCHER_H_
diff --git a/imposed_ensemble_matcher_unittest.cc b/imposed_ensemble_matcher_unittest.cc
new file mode 100644
index 0000000..97a8898
--- /dev/null
+++ b/imposed_ensemble_matcher_unittest.cc
@@ -0,0 +1,214 @@
+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "components/zucchini/imposed_ensemble_matcher.h"
+
+#include "base/bind.h"
+#include "base/bind_helpers.h"
+#include "base/logging.h"
+#include "base/optional.h"
+#include "components/zucchini/buffer_view.h"
+#include "components/zucchini/disassembler.h"
+#include "components/zucchini/element_detection.h"
+#include "components/zucchini/image_utils.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace zucchini {
+
+namespace {
+
+// This test uses a mock archive format where regions are determined by their
+// consecutive byte values rather than parsing real executables. In fact, since
+// elements are imposed, only the first byte of the element is used to specify
+// executable type of the mock data:
+// - 'W' and 'w' specify kExeTypeWin32X86.
+// - 'E' and 'e' specify kExeTypeElfX86.
+// - Everything else specify kExeTypeUnknown.
+class TestElementDetector {
+ public:
+  TestElementDetector() {}
+
+  base::Optional<Element> Run(ConstBufferView image) const {
+    DCHECK_GT(image.size(), 0U);
+    char first_char = *image.begin();
+    if (first_char == 'W' || first_char == 'w')
+      return Element(image.local_region(), kExeTypeWin32X86);
+    if (first_char == 'E' || first_char == 'e')
+      return Element(image.local_region(), kExeTypeElfX86);
+    return base::nullopt;
+  }
+};
+
+}  // namespace
+
+TEST(ImposedMatchParserTest, ImposedMatchParser) {
+  std::vector<uint8_t> old_data;
+  std::vector<uint8_t> new_data;
+  auto populate = [](const std::string& s, std::vector<uint8_t>* data) {
+    for (char ch : s)
+      data->push_back(static_cast<uint8_t>(ch));
+  };
+  // Pos:             11111111
+  //        012345678901234567
+  populate("1WW222EEEE", &old_data);
+  populate("33eee2222222wwww44", &new_data);
+
+  ConstBufferView old_image(&old_data[0], old_data.size());
+  ConstBufferView new_image(&new_data[0], new_data.size());
+
+  TestElementDetector detector;
+
+  // Reusable output values.
+  std::string prev_imposed_matches;
+  ImposedMatchParser::Status status;
+  size_t num_identical;
+  std::vector<ElementMatch> matches;
+  std::vector<ElementMatch> bad_matches;
+
+  auto run_test = [&](const std::string& imposed_matches) -> bool {
+    prev_imposed_matches = imposed_matches;
+    status = ImposedMatchParser::kSuccess;
+    num_identical = 0;
+    matches.clear();
+    bad_matches.clear();
+    ImposedMatchParser parser;
+    status = parser.Parse(imposed_matches, old_image, new_image,
+                          base::BindRepeating(&TestElementDetector::Run,
+                                              base::Unretained(&detector)));
+    num_identical = parser.num_identical();
+    matches = std::move(*parser.mutable_matches());
+    bad_matches = std::move(*parser.mutable_bad_matches());
+    return status == ImposedMatchParser::kSuccess;
+  };
+
+  auto run_check = [&](const ElementMatch& match, ExecutableType exe_type,
+                       offset_t old_offset, size_t old_size,
+                       offset_t new_offset, size_t new_size) {
+    EXPECT_EQ(exe_type, match.exe_type()) << prev_imposed_matches;
+    EXPECT_EQ(exe_type, match.old_element.exe_type) << prev_imposed_matches;
+    EXPECT_EQ(old_offset, match.old_element.offset) << prev_imposed_matches;
+    EXPECT_EQ(old_size, match.old_element.size) << prev_imposed_matches;
+    EXPECT_EQ(exe_type, match.new_element.exe_type) << prev_imposed_matches;
+    EXPECT_EQ(new_offset, match.new_element.offset) << prev_imposed_matches;
+    EXPECT_EQ(new_size, match.new_element.size) << prev_imposed_matches;
+  };
+
+  // Empty string: Vacuous but valid.
+  EXPECT_TRUE(run_test(""));
+  EXPECT_EQ(0U, num_identical);
+  EXPECT_EQ(0U, matches.size());
+  EXPECT_EQ(0U, bad_matches.size());
+
+  // Full matches. Different permutations give same result.
+  for (const std::string& imposed_matches :
+       {"1+2=12+4,4+2=5+2,6+4=2+3", "1+2=12+4,6+4=2+3,4+2=5+2",
+        "4+2=5+2,1+2=12+4,6+4=2+3", "4+2=5+2,6+4=2+3,1+2=12+4",
+        "6+4=2+3,1+2=12+4,4+2=5+2", "6+4=2+3,1+2=12+4,4+2=5+2"}) {
+    EXPECT_TRUE(run_test(imposed_matches));
+    EXPECT_EQ(1U, num_identical);  // "4+2=5+2"
+    EXPECT_EQ(2U, matches.size());
+    // Results are sorted by "new" offsets.
+    run_check(matches[0], kExeTypeElfX86, 6, 4, 2, 3);
+    run_check(matches[1], kExeTypeWin32X86, 1, 2, 12, 4);
+    EXPECT_EQ(0U, bad_matches.size());
+  }
+
+  // Single subregion match.
+  EXPECT_TRUE(run_test("1+2=12+4"));
+  EXPECT_EQ(0U, num_identical);
+  EXPECT_EQ(1U, matches.size());
+  run_check(matches[0], kExeTypeWin32X86, 1, 2, 12, 4);
+  EXPECT_EQ(0U, bad_matches.size());
+
+  // Single subregion match. We're lax with redundant 0.
+  EXPECT_TRUE(run_test("6+04=02+10"));
+  EXPECT_EQ(0U, num_identical);
+  EXPECT_EQ(1U, matches.size());
+  run_check(matches[0], kExeTypeElfX86, 6, 4, 2, 10);
+  EXPECT_EQ(0U, bad_matches.size());
+
+  // Successive elements, no overlap.
+  EXPECT_TRUE(run_test("1+1=12+1,2+1=13+1"));
+  EXPECT_EQ(0U, num_identical);
+  EXPECT_EQ(2U, matches.size());
+  run_check(matches[0], kExeTypeWin32X86, 1, 1, 12, 1);
+  run_check(matches[1], kExeTypeWin32X86, 2, 1, 13, 1);
+  EXPECT_EQ(0U, bad_matches.size());
+
+  // Overlap in "old" file is okay.
+  EXPECT_TRUE(run_test("1+2=12+2,1+2=14+2"));
+  EXPECT_EQ(0U, num_identical);
+  EXPECT_EQ(2U, matches.size());
+  run_check(matches[0], kExeTypeWin32X86, 1, 2, 12, 2);
+  run_check(matches[1], kExeTypeWin32X86, 1, 2, 14, 2);
+  EXPECT_EQ(0U, bad_matches.size());
+
+  // Entire files: Have unknown type, so are recognized as such, and ignored.
+  EXPECT_TRUE(run_test("0+10=0+18"));
+  EXPECT_EQ(0U, num_identical);
+  EXPECT_EQ(0U, matches.size());
+  EXPECT_EQ(1U, bad_matches.size());
+  run_check(bad_matches[0], kExeTypeUnknown, 0, 10, 0, 18);
+
+  // Forgive matches that mix known type with unknown type.
+  EXPECT_TRUE(run_test("1+2=0+18"));
+  EXPECT_EQ(0U, num_identical);
+  EXPECT_EQ(0U, matches.size());
+  EXPECT_EQ(1U, bad_matches.size());
+  run_check(bad_matches[0], kExeTypeUnknown, 1, 2, 0, 18);
+
+  EXPECT_TRUE(run_test("0+10=12+4"));
+  EXPECT_EQ(0U, num_identical);
+  EXPECT_EQ(0U, matches.size());
+  EXPECT_EQ(1U, bad_matches.size());
+  run_check(bad_matches[0], kExeTypeUnknown, 0, 10, 12, 4);
+
+  // Test invalid delimiter.
+  for (const std::string& imposed_matches :
+       {"1+2=12+4,4+2=5+2x", "1+2=12+4 4+2=5+2", "1+2=12+4,4+2=5+2 ",
+        "1+2=12+4 "}) {
+    EXPECT_FALSE(run_test(imposed_matches));
+    EXPECT_EQ(ImposedMatchParser::kInvalidDelimiter, status);
+  }
+
+  // Test parse errors, including uint32_t overflow.
+  for (const std::string& imposed_matches :
+       {"x1+2=12+4,4+2=5+2,6+4=2+3", "x1+2=12+4,4+2=5+2,6+4=2+3x", ",", " ",
+        "+2=12+4", "1+2+12+4", "1=2+12+4", " 1+2=12+4", "1+2= 12+4", "1", "1+2",
+        "1+2=", "1+2=12", "1+2=12+", "4294967296+2=12+4"}) {
+    EXPECT_FALSE(run_test(imposed_matches));
+    EXPECT_EQ(ImposedMatchParser::kParseError, status);
+  }
+
+  // Test bound errors, include 0-size.
+  for (const std::string& imposed_matches :
+       {"1+10=12+4", "1+2=12+7", "0+11=0+18", "0+12=0+17", "10+1=0+18",
+        "0+10=18+1", "0+0=0+18", "0+10=0+0", "1000000000+1=0+1000000000"}) {
+    EXPECT_FALSE(run_test(imposed_matches));
+    EXPECT_EQ(ImposedMatchParser::kOutOfBound, status);
+  }
+
+  // Test overlap errors. Matches that get ignored are still tested.
+  for (const std::string& imposed_matches :
+       {"1+2=12+4,4+2=5+2,6+4=2+4", "0+10=0+18,1+2=12+4", "6+4=2+10,3+2=5+2"}) {
+    EXPECT_FALSE(run_test(imposed_matches));
+    EXPECT_EQ(ImposedMatchParser::kOverlapInNew, status);
+  }
+
+  // Test type mismatch errors.
+  EXPECT_FALSE(run_test("1+2=2+3"));
+  EXPECT_EQ(ImposedMatchParser::kTypeMismatch, status);
+
+  EXPECT_FALSE(run_test("6+4=12+4"));
+  EXPECT_EQ(ImposedMatchParser::kTypeMismatch, status);
+}
+
+}  // namespace zucchini
diff --git a/main_utils.cc b/main_utils.cc
index 7a28388..b6b5642 100644
--- a/main_utils.cc
+++ b/main_utils.cc
@@ -63,13 +63,16 @@ struct Command {
 /******** List of Zucchini commands ********/
 
 constexpr Command kCommands[] = {
-    {"gen", "-gen <old_file> <new_file> <patch_file> [-raw] [-keep]", 3,
-     &MainGen},
+    {"gen",
+     "-gen <old_file> <new_file> <patch_file> [-raw] [-keep]"
+     " [-impose=#+#=#+#,#+#=#+#,...]",
+     3, &MainGen},
     {"apply", "-apply <old_file> <patch_file> <new_file> [-keep]", 3,
      &MainApply},
     {"read", "-read <exe> [-dump]", 1, &MainRead},
     {"detect", "-detect <archive_file> [-dd=format#]", 1, &MainDetect},
-    {"match", "-match <old_file> <new_file>", 2, &MainMatch},
+    {"match", "-match <old_file> <new_file> [-impose=#+#=#+#,#+#=#+#,...]", 2,
+     &MainMatch},
     {"crc32", "-crc32 <file>", 1, &MainCrc32},
 };
 
diff --git a/patch_reader.cc b/patch_reader.cc
index 7e0b815..179f0fd 100644
--- a/patch_reader.cc
+++ b/patch_reader.cc
@@ -4,7 +4,6 @@
 
 #include "components/zucchini/patch_reader.h"
 
-#include <algorithm>
 #include <type_traits>
 #include <utility>
 
diff --git a/zucchini.h b/zucchini.h
index 9100709..e9093eb 100644
--- a/zucchini.h
+++ b/zucchini.h
@@ -5,6 +5,8 @@
 #ifndef COMPONENTS_ZUCCHINI_ZUCCHINI_H_
 #define COMPONENTS_ZUCCHINI_ZUCCHINI_H_
 
+#include <string>
+
 #include "components/zucchini/buffer_view.h"
 #include "components/zucchini/patch_reader.h"
 #include "components/zucchini/patch_writer.h"
@@ -31,12 +33,25 @@ enum Code {
 
 }  // namespace status
 
-// Generates ensemble patch from |old_image| to |new_image|, and writes it to
-// |patch_writer|.
+// Generates ensemble patch from |old_image| to |new_image| using the default
+// element detection and matching heuristics, writes the results to
+// |patch_writer|, and returns a status::Code.
 status::Code GenerateEnsemble(ConstBufferView old_image,
                               ConstBufferView new_image,
                               EnsemblePatchWriter* patch_writer);
 
+// Same as GenerateEnsemble(), but if |imposed_matches| is non-empty, then
+// overrides default element detection and matching heuristics with custom
+// element matching encoded in |imposed_matches|, which should be formatted as:
+//   "#+#=#+#,#+#=#+#,..."  (e.g., "1+2=3+4", "1+2=3+4,5+6=7+8"),
+// where "#+#=#+#" encodes a match as 4 unsigned integers:
+//   [offset in "old", size in "old", offset in "new", size in "new"].
+status::Code GenerateEnsembleWithImposedMatches(
+    ConstBufferView old_image,
+    ConstBufferView new_image,
+    std::string imposed_matches,
+    EnsemblePatchWriter* patch_writer);
+
 // Generates raw patch from |old_image| to |new_image|, and writes it to
 // |patch_writer|.
 status::Code GenerateRaw(ConstBufferView old_image,
diff --git a/zucchini_commands.cc b/zucchini_commands.cc
index 2d4b156..62dd20d 100644
--- a/zucchini_commands.cc
+++ b/zucchini_commands.cc
@@ -8,6 +8,7 @@
 #include <stdint.h>
 
 #include <ostream>
+#include <string>
 #include <utility>
 
 #include "base/command_line.h"
@@ -29,6 +30,7 @@ namespace {
 /******** Command-line Switches ********/
 
 constexpr char kSwitchDump[] = "dump";
+constexpr char kSwitchImpose[] = "impose";
 constexpr char kSwitchKeep[] = "keep";
 constexpr char kSwitchRaw[] = "raw";
 
@@ -56,11 +58,18 @@ zucchini::status::Code MainGen(MainParams params) {
   zucchini::EnsemblePatchWriter patch_writer(old_image.region(),
                                              new_image.region());
 
-  auto generate = params.command_line.HasSwitch(kSwitchRaw)
-                      ? zucchini::GenerateRaw
-                      : zucchini::GenerateEnsemble;
-  zucchini::status::Code result =
-      generate(old_image.region(), new_image.region(), &patch_writer);
+  zucchini::status::Code result = zucchini::status::kStatusSuccess;
+  if (params.command_line.HasSwitch(kSwitchRaw)) {
+    result = GenerateRaw(old_image.region(), new_image.region(), &patch_writer);
+  } else {
+    // May be empty.
+    std::string imposed_matches =
+        params.command_line.GetSwitchValueASCII(kSwitchImpose);
+    result = GenerateEnsembleWithImposedMatches(
+        old_image.region(), new_image.region(), std::move(imposed_matches),
+        &patch_writer);
+  }
+
   if (result != zucchini::status::kStatusSuccess) {
     params.out << "Fatal error encountered when generating patch." << std::endl;
     return result;
@@ -154,9 +163,13 @@ zucchini::status::Code MainMatch(MainParams params) {
                << new_image.error();
     return zucchini::status::kStatusFileReadError;
   }
+
+  std::string imposed_matches =
+      params.command_line.GetSwitchValueASCII(kSwitchImpose);
   zucchini::status::Code status =
       zucchini::MatchAll({old_image.data(), old_image.length()},
-                         {new_image.data(), new_image.length()}, params.out);
+                         {new_image.data(), new_image.length()},
+                         std::move(imposed_matches), params.out);
   if (status != zucchini::status::kStatusSuccess)
     params.err << "Fatal error found when matching executables." << std::endl;
   return status;
diff --git a/zucchini_gen.cc b/zucchini_gen.cc
index e88d37b..29be814 100644
--- a/zucchini_gen.cc
+++ b/zucchini_gen.cc
@@ -10,6 +10,7 @@
 #include <algorithm>
 #include <map>
 #include <memory>
+#include <string>
 #include <utility>
 
 #include "base/logging.h"
@@ -21,6 +22,7 @@
 #include "components/zucchini/equivalence_map.h"
 #include "components/zucchini/heuristic_ensemble_matcher.h"
 #include "components/zucchini/image_index.h"
+#include "components/zucchini/imposed_ensemble_matcher.h"
 #include "components/zucchini/patch_writer.h"
 #include "components/zucchini/reference_bytes_mixer.h"
 #include "components/zucchini/suffix_array.h"
@@ -317,13 +319,10 @@ bool GenerateExecutableElement(ExecutableType exe_type,
                           reference_bytes_mixer.get(), patch_writer);
 }
 
-/******** Exported Functions ********/
-
-status::Code GenerateEnsemble(ConstBufferView old_image,
-                              ConstBufferView new_image,
-                              EnsemblePatchWriter* patch_writer) {
-  std::unique_ptr<EnsembleMatcher> matcher =
-      std::make_unique<HeuristicEnsembleMatcher>(nullptr);
+status::Code GenerateEnsembleCommon(ConstBufferView old_image,
+                                    ConstBufferView new_image,
+                                    std::unique_ptr<EnsembleMatcher> matcher,
+                                    EnsemblePatchWriter* patch_writer) {
   if (!matcher->RunMatch(old_image, new_image)) {
     LOG(INFO) << "RunMatch() failed, generating raw patch.";
     return GenerateRaw(old_image, new_image, patch_writer);
@@ -420,6 +419,29 @@ status::Code GenerateEnsemble(ConstBufferView old_image,
   return status::kStatusSuccess;
 }
 
+/******** Exported Functions ********/
+
+status::Code GenerateEnsemble(ConstBufferView old_image,
+                              ConstBufferView new_image,
+                              EnsemblePatchWriter* patch_writer) {
+  return GenerateEnsembleCommon(
+      old_image, new_image, std::make_unique<HeuristicEnsembleMatcher>(nullptr),
+      patch_writer);
+}
+
+status::Code GenerateEnsembleWithImposedMatches(
+    ConstBufferView old_image,
+    ConstBufferView new_image,
+    std::string imposed_matches,
+    EnsemblePatchWriter* patch_writer) {
+  if (imposed_matches.empty())
+    return GenerateEnsemble(old_image, new_image, patch_writer);
+
+  return GenerateEnsembleCommon(
+      old_image, new_image,
+      std::make_unique<ImposedEnsembleMatcher>(imposed_matches), patch_writer);
+}
+
 status::Code GenerateRaw(ConstBufferView old_image,
                          ConstBufferView new_image,
                          EnsemblePatchWriter* patch_writer) {
diff --git a/zucchini_tools.cc b/zucchini_tools.cc
index 57ff0b2..5fcf066 100644
--- a/zucchini_tools.cc
+++ b/zucchini_tools.cc
@@ -10,7 +10,7 @@
 #include <algorithm>
 #include <memory>
 #include <ostream>
-#include <string>
+#include <utility>
 
 #include "base/bind.h"
 #include "base/logging.h"
@@ -19,6 +19,7 @@
 #include "components/zucchini/element_detection.h"
 #include "components/zucchini/ensemble_matcher.h"
 #include "components/zucchini/heuristic_ensemble_matcher.h"
+#include "components/zucchini/imposed_ensemble_matcher.h"
 #include "components/zucchini/io_utils.h"
 
 namespace zucchini {
@@ -44,11 +45,10 @@ status::Code ReadReferences(ConstBufferView image,
     targets.erase(std::unique(targets.begin(), targets.end()), targets.end());
     size_t num_targets = targets.size();
 
-    out << "Type " << int(group.type_tag().value());
-    out << ": Pool=" << static_cast<uint32_t>(group.pool_tag().value());
-    out << ", width=" << group.width();
-    out << ", #locations=" << num_locations;
-    out << ", #targets=" << num_targets;
+    out << "Type " << int(group.type_tag().value())
+        << ": Pool=" << static_cast<uint32_t>(group.pool_tag().value())
+        << ", width=" << group.width() << ", #locations=" << num_locations
+        << ", #targets=" << num_targets;
     if (num_targets > 0) {
       double ratio = static_cast<double>(num_locations) / num_targets;
       out << " (ratio=" << base::StringPrintf("%.4f", ratio) << ")";
@@ -59,8 +59,8 @@ status::Code ReadReferences(ConstBufferView image,
       refs = group.GetReader(disasm.get());
 
       for (auto ref = refs->GetNext(); ref; ref = refs->GetNext()) {
-        out << "  " << AsHex<8>(ref->location);
-        out << " " << AsHex<8>(ref->target) << std::endl;
+        out << "  " << AsHex<8>(ref->location) << " " << AsHex<8>(ref->target)
+            << std::endl;
       }
     }
   }
@@ -112,14 +112,27 @@ status::Code DetectAll(ConstBufferView image,
 
 status::Code MatchAll(ConstBufferView old_image,
                       ConstBufferView new_image,
+                      std::string imposed_matches,
                       std::ostream& out) {
-  HeuristicEnsembleMatcher matcher(&out);
-  if (!matcher.RunMatch(old_image, new_image)) {
+  std::unique_ptr<EnsembleMatcher> matcher;
+  if (imposed_matches.empty()) {
+    matcher = std::make_unique<HeuristicEnsembleMatcher>(&out);
+  } else {
+    matcher =
+        std::make_unique<ImposedEnsembleMatcher>(std::move(imposed_matches));
+  }
+  if (!matcher->RunMatch(old_image, new_image)) {
     out << "RunMatch() failed.";
     return status::kStatusFatal;
   }
-  out << "Found " << matcher.matches().size() << " nontrivial matches and "
-      << matcher.num_identical() << " identical matches." << std::endl;
+  out << "Found " << matcher->matches().size() << " nontrivial matches and "
+      << matcher->num_identical() << " identical matches." << std::endl
+      << "To impose the same matches by command line, use: " << std::endl
+      << "  -impose=";
+  PrefixSep sep(",");
+  for (const ElementMatch& match : matcher->matches())
+    out << sep << match.ToString();
+  out << std::endl;
 
   return status::kStatusSuccess;
 }
diff --git a/zucchini_tools.h b/zucchini_tools.h
index 6268745..bf9a95c 100644
--- a/zucchini_tools.h
+++ b/zucchini_tools.h
@@ -6,6 +6,7 @@
 #define COMPONENTS_ZUCCHINI_ZUCCHINI_TOOLS_H_
 
 #include <iosfwd>
+#include <string>
 #include <vector>
 
 #include "components/zucchini/buffer_view.h"
@@ -29,8 +30,14 @@ status::Code DetectAll(ConstBufferView image,
                        std::vector<ConstBufferView>* sub_image_list);
 
 // Prints all matched regions from |old_image| to |new_image|.
+// |imposed_matches|, if non-empty, encodes custom element matching to override
+// the default element detection and matching heuristics, and is formatted as:
+//   "#+#=#+#,#+#=#+#,..."  (e.g., "1+2=3+4", "1+2=3+4,5+6=7+8"),
+// where "#+#=#+#" encodes a match as 4 unsigned integers:
+//   [offset in "old", size in "old", offset in "new", size in "new"].
 status::Code MatchAll(ConstBufferView old_image,
                       ConstBufferView new_image,
+                      std::string imposed_matches,
                       std::ostream& out);
 
 }  // namespace zucchini