1 files changed, 101 insertions, 12 deletions

diff --git a/cgpt/cgpt_find.c b/cgpt/cgpt_find.c
index a8d31863..df6c4d39 100644
--- a/cgpt/cgpt_find.c
+++ b/cgpt/cgpt_find.c
@@ -83,27 +83,37 @@ static void showmatch(CgptFindParams *params, char *filename,
     EntryDetails(entry, partnum - 1, params->numeric);
 }
 
+// This needs to handle /dev/mmcblk0 -> /dev/mmcblk0p3, /dev/sda -> /dev/sda3
+static void mtd_showmatch(CgptFindParams *params, char *filename,
+                           int partnum, MtdDiskPartition *entry) {
+  char * format = "%s%d\n";
+  if (strncmp("/dev/mmcblk", filename, 11) == 0)
+    format = "%sp%d\n";
+  if (params->numeric)
+    printf("%d\n", partnum);
+  else
+    printf(format, filename, partnum);
+  if (params->verbose > 0)
+    MtdEntryDetails(entry, partnum - 1, params->numeric);
+}
+
 // This returns true if a GPT partition matches the search criteria. If a match
 // isn't found (or if the file doesn't contain a GPT), it returns false. The
 // filename and partition number that matched is left in a global, since we
 // could have multiple hits.
-static int do_search(CgptFindParams *params, char *fileName) {
-  int retval = 0;
+static int gpt_search(CgptFindParams *params, struct drive *drive,
+                      char *filename) {
   int i;
-  struct drive drive;
   GptEntry *entry;
+  int retval = 0;
   char partlabel[GPT_PARTNAME_LEN];
 
-  if (CGPT_OK != DriveOpen(fileName, &drive, O_RDONLY))
-    return 0;
-
-  if (GPT_SUCCESS != GptSanityCheck(&drive.gpt)) {
-    (void) DriveClose(&drive, 0);
+  if (GPT_SUCCESS != GptSanityCheck(&drive->gpt)) {
     return 0;
   }
 
-  for (i = 0; i < GetNumberOfEntries(&drive); ++i) {
-    entry = GetEntry(&drive.gpt, ANY_VALID, i);
+  for (i = 0; i < GetNumberOfEntries(drive); ++i) {
+    entry = GetEntry(&drive->gpt, ANY_VALID, i);
 
     if (GuidIsZero(&entry->type))
       continue;
@@ -122,15 +132,94 @@ static int do_search(CgptFindParams *params, char *fileName) {
       if (!strncmp(params->label, partlabel, sizeof(partlabel)))
         found = 1;
     }
-    if (found && match_content(params, &drive, entry)) {
+    if (found && match_content(params, drive, entry)) {
       params->hits++;
       retval++;
-      showmatch(params, fileName, i+1, entry);
+      showmatch(params, filename, i+1, entry);
       if (!params->match_partnum)
         params->match_partnum = i+1;
     }
   }
 
+  return retval;
+}
+
+static int mtd_match_type_to_guid(const MtdDiskPartition *e, const Guid *guid) {
+  return LookupMtdTypeForGuid(guid) == MtdGetEntryType(e);
+}
+
+static int mtd_match_content(CgptFindParams *params, struct drive *drive,
+                             MtdDiskPartition *entry) {
+  uint64_t part_size;
+
+  if (!params->matchlen)
+    return 1;
+
+  // Ensure that the region we want to match against is inside the partition.
+  part_size = LBA_SIZE * (entry->ending_lba - entry->starting_lba + 1);
+  if (params->matchoffset + params->matchlen > part_size) {
+    return 0;
+  }
+
+  // Read the partition data.
+  if (!FillBuffer(params,
+                  drive->fd,
+                  (LBA_SIZE * entry->starting_lba) + params->matchoffset,
+                  params->matchlen)) {
+    Error("unable to read partition data\n");
+    return 0;
+  }
+
+  // Compare it
+  if (0 == memcmp(params->matchbuf, params->comparebuf, params->matchlen)) {
+    return 1;
+  }
+
+  // Nope.
+  return 0;
+}
+
+static int mtd_search(CgptFindParams *params, struct drive *drive,
+                      char *filename) {
+  int i;
+  int retval = 0;
+  for (i = 0; i < GetNumberOfEntries(drive); ++i) {
+    MtdDiskPartition *e = MtdGetEntry(&drive->mtd, ANY_VALID, i);
+
+    if (IsUnused(drive, ANY_VALID, i))
+      continue;
+
+    int found = 0;
+
+    // Only searches by type are possible right now
+    if (params->set_type && mtd_match_type_to_guid(e, &params->type_guid)) {
+      found = 1;
+    }
+
+    if (found && mtd_match_content(params, drive, e)) {
+      params->hits++;
+      retval++;
+      mtd_showmatch(params, filename, i+1, e);
+      if (!params->match_partnum)
+        params->match_partnum = i+1;
+    }
+  }
+  return retval;
+}
+
+static int do_search(CgptFindParams *params, char *fileName) {
+  int retval;
+  struct drive drive;
+
+  if (CGPT_OK != DriveOpen(fileName, &drive, O_RDONLY))
+    return 0;
+
+  if (drive.is_mtd) {
+    retval = mtd_search(params, &drive, fileName);
+  } else {
+    retval = gpt_search(params, &drive, fileName);
+  }
+
   (void) DriveClose(&drive, 0);
 
   return retval;