add bonus data feature to imgdiff/imgpatch/applypatch

The bonus data option lets you give an additional blob of uncompressed
data to be used when constructing a patch for chunk #1 of an image.
The same blob must be available at patch time, and can be passed to
the command-line applypatch tool (this feature is not accessible from
edify scripts).

This will be used to reduce the size of recovery-from-boot patches by
storing parts of the recovery ramdisk (the UI images) on the system
partition.

Change-Id: Iac1959cdf7f5e4582f8d434e83456e483b64c02c

6 files changed, 100 insertions, 22 deletions

diff --git a/applypatch/applypatch.c b/applypatch/applypatch.c
index 488fd8c..7b8a010 100644
--- a/applypatch/applypatch.c
+++ b/applypatch/applypatch.c
@@ -39,7 +39,8 @@ static int GenerateTarget(FileContents* source_file,
                           const char* source_filename,
                           const char* target_filename,
                           const uint8_t target_sha1[SHA_DIGEST_SIZE],
-                          size_t target_size);
+                          size_t target_size,
+                          const Value* bonus_data);
 
 static int mtd_partitions_scanned = 0;
 
@@ -617,7 +618,8 @@ int applypatch(const char* source_filename,
                size_t target_size,
                int num_patches,
                char** const patch_sha1_str,
-               Value** patch_data) {
+               Value** patch_data,
+               Value* bonus_data) {
     printf("\napplying patch to %s\n", source_filename);
 
     if (target_filename[0] == '-' &&
@@ -699,7 +701,7 @@ int applypatch(const char* source_filename,
     int result = GenerateTarget(&source_file, source_patch_value,
                                 &copy_file, copy_patch_value,
                                 source_filename, target_filename,
-                                target_sha1, target_size);
+                                target_sha1, target_size, bonus_data);
     free(source_file.data);
     free(copy_file.data);
 
@@ -713,7 +715,8 @@ static int GenerateTarget(FileContents* source_file,
                           const char* source_filename,
                           const char* target_filename,
                           const uint8_t target_sha1[SHA_DIGEST_SIZE],
-                          size_t target_size) {
+                          size_t target_size,
+                          const Value* bonus_data) {
     int retry = 1;
     SHA_CTX ctx;
     int output;
@@ -867,7 +870,7 @@ static int GenerateTarget(FileContents* source_file,
         } else if (header_bytes_read >= 8 &&
                    memcmp(header, "IMGDIFF2", 8) == 0) {
             result = ApplyImagePatch(source_to_use->data, source_to_use->size,
-                                     patch, sink, token, &ctx);
+                                     patch, sink, token, &ctx, bonus_data);
         } else {
             printf("Unknown patch file format\n");
             return 1;
diff --git a/applypatch/applypatch.h b/applypatch/applypatch.h
index fb58843..d1a0232 100644
--- a/applypatch/applypatch.h
+++ b/applypatch/applypatch.h
@@ -55,7 +55,8 @@ int applypatch(const char* source_filename,
                size_t target_size,
                int num_patches,
                char** const patch_sha1_str,
-               Value** patch_data);
+               Value** patch_data,
+               Value* bonus_data);
 int applypatch_check(const char* filename,
                      int num_patches,
                      char** const patch_sha1_str);
@@ -79,7 +80,8 @@ int ApplyBSDiffPatchMem(const unsigned char* old_data, ssize_t old_size,
 // imgpatch.c
 int ApplyImagePatch(const unsigned char* old_data, ssize_t old_size,
                     const Value* patch,
-                    SinkFn sink, void* token, SHA_CTX* ctx);
+                    SinkFn sink, void* token, SHA_CTX* ctx,
+                    const Value* bonus_data);
 
 // freecache.c
 int MakeFreeSpaceOnCache(size_t bytes_needed);
diff --git a/applypatch/imgdiff.c b/applypatch/imgdiff.c
index 6b9ebee..05c4f25 100644
--- a/applypatch/imgdiff.c
+++ b/applypatch/imgdiff.c
@@ -111,6 +111,14 @@
  *
  * After the header there are 'chunk count' bsdiff patches; the offset
  * of each from the beginning of the file is specified in the header.
+ *
+ * This tool can take an optional file of "bonus data".  This is an
+ * extra file of data that is appended to chunk #1 after it is
+ * compressed (it must be a CHUNK_DEFLATE chunk).  The same file must
+ * be available (and passed to applypatch with -b) when applying the
+ * patch.  This is used to reduce the size of recovery-from-boot
+ * patches by combining the boot image with recovery ramdisk
+ * information that is stored on the system partition.
  */
 
 #include <errno.h>
@@ -772,21 +780,45 @@ void DumpChunks(ImageChunk* chunks, int num_chunks) {
 }
 
 int main(int argc, char** argv) {
-  if (argc != 4 && argc != 5) {
-    usage:
-    printf("usage: %s [-z] <src-img> <tgt-img> <patch-file>\n",
-            argv[0]);
-    return 2;
-  }
-
   int zip_mode = 0;
 
-  if (strcmp(argv[1], "-z") == 0) {
+  if (argc >= 2 && strcmp(argv[1], "-z") == 0) {
     zip_mode = 1;
     --argc;
     ++argv;
   }
 
+  size_t bonus_size = 0;
+  unsigned char* bonus_data = NULL;
+  if (argc >= 3 && strcmp(argv[1], "-b") == 0) {
+    struct stat st;
+    if (stat(argv[2], &st) != 0) {
+      printf("failed to stat bonus file %s: %s\n", argv[2], strerror(errno));
+      return 1;
+    }
+    bonus_size = st.st_size;
+    bonus_data = malloc(bonus_size);
+    FILE* f = fopen(argv[2], "rb");
+    if (f == NULL) {
+      printf("failed to open bonus file %s: %s\n", argv[2], strerror(errno));
+      return 1;
+    }
+    if (fread(bonus_data, 1, bonus_size, f) != bonus_size) {
+      printf("failed to read bonus file %s: %s\n", argv[2], strerror(errno));
+      return 1;
+    }
+    fclose(f);
+
+    argc -= 2;
+    argv += 2;
+  }
+
+  if (argc != 4) {
+    usage:
+    printf("usage: %s [-z] [-b <bonus-file>] <src-img> <tgt-img> <patch-file>\n",
+            argv[0]);
+    return 2;
+  }
 
   int num_src_chunks;
   ImageChunk* src_chunks;
@@ -909,6 +941,8 @@ int main(int argc, char** argv) {
   // Compute bsdiff patches for each chunk's data (the uncompressed
   // data, in the case of deflate chunks).
 
+  DumpChunks(src_chunks, num_src_chunks);
+
   printf("Construct patches for %d chunks...\n", num_tgt_chunks);
   unsigned char** patch_data = malloc(num_tgt_chunks * sizeof(unsigned char*));
   size_t* patch_size = malloc(num_tgt_chunks * sizeof(size_t));
@@ -923,6 +957,13 @@ int main(int argc, char** argv) {
         patch_data[i] = MakePatch(src_chunks, tgt_chunks+i, patch_size+i);
       }
     } else {
+      if (i == 1 && bonus_data) {
+        printf("  using %d bytes of bonus data for chunk %d\n", bonus_size, i);
+        src_chunks[i].data = realloc(src_chunks[i].data, src_chunks[i].len + bonus_size);
+        memcpy(src_chunks[i].data+src_chunks[i].len, bonus_data, bonus_size);
+        src_chunks[i].len += bonus_size;
+     }
+
       patch_data[i] = MakePatch(src_chunks+i, tgt_chunks+i, patch_size+i);
     }
     printf("patch %3d is %d bytes (of %d)\n",
diff --git a/applypatch/imgpatch.c b/applypatch/imgpatch.c
index e3ee80a..3a1df38 100644
--- a/applypatch/imgpatch.c
+++ b/applypatch/imgpatch.c
@@ -37,7 +37,8 @@
  */
 int ApplyImagePatch(const unsigned char* old_data, ssize_t old_size,
                     const Value* patch,
-                    SinkFn sink, void* token, SHA_CTX* ctx) {
+                    SinkFn sink, void* token, SHA_CTX* ctx,
+                    const Value* bonus_data) {
     ssize_t pos = 12;
     char* header = patch->data;
     if (patch->size < 12) {
@@ -123,6 +124,12 @@ int ApplyImagePatch(const unsigned char* old_data, ssize_t old_size,
             // Decompress the source data; the chunk header tells us exactly
             // how big we expect it to be when decompressed.
 
+            // Note: expanded_len will include the bonus data size if
+            // the patch was constructed with bonus data.  The
+            // deflation will come up 'bonus_size' bytes short; these
+            // must be appended from the bonus_data value.
+            size_t bonus_size = (i == 1 && bonus_data != NULL) ? bonus_data->size : 0;
+
             unsigned char* expanded_source = malloc(expanded_len);
             if (expanded_source == NULL) {
                 printf("failed to allocate %d bytes for expanded_source\n",
@@ -153,13 +160,19 @@ int ApplyImagePatch(const unsigned char* old_data, ssize_t old_size,
                 printf("source inflation returned %d\n", ret);
                 return -1;
             }
-            // We should have filled the output buffer exactly.
-            if (strm.avail_out != 0) {
-                printf("source inflation short by %d bytes\n", strm.avail_out);
+            // We should have filled the output buffer exactly, except
+            // for the bonus_size.
+            if (strm.avail_out != bonus_size) {
+                printf("source inflation short by %d bytes\n", strm.avail_out-bonus_size);
                 return -1;
             }
             inflateEnd(&strm);
 
+            if (bonus_size) {
+                memcpy(expanded_source + (expanded_len - bonus_size),
+                       bonus_data->data, bonus_size);
+            }
+
             // Next, apply the bsdiff patch (in memory) to the uncompressed
             // data.
             unsigned char* uncompressed_target_data;
diff --git a/applypatch/main.c b/applypatch/main.c
index 7025a2e..f61db5d 100644
--- a/applypatch/main.c
+++ b/applypatch/main.c
@@ -100,6 +100,21 @@ static int ParsePatchArgs(int argc, char** argv,
 }
 
 int PatchMode(int argc, char** argv) {
+    Value* bonus = NULL;
+    if (argc >= 3 && strcmp(argv[1], "-b") == 0) {
+        FileContents fc;
+        if (LoadFileContents(argv[2], &fc, RETOUCH_DONT_MASK) != 0) {
+            printf("failed to load bonus file %s\n", argv[2]);
+            return 1;
+        }
+        bonus = malloc(sizeof(Value));
+        bonus->type = VAL_BLOB;
+        bonus->size = fc.size;
+        bonus->data = (char*)fc.data;
+        argc -= 2;
+        argv += 2;
+    }
+
     if (argc < 6) {
         return 2;
     }
@@ -120,7 +135,7 @@ int PatchMode(int argc, char** argv) {
     }
 
     int result = applypatch(argv[1], argv[2], argv[3], target_size,
-                            num_patches, sha1s, patches);
+                            num_patches, sha1s, patches, bonus);
 
     int i;
     for (i = 0; i < num_patches; ++i) {
@@ -130,6 +145,10 @@ int PatchMode(int argc, char** argv) {
             free(p);
         }
     }
+    if (bonus) {
+        free(bonus->data);
+        free(bonus);
+    }
     free(sha1s);
     free(patches);
 
@@ -163,7 +182,7 @@ int main(int argc, char** argv) {
     if (argc < 2) {
       usage:
         printf(
-            "usage: %s <src-file> <tgt-file> <tgt-sha1> <tgt-size> "
+            "usage: %s [-b <bonus-file>] <src-file> <tgt-file> <tgt-sha1> <tgt-size> "
             "[<src-sha1>:<patch> ...]\n"
             "   or  %s -c <file> [<sha1> ...]\n"
             "   or  %s -s <bytes>\n"
diff --git a/updater/install.c b/updater/install.c
index ba27e9f..41f053d 100644
--- a/updater/install.c
+++ b/updater/install.c
@@ -901,7 +901,7 @@ Value* ApplyPatchFn(const char* name, State* state, int argc, Expr* argv[]) {
 
     int result = applypatch(source_filename, target_filename,
                             target_sha1, target_size,
-                            patchcount, patch_sha_str, patches);
+                            patchcount, patch_sha_str, patches, NULL);
 
     for (i = 0; i < patchcount; ++i) {
         FreeValue(patches[i]);