1 files changed, 122 insertions, 101 deletions

diff --git a/applypatch/applypatch.cpp b/applypatch/applypatch.cpp
index 96bd88e88..1767761a8 100644
--- a/applypatch/applypatch.cpp
+++ b/applypatch/applypatch.cpp
@@ -25,10 +25,13 @@
 #include <sys/types.h>
 #include <unistd.h>
 
+#include <base/strings.h>
+
 #include "mincrypt/sha.h"
 #include "applypatch.h"
 #include "mtdutils/mtdutils.h"
 #include "edify/expr.h"
+#include "print_sha1.h"
 
 static int LoadPartitionContents(const char* filename, FileContents* file);
 static ssize_t FileSink(const unsigned char* data, ssize_t len, void* token);
@@ -42,7 +45,7 @@ static int GenerateTarget(FileContents* source_file,
                           size_t target_size,
                           const Value* bonus_data);
 
-static int mtd_partitions_scanned = 0;
+static bool mtd_partitions_scanned = false;
 
 // Read a file into memory; store the file contents and associated
 // metadata in *file.
@@ -87,21 +90,6 @@ int LoadFileContents(const char* filename, FileContents* file) {
     return 0;
 }
 
-static size_t* size_array;
-// comparison function for qsort()ing an int array of indexes into
-// size_array[].
-static int compare_size_indices(const void* a, const void* b) {
-    const int aa = *reinterpret_cast<const int*>(a);
-    const int bb = *reinterpret_cast<const int*>(b);
-    if (size_array[aa] < size_array[bb]) {
-        return -1;
-    } else if (size_array[aa] > size_array[bb]) {
-        return 1;
-    } else {
-        return 0;
-    }
-}
-
 // Load the contents of an MTD or EMMC partition into the provided
 // FileContents.  filename should be a string of the form
 // "MTD:<partition_name>:<size_1>:<sha1_1>:<size_2>:<sha1_2>:..."  (or
@@ -120,53 +108,45 @@ static int compare_size_indices(const void* a, const void* b) {
 enum PartitionType { MTD, EMMC };
 
 static int LoadPartitionContents(const char* filename, FileContents* file) {
-    char* copy = strdup(filename);
-    const char* magic = strtok(copy, ":");
+    std::string copy(filename);
+    std::vector<std::string> pieces = android::base::Split(copy, ":");
+    if (pieces.size() < 4 || pieces.size() % 2 != 0) {
+        printf("LoadPartitionContents called with bad filename (%s)\n", filename);
+        return -1;
+    }
 
     enum PartitionType type;
-
-    if (strcmp(magic, "MTD") == 0) {
+    if (pieces[0] == "MTD") {
         type = MTD;
-    } else if (strcmp(magic, "EMMC") == 0) {
+    } else if (pieces[0] == "EMMC") {
         type = EMMC;
     } else {
         printf("LoadPartitionContents called with bad filename (%s)\n", filename);
         return -1;
     }
-    const char* partition = strtok(NULL, ":");
-
-    int i;
-    int colons = 0;
-    for (i = 0; filename[i] != '\0'; ++i) {
-        if (filename[i] == ':') {
-            ++colons;
-        }
-    }
-    if (colons < 3 || colons%2 == 0) {
-        printf("LoadPartitionContents called with bad filename (%s)\n",
-               filename);
-    }
+    const char* partition = pieces[1].c_str();
 
-    int pairs = (colons-1)/2;     // # of (size,sha1) pairs in filename
-    int* index = reinterpret_cast<int*>(malloc(pairs * sizeof(int)));
-    size_t* size = reinterpret_cast<size_t*>(malloc(pairs * sizeof(size_t)));
-    char** sha1sum = reinterpret_cast<char**>(malloc(pairs * sizeof(char*)));
+    size_t pairs = (pieces.size() - 2) / 2;    // # of (size, sha1) pairs in filename
+    std::vector<size_t> index(pairs);
+    std::vector<size_t> size(pairs);
+    std::vector<std::string> sha1sum(pairs);
 
-    for (i = 0; i < pairs; ++i) {
-        const char* size_str = strtok(NULL, ":");
-        size[i] = strtol(size_str, NULL, 10);
+    for (size_t i = 0; i < pairs; ++i) {
+        size[i] = strtol(pieces[i*2+2].c_str(), NULL, 10);
         if (size[i] == 0) {
             printf("LoadPartitionContents called with bad size (%s)\n", filename);
             return -1;
         }
-        sha1sum[i] = strtok(NULL, ":");
+        sha1sum[i] = pieces[i*2+3].c_str();
         index[i] = i;
     }
 
-    // sort the index[] array so it indexes the pairs in order of
-    // increasing size.
-    size_array = size;
-    qsort(index, pairs, sizeof(int), compare_size_indices);
+    // Sort the index[] array so it indexes the pairs in order of increasing size.
+    sort(index.begin(), index.end(),
+        [&](const size_t& i, const size_t& j) {
+            return (size[i] < size[j]);
+        }
+    );
 
     MtdReadContext* ctx = NULL;
     FILE* dev = NULL;
@@ -175,20 +155,18 @@ static int LoadPartitionContents(const char* filename, FileContents* file) {
         case MTD: {
             if (!mtd_partitions_scanned) {
                 mtd_scan_partitions();
-                mtd_partitions_scanned = 1;
+                mtd_partitions_scanned = true;
             }
 
             const MtdPartition* mtd = mtd_find_partition_by_name(partition);
             if (mtd == NULL) {
-                printf("mtd partition \"%s\" not found (loading %s)\n",
-                       partition, filename);
+                printf("mtd partition \"%s\" not found (loading %s)\n", partition, filename);
                 return -1;
             }
 
             ctx = mtd_read_partition(mtd);
             if (ctx == NULL) {
-                printf("failed to initialize read of mtd partition \"%s\"\n",
-                       partition);
+                printf("failed to initialize read of mtd partition \"%s\"\n", partition);
                 return -1;
             }
             break;
@@ -197,8 +175,7 @@ static int LoadPartitionContents(const char* filename, FileContents* file) {
         case EMMC:
             dev = fopen(partition, "rb");
             if (dev == NULL) {
-                printf("failed to open emmc partition \"%s\": %s\n",
-                       partition, strerror(errno));
+                printf("failed to open emmc partition \"%s\": %s\n", partition, strerror(errno));
                 return -1;
             }
     }
@@ -207,15 +184,15 @@ static int LoadPartitionContents(const char* filename, FileContents* file) {
     SHA_init(&sha_ctx);
     uint8_t parsed_sha[SHA_DIGEST_SIZE];
 
-    // allocate enough memory to hold the largest size.
+    // Allocate enough memory to hold the largest size.
     file->data = reinterpret_cast<unsigned char*>(malloc(size[index[pairs-1]]));
     char* p = (char*)file->data;
     file->size = 0;                // # bytes read so far
+    bool found = false;
 
-    for (i = 0; i < pairs; ++i) {
-        // Read enough additional bytes to get us up to the next size
-        // (again, we're trying the possibilities in order of increasing
-        // size).
+    for (size_t i = 0; i < pairs; ++i) {
+        // Read enough additional bytes to get us up to the next size. (Again,
+        // we're trying the possibilities in order of increasing size).
         size_t next = size[index[i]] - file->size;
         size_t read = 0;
         if (next > 0) {
@@ -245,8 +222,8 @@ static int LoadPartitionContents(const char* filename, FileContents* file) {
         memcpy(&temp_ctx, &sha_ctx, sizeof(SHA_CTX));
         const uint8_t* sha_so_far = SHA_final(&temp_ctx);
 
-        if (ParseSha1(sha1sum[index[i]], parsed_sha) != 0) {
-            printf("failed to parse sha1 %s in %s\n", sha1sum[index[i]], filename);
+        if (ParseSha1(sha1sum[index[i]].c_str(), parsed_sha) != 0) {
+            printf("failed to parse sha1 %s in %s\n", sha1sum[index[i]].c_str(), filename);
             free(file->data);
             file->data = NULL;
             return -1;
@@ -256,7 +233,8 @@ static int LoadPartitionContents(const char* filename, FileContents* file) {
             // we have a match.  stop reading the partition; we'll return
             // the data we've read so far.
             printf("partition read matched size %zu sha %s\n",
-                   size[index[i]], sha1sum[index[i]]);
+                   size[index[i]], sha1sum[index[i]].c_str());
+            found = true;
             break;
         }
 
@@ -274,9 +252,8 @@ static int LoadPartitionContents(const char* filename, FileContents* file) {
     }
 
 
-    if (i == pairs) {
-        // Ran off the end of the list of (size,sha1) pairs without
-        // finding a match.
+    if (!found) {
+        // Ran off the end of the list of (size,sha1) pairs without finding a match.
         printf("contents of partition \"%s\" didn't match %s\n", partition, filename);
         free(file->data);
         file->data = NULL;
@@ -293,11 +270,6 @@ static int LoadPartitionContents(const char* filename, FileContents* file) {
     file->st.st_uid = 0;
     file->st.st_gid = 0;
 
-    free(copy);
-    free(index);
-    free(size);
-    free(sha1sum);
-
     return 0;
 }
 
@@ -340,33 +312,34 @@ int SaveFileContents(const char* filename, const FileContents* file) {
 }
 
 // Write a memory buffer to 'target' partition, a string of the form
-// "MTD:<partition>[:...]" or "EMMC:<partition_device>:".  Return 0 on
-// success.
+// "MTD:<partition>[:...]" or "EMMC:<partition_device>[:...]". The target name
+// might contain multiple colons, but WriteToPartition() only uses the first
+// two and ignores the rest. Return 0 on success.
 int WriteToPartition(unsigned char* data, size_t len, const char* target) {
-    char* copy = strdup(target);
-    const char* magic = strtok(copy, ":");
+    std::string copy(target);
+    std::vector<std::string> pieces = android::base::Split(copy, ":");
+
+    if (pieces.size() < 2) {
+        printf("WriteToPartition called with bad target (%s)\n", target);
+        return -1;
+    }
 
     enum PartitionType type;
-    if (strcmp(magic, "MTD") == 0) {
+    if (pieces[0] == "MTD") {
         type = MTD;
-    } else if (strcmp(magic, "EMMC") == 0) {
+    } else if (pieces[0] == "EMMC") {
         type = EMMC;
     } else {
         printf("WriteToPartition called with bad target (%s)\n", target);
         return -1;
     }
-    const char* partition = strtok(NULL, ":");
-
-    if (partition == NULL) {
-        printf("bad partition target name \"%s\"\n", target);
-        return -1;
-    }
+    const char* partition = pieces[1].c_str();
 
     switch (type) {
         case MTD: {
             if (!mtd_partitions_scanned) {
                 mtd_scan_partitions();
-                mtd_partitions_scanned = 1;
+                mtd_partitions_scanned = true;
             }
 
             const MtdPartition* mtd = mtd_find_partition_by_name(partition);
@@ -410,7 +383,7 @@ int WriteToPartition(unsigned char* data, size_t len, const char* target) {
                 return -1;
             }
 
-            for (int attempt = 0; attempt < 2; ++attempt) {
+            for (size_t attempt = 0; attempt < 2; ++attempt) {
                 if (TEMP_FAILURE_RETRY(lseek(fd, start, SEEK_SET)) == -1) {
                     printf("failed seek on %s: %s\n", partition, strerror(errno));
                     return -1;
@@ -490,7 +463,7 @@ int WriteToPartition(unsigned char* data, size_t len, const char* target) {
                 }
 
                 if (start == len) {
-                    printf("verification read succeeded (attempt %d)\n", attempt+1);
+                    printf("verification read succeeded (attempt %zu)\n", attempt+1);
                     success = true;
                     break;
                 }
@@ -510,7 +483,6 @@ int WriteToPartition(unsigned char* data, size_t len, const char* target) {
         }
     }
 
-    free(copy);
     return 0;
 }
 
@@ -658,14 +630,6 @@ int CacheSizeCheck(size_t bytes) {
     }
 }
 
-static void print_short_sha1(const uint8_t sha1[SHA_DIGEST_SIZE]) {
-    const char* hex = "0123456789abcdef";
-    for (size_t i = 0; i < 4; ++i) {
-        putchar(hex[(sha1[i]>>4) & 0xf]);
-        putchar(hex[sha1[i] & 0xf]);
-    }
-}
-
 // This function applies binary patches to files in a way that is safe
 // (the original file is not touched until we have the desired
 // replacement for it) and idempotent (it's okay to run this program
@@ -678,7 +642,7 @@ static void print_short_sha1(const uint8_t sha1[SHA_DIGEST_SIZE]) {
 //   entries in <patch_sha1_str>, the corresponding patch from
 //   <patch_data> (which must be a VAL_BLOB) is applied to produce a
 //   new file (the type of patch is automatically detected from the
-//   blob daat).  If that new file has sha1 hash <target_sha1_str>,
+//   blob data).  If that new file has sha1 hash <target_sha1_str>,
 //   moves it to replace <target_filename>, and exits successfully.
 //   Note that if <source_filename> and <target_filename> are not the
 //   same, <source_filename> is NOT deleted on success.
@@ -689,7 +653,7 @@ static void print_short_sha1(const uint8_t sha1[SHA_DIGEST_SIZE]) {
 //   status.
 //
 // <source_filename> may refer to a partition to read the source data.
-// See the comments for the LoadPartition Contents() function above
+// See the comments for the LoadPartitionContents() function above
 // for the format of such a filename.
 
 int applypatch(const char* source_filename,
@@ -724,9 +688,7 @@ int applypatch(const char* source_filename,
         if (memcmp(source_file.sha1, target_sha1, SHA_DIGEST_SIZE) == 0) {
             // The early-exit case:  the patch was already applied, this file
             // has the desired hash, nothing for us to do.
-            printf("already ");
-            print_short_sha1(target_sha1);
-            putchar('\n');
+            printf("already %s\n", short_sha1(target_sha1).c_str());
             free(source_file.data);
             return 0;
         }
@@ -783,6 +745,67 @@ int applypatch(const char* source_filename,
     return result;
 }
 
+/*
+ * This function flashes a given image to the target partition. It verifies
+ * the target cheksum first, and will return if target has the desired hash.
+ * It checks the checksum of the given source image before flashing, and
+ * verifies the target partition afterwards. The function is idempotent.
+ * Returns zero on success.
+ */
+int applypatch_flash(const char* source_filename, const char* target_filename,
+                     const char* target_sha1_str, size_t target_size) {
+    printf("flash %s: ", target_filename);
+
+    uint8_t target_sha1[SHA_DIGEST_SIZE];
+    if (ParseSha1(target_sha1_str, target_sha1) != 0) {
+        printf("failed to parse tgt-sha1 \"%s\"\n", target_sha1_str);
+        return 1;
+    }
+
+    FileContents source_file;
+    source_file.data = NULL;
+    std::string target_str(target_filename);
+
+    std::vector<std::string> pieces = android::base::Split(target_str, ":");
+    if (pieces.size() != 2 || (pieces[0] != "MTD" && pieces[0] != "EMMC")) {
+        printf("invalid target name \"%s\"", target_filename);
+        return 1;
+    }
+
+    // Load the target into the source_file object to see if already applied.
+    pieces.push_back(std::to_string(target_size));
+    pieces.push_back(target_sha1_str);
+    std::string fullname = android::base::Join(pieces, ':');
+    if (LoadPartitionContents(fullname.c_str(), &source_file) == 0 &&
+        memcmp(source_file.sha1, target_sha1, SHA_DIGEST_SIZE) == 0) {
+        // The early-exit case: the image was already applied, this partition
+        // has the desired hash, nothing for us to do.
+        printf("already %s\n", short_sha1(target_sha1).c_str());
+        free(source_file.data);
+        return 0;
+    }
+
+    if (LoadFileContents(source_filename, &source_file) == 0) {
+        if (memcmp(source_file.sha1, target_sha1, SHA_DIGEST_SIZE) != 0) {
+            // The source doesn't have desired checksum.
+            printf("source \"%s\" doesn't have expected sha1 sum\n", source_filename);
+            printf("expected: %s, found: %s\n", short_sha1(target_sha1).c_str(),
+                    short_sha1(source_file.sha1).c_str());
+            free(source_file.data);
+            return 1;
+        }
+    }
+
+    if (WriteToPartition(source_file.data, target_size, target_filename) != 0) {
+        printf("write of copied data to %s failed\n", target_filename);
+        free(source_file.data);
+        return 1;
+    }
+
+    free(source_file.data);
+    return 0;
+}
+
 static int GenerateTarget(FileContents* source_file,
                           const Value* source_patch_value,
                           FileContents* copy_file,
@@ -983,9 +1006,7 @@ static int GenerateTarget(FileContents* source_file,
         printf("patch did not produce expected sha1\n");
         return 1;
     } else {
-        printf("now ");
-        print_short_sha1(target_sha1);
-        putchar('\n');
+        printf("now %s\n", short_sha1(target_sha1).c_str());
     }
 
     if (output < 0) {