1 /*
2  * Copyright (C) 2009 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 /*
18  * This program constructs binary patches for images -- such as boot.img
19  * and recovery.img -- that consist primarily of large chunks of gzipped
20  * data interspersed with uncompressed data.  Doing a naive bsdiff of
21  * these files is not useful because small changes in the data lead to
22  * large changes in the compressed bitstream; bsdiff patches of gzipped
23  * data are typically as large as the data itself.
24  *
25  * To patch these usefully, we break the source and target images up into
26  * chunks of two types: "normal" and "gzip".  Normal chunks are simply
27  * patched using a plain bsdiff.  Gzip chunks are first expanded, then a
28  * bsdiff is applied to the uncompressed data, then the patched data is
29  * gzipped using the same encoder parameters.  Patched chunks are
30  * concatenated together to create the output file; the output image
31  * should be *exactly* the same series of bytes as the target image used
32  * originally to generate the patch.
33  *
34  * To work well with this tool, the gzipped sections of the target
35  * image must have been generated using the same deflate encoder that
36  * is available in applypatch, namely, the one in the zlib library.
37  * In practice this means that images should be compressed using the
38  * "minigzip" tool included in the zlib distribution, not the GNU gzip
39  * program.
40  *
41  * An "imgdiff" patch consists of a header describing the chunk structure
42  * of the file and any encoding parameters needed for the gzipped
43  * chunks, followed by N bsdiff patches, one per chunk.
44  *
45  * For a diff to be generated, the source and target images must have the
46  * same "chunk" structure: that is, the same number of gzipped and normal
47  * chunks in the same order.  Android boot and recovery images currently
48  * consist of five chunks:  a small normal header, a gzipped kernel, a
49  * small normal section, a gzipped ramdisk, and finally a small normal
50  * footer.
51  *
52  * Caveats:  we locate gzipped sections within the source and target
53  * images by searching for the byte sequence 1f8b0800:  1f8b is the gzip
54  * magic number; 08 specifies the "deflate" encoding [the only encoding
55  * supported by the gzip standard]; and 00 is the flags byte.  We do not
56  * currently support any extra header fields (which would be indicated by
57  * a nonzero flags byte).  We also don't handle the case when that byte
58  * sequence appears spuriously in the file.  (Note that it would have to
59  * occur spuriously within a normal chunk to be a problem.)
60  *
61  *
62  * The imgdiff patch header looks like this:
63  *
64  *    "IMGDIFF1"                  (8)   [magic number and version]
65  *    chunk count                 (4)
66  *    for each chunk:
67  *        chunk type              (4)   [CHUNK_{NORMAL, GZIP, DEFLATE, RAW}]
68  *        if chunk type == CHUNK_NORMAL:
69  *           source start         (8)
70  *           source len           (8)
71  *           bsdiff patch offset  (8)   [from start of patch file]
72  *        if chunk type == CHUNK_GZIP:      (version 1 only)
73  *           source start         (8)
74  *           source len           (8)
75  *           bsdiff patch offset  (8)   [from start of patch file]
76  *           source expanded len  (8)   [size of uncompressed source]
77  *           target expected len  (8)   [size of uncompressed target]
78  *           gzip level           (4)
79  *                method          (4)
80  *                windowBits      (4)
81  *                memLevel        (4)
82  *                strategy        (4)
83  *           gzip header len      (4)
84  *           gzip header          (gzip header len)
85  *           gzip footer          (8)
86  *        if chunk type == CHUNK_DEFLATE:   (version 2 only)
87  *           source start         (8)
88  *           source len           (8)
89  *           bsdiff patch offset  (8)   [from start of patch file]
90  *           source expanded len  (8)   [size of uncompressed source]
91  *           target expected len  (8)   [size of uncompressed target]
92  *           gzip level           (4)
93  *                method          (4)
94  *                windowBits      (4)
95  *                memLevel        (4)
96  *                strategy        (4)
97  *        if chunk type == RAW:             (version 2 only)
98  *           target len           (4)
99  *           data                 (target len)
100  *
101  * All integers are little-endian.  "source start" and "source len"
102  * specify the section of the input image that comprises this chunk,
103  * including the gzip header and footer for gzip chunks.  "source
104  * expanded len" is the size of the uncompressed source data.  "target
105  * expected len" is the size of the uncompressed data after applying
106  * the bsdiff patch.  The next five parameters specify the zlib
107  * parameters to be used when compressing the patched data, and the
108  * next three specify the header and footer to be wrapped around the
109  * compressed data to create the output chunk (so that header contents
110  * like the timestamp are recreated exactly).
111  *
112  * After the header there are 'chunk count' bsdiff patches; the offset
113  * of each from the beginning of the file is specified in the header.
114  *
115  * This tool can take an optional file of "bonus data".  This is an
116  * extra file of data that is appended to chunk #1 after it is
117  * compressed (it must be a CHUNK_DEFLATE chunk).  The same file must
118  * be available (and passed to applypatch with -b) when applying the
119  * patch.  This is used to reduce the size of recovery-from-boot
120  * patches by combining the boot image with recovery ramdisk
121  * information that is stored on the system partition.
122  */
123 
124 #include <errno.h>
125 #include <inttypes.h>
126 #include <stdio.h>
127 #include <stdlib.h>
128 #include <string.h>
129 #include <sys/stat.h>
130 #include <unistd.h>
131 #include <sys/types.h>
132 
133 #include "zlib.h"
134 #include "imgdiff.h"
135 #include "utils.h"
136 
137 typedef struct {
138   int type;             // CHUNK_NORMAL, CHUNK_DEFLATE
139   size_t start;         // offset of chunk in original image file
140 
141   size_t len;
142   unsigned char* data;  // data to be patched (uncompressed, for deflate chunks)
143 
144   size_t source_start;
145   size_t source_len;
146 
147   off_t* I;             // used by bsdiff
148 
149   // --- for CHUNK_DEFLATE chunks only: ---
150 
151   // original (compressed) deflate data
152   size_t deflate_len;
153   unsigned char* deflate_data;
154 
155   char* filename;       // used for zip entries
156 
157   // deflate encoder parameters
158   int level, method, windowBits, memLevel, strategy;
159 
160   size_t source_uncompressed_len;
161 } ImageChunk;
162 
163 typedef struct {
164   int data_offset;
165   int deflate_len;
166   int uncomp_len;
167   char* filename;
168 } ZipFileEntry;
169 
fileentry_compare(const void * a,const void * b)170 static int fileentry_compare(const void* a, const void* b) {
171   int ao = ((ZipFileEntry*)a)->data_offset;
172   int bo = ((ZipFileEntry*)b)->data_offset;
173   if (ao < bo) {
174     return -1;
175   } else if (ao > bo) {
176     return 1;
177   } else {
178     return 0;
179   }
180 }
181 
182 // from bsdiff.c
183 int bsdiff(u_char* old, off_t oldsize, off_t** IP, u_char* newdata, off_t newsize,
184            const char* patch_filename);
185 
ReadZip(const char * filename,int * num_chunks,ImageChunk ** chunks,int include_pseudo_chunk)186 unsigned char* ReadZip(const char* filename,
187                        int* num_chunks, ImageChunk** chunks,
188                        int include_pseudo_chunk) {
189   struct stat st;
190   if (stat(filename, &st) != 0) {
191     printf("failed to stat \"%s\": %s\n", filename, strerror(errno));
192     return NULL;
193   }
194 
195   size_t sz = static_cast<size_t>(st.st_size);
196   unsigned char* img = reinterpret_cast<unsigned char*>(malloc(sz));
197   FILE* f = fopen(filename, "rb");
198   if (fread(img, 1, sz, f) != sz) {
199     printf("failed to read \"%s\" %s\n", filename, strerror(errno));
200     fclose(f);
201     return NULL;
202   }
203   fclose(f);
204 
205   // look for the end-of-central-directory record.
206 
207   int i;
208   for (i = st.st_size-20; i >= 0 && i > st.st_size - 65600; --i) {
209     if (img[i] == 0x50 && img[i+1] == 0x4b &&
210         img[i+2] == 0x05 && img[i+3] == 0x06) {
211       break;
212     }
213   }
214   // double-check: this archive consists of a single "disk"
215   if (!(img[i+4] == 0 && img[i+5] == 0 && img[i+6] == 0 && img[i+7] == 0)) {
216     printf("can't process multi-disk archive\n");
217     return NULL;
218   }
219 
220   int cdcount = Read2(img+i+8);
221   int cdoffset = Read4(img+i+16);
222 
223   ZipFileEntry* temp_entries = reinterpret_cast<ZipFileEntry*>(malloc(
224       cdcount * sizeof(ZipFileEntry)));
225   int entrycount = 0;
226 
227   unsigned char* cd = img+cdoffset;
228   for (i = 0; i < cdcount; ++i) {
229     if (!(cd[0] == 0x50 && cd[1] == 0x4b && cd[2] == 0x01 && cd[3] == 0x02)) {
230       printf("bad central directory entry %d\n", i);
231       return NULL;
232     }
233 
234     int clen = Read4(cd+20);   // compressed len
235     int ulen = Read4(cd+24);   // uncompressed len
236     int nlen = Read2(cd+28);   // filename len
237     int xlen = Read2(cd+30);   // extra field len
238     int mlen = Read2(cd+32);   // file comment len
239     int hoffset = Read4(cd+42);   // local header offset
240 
241     char* filename = reinterpret_cast<char*>(malloc(nlen+1));
242     memcpy(filename, cd+46, nlen);
243     filename[nlen] = '\0';
244 
245     int method = Read2(cd+10);
246 
247     cd += 46 + nlen + xlen + mlen;
248 
249     if (method != 8) {  // 8 == deflate
250       free(filename);
251       continue;
252     }
253 
254     unsigned char* lh = img + hoffset;
255 
256     if (!(lh[0] == 0x50 && lh[1] == 0x4b && lh[2] == 0x03 && lh[3] == 0x04)) {
257       printf("bad local file header entry %d\n", i);
258       return NULL;
259     }
260 
261     if (Read2(lh+26) != nlen || memcmp(lh+30, filename, nlen) != 0) {
262       printf("central dir filename doesn't match local header\n");
263       return NULL;
264     }
265 
266     xlen = Read2(lh+28);   // extra field len; might be different from CD entry?
267 
268     temp_entries[entrycount].data_offset = hoffset+30+nlen+xlen;
269     temp_entries[entrycount].deflate_len = clen;
270     temp_entries[entrycount].uncomp_len = ulen;
271     temp_entries[entrycount].filename = filename;
272     ++entrycount;
273   }
274 
275   qsort(temp_entries, entrycount, sizeof(ZipFileEntry), fileentry_compare);
276 
277 #if 0
278   printf("found %d deflated entries\n", entrycount);
279   for (i = 0; i < entrycount; ++i) {
280     printf("off %10d  len %10d unlen %10d   %p %s\n",
281            temp_entries[i].data_offset,
282            temp_entries[i].deflate_len,
283            temp_entries[i].uncomp_len,
284            temp_entries[i].filename,
285            temp_entries[i].filename);
286   }
287 #endif
288 
289   *num_chunks = 0;
290   *chunks = reinterpret_cast<ImageChunk*>(malloc((entrycount*2+2) * sizeof(ImageChunk)));
291   ImageChunk* curr = *chunks;
292 
293   if (include_pseudo_chunk) {
294     curr->type = CHUNK_NORMAL;
295     curr->start = 0;
296     curr->len = st.st_size;
297     curr->data = img;
298     curr->filename = NULL;
299     curr->I = NULL;
300     ++curr;
301     ++*num_chunks;
302   }
303 
304   int pos = 0;
305   int nextentry = 0;
306 
307   while (pos < st.st_size) {
308     if (nextentry < entrycount && pos == temp_entries[nextentry].data_offset) {
309       curr->type = CHUNK_DEFLATE;
310       curr->start = pos;
311       curr->deflate_len = temp_entries[nextentry].deflate_len;
312       curr->deflate_data = img + pos;
313       curr->filename = temp_entries[nextentry].filename;
314       curr->I = NULL;
315 
316       curr->len = temp_entries[nextentry].uncomp_len;
317       curr->data = reinterpret_cast<unsigned char*>(malloc(curr->len));
318 
319       z_stream strm;
320       strm.zalloc = Z_NULL;
321       strm.zfree = Z_NULL;
322       strm.opaque = Z_NULL;
323       strm.avail_in = curr->deflate_len;
324       strm.next_in = curr->deflate_data;
325 
326       // -15 means we are decoding a 'raw' deflate stream; zlib will
327       // not expect zlib headers.
328       int ret = inflateInit2(&strm, -15);
329 
330       strm.avail_out = curr->len;
331       strm.next_out = curr->data;
332       ret = inflate(&strm, Z_NO_FLUSH);
333       if (ret != Z_STREAM_END) {
334         printf("failed to inflate \"%s\"; %d\n", curr->filename, ret);
335         return NULL;
336       }
337 
338       inflateEnd(&strm);
339 
340       pos += curr->deflate_len;
341       ++nextentry;
342       ++*num_chunks;
343       ++curr;
344       continue;
345     }
346 
347     // use a normal chunk to take all the data up to the start of the
348     // next deflate section.
349 
350     curr->type = CHUNK_NORMAL;
351     curr->start = pos;
352     if (nextentry < entrycount) {
353       curr->len = temp_entries[nextentry].data_offset - pos;
354     } else {
355       curr->len = st.st_size - pos;
356     }
357     curr->data = img + pos;
358     curr->filename = NULL;
359     curr->I = NULL;
360     pos += curr->len;
361 
362     ++*num_chunks;
363     ++curr;
364   }
365 
366   free(temp_entries);
367   return img;
368 }
369 
370 /*
371  * Read the given file and break it up into chunks, putting the number
372  * of chunks and their info in *num_chunks and **chunks,
373  * respectively.  Returns a malloc'd block of memory containing the
374  * contents of the file; various pointers in the output chunk array
375  * will point into this block of memory.  The caller should free the
376  * return value when done with all the chunks.  Returns NULL on
377  * failure.
378  */
ReadImage(const char * filename,int * num_chunks,ImageChunk ** chunks)379 unsigned char* ReadImage(const char* filename,
380                          int* num_chunks, ImageChunk** chunks) {
381   struct stat st;
382   if (stat(filename, &st) != 0) {
383     printf("failed to stat \"%s\": %s\n", filename, strerror(errno));
384     return NULL;
385   }
386 
387   size_t sz = static_cast<size_t>(st.st_size);
388   unsigned char* img = reinterpret_cast<unsigned char*>(malloc(sz + 4));
389   FILE* f = fopen(filename, "rb");
390   if (fread(img, 1, sz, f) != sz) {
391     printf("failed to read \"%s\" %s\n", filename, strerror(errno));
392     fclose(f);
393     return NULL;
394   }
395   fclose(f);
396 
397   // append 4 zero bytes to the data so we can always search for the
398   // four-byte string 1f8b0800 starting at any point in the actual
399   // file data, without special-casing the end of the data.
400   memset(img+sz, 0, 4);
401 
402   size_t pos = 0;
403 
404   *num_chunks = 0;
405   *chunks = NULL;
406 
407   while (pos < sz) {
408     unsigned char* p = img+pos;
409 
410     bool processed_deflate = false;
411     if (sz - pos >= 4 &&
412         p[0] == 0x1f && p[1] == 0x8b &&
413         p[2] == 0x08 &&    // deflate compression
414         p[3] == 0x00) {    // no header flags
415       // 'pos' is the offset of the start of a gzip chunk.
416       size_t chunk_offset = pos;
417 
418       *num_chunks += 3;
419       *chunks = reinterpret_cast<ImageChunk*>(realloc(*chunks,
420           *num_chunks * sizeof(ImageChunk)));
421       ImageChunk* curr = *chunks + (*num_chunks-3);
422 
423       // create a normal chunk for the header.
424       curr->start = pos;
425       curr->type = CHUNK_NORMAL;
426       curr->len = GZIP_HEADER_LEN;
427       curr->data = p;
428       curr->I = NULL;
429 
430       pos += curr->len;
431       p += curr->len;
432       ++curr;
433 
434       curr->type = CHUNK_DEFLATE;
435       curr->filename = NULL;
436       curr->I = NULL;
437 
438       // We must decompress this chunk in order to discover where it
439       // ends, and so we can put the uncompressed data and its length
440       // into curr->data and curr->len.
441 
442       size_t allocated = 32768;
443       curr->len = 0;
444       curr->data = reinterpret_cast<unsigned char*>(malloc(allocated));
445       curr->start = pos;
446       curr->deflate_data = p;
447 
448       z_stream strm;
449       strm.zalloc = Z_NULL;
450       strm.zfree = Z_NULL;
451       strm.opaque = Z_NULL;
452       strm.avail_in = sz - pos;
453       strm.next_in = p;
454 
455       // -15 means we are decoding a 'raw' deflate stream; zlib will
456       // not expect zlib headers.
457       int ret = inflateInit2(&strm, -15);
458 
459       do {
460         strm.avail_out = allocated - curr->len;
461         strm.next_out = curr->data + curr->len;
462         ret = inflate(&strm, Z_NO_FLUSH);
463         if (ret < 0) {
464           if (!processed_deflate) {
465             // This is the first chunk, assume that it's just a spurious
466             // gzip header instead of a real one.
467             break;
468           }
469           printf("Error: inflate failed [%s] at file offset [%zu]\n"
470                  "imgdiff only supports gzip kernel compression,"
471                  " did you try CONFIG_KERNEL_LZO?\n",
472                  strm.msg, chunk_offset);
473           free(img);
474           return NULL;
475         }
476         curr->len = allocated - strm.avail_out;
477         if (strm.avail_out == 0) {
478           allocated *= 2;
479           curr->data = reinterpret_cast<unsigned char*>(realloc(curr->data, allocated));
480         }
481         processed_deflate = true;
482       } while (ret != Z_STREAM_END);
483 
484       curr->deflate_len = sz - strm.avail_in - pos;
485       inflateEnd(&strm);
486       pos += curr->deflate_len;
487       p += curr->deflate_len;
488       ++curr;
489 
490       // create a normal chunk for the footer
491 
492       curr->type = CHUNK_NORMAL;
493       curr->start = pos;
494       curr->len = GZIP_FOOTER_LEN;
495       curr->data = img+pos;
496       curr->I = NULL;
497 
498       pos += curr->len;
499       p += curr->len;
500       ++curr;
501 
502       // The footer (that we just skipped over) contains the size of
503       // the uncompressed data.  Double-check to make sure that it
504       // matches the size of the data we got when we actually did
505       // the decompression.
506       size_t footer_size = Read4(p-4);
507       if (footer_size != curr[-2].len) {
508         printf("Error: footer size %zu != decompressed size %zu\n",
509             footer_size, curr[-2].len);
510         free(img);
511         return NULL;
512       }
513     } else {
514       // Reallocate the list for every chunk; we expect the number of
515       // chunks to be small (5 for typical boot and recovery images).
516       ++*num_chunks;
517       *chunks = reinterpret_cast<ImageChunk*>(realloc(*chunks, *num_chunks * sizeof(ImageChunk)));
518       ImageChunk* curr = *chunks + (*num_chunks-1);
519       curr->start = pos;
520       curr->I = NULL;
521 
522       // 'pos' is not the offset of the start of a gzip chunk, so scan
523       // forward until we find a gzip header.
524       curr->type = CHUNK_NORMAL;
525       curr->data = p;
526 
527       for (curr->len = 0; curr->len < (sz - pos); ++curr->len) {
528         if (p[curr->len] == 0x1f &&
529             p[curr->len+1] == 0x8b &&
530             p[curr->len+2] == 0x08 &&
531             p[curr->len+3] == 0x00) {
532           break;
533         }
534       }
535       pos += curr->len;
536     }
537   }
538 
539   return img;
540 }
541 
542 #define BUFFER_SIZE 32768
543 
544 /*
545  * Takes the uncompressed data stored in the chunk, compresses it
546  * using the zlib parameters stored in the chunk, and checks that it
547  * matches exactly the compressed data we started with (also stored in
548  * the chunk).  Return 0 on success.
549  */
TryReconstruction(ImageChunk * chunk,unsigned char * out)550 int TryReconstruction(ImageChunk* chunk, unsigned char* out) {
551   size_t p = 0;
552 
553 #if 0
554   printf("trying %d %d %d %d %d\n",
555           chunk->level, chunk->method, chunk->windowBits,
556           chunk->memLevel, chunk->strategy);
557 #endif
558 
559   z_stream strm;
560   strm.zalloc = Z_NULL;
561   strm.zfree = Z_NULL;
562   strm.opaque = Z_NULL;
563   strm.avail_in = chunk->len;
564   strm.next_in = chunk->data;
565   int ret;
566   ret = deflateInit2(&strm, chunk->level, chunk->method, chunk->windowBits,
567                      chunk->memLevel, chunk->strategy);
568   do {
569     strm.avail_out = BUFFER_SIZE;
570     strm.next_out = out;
571     ret = deflate(&strm, Z_FINISH);
572     size_t have = BUFFER_SIZE - strm.avail_out;
573 
574     if (memcmp(out, chunk->deflate_data+p, have) != 0) {
575       // mismatch; data isn't the same.
576       deflateEnd(&strm);
577       return -1;
578     }
579     p += have;
580   } while (ret != Z_STREAM_END);
581   deflateEnd(&strm);
582   if (p != chunk->deflate_len) {
583     // mismatch; ran out of data before we should have.
584     return -1;
585   }
586   return 0;
587 }
588 
589 /*
590  * Verify that we can reproduce exactly the same compressed data that
591  * we started with.  Sets the level, method, windowBits, memLevel, and
592  * strategy fields in the chunk to the encoding parameters needed to
593  * produce the right output.  Returns 0 on success.
594  */
ReconstructDeflateChunk(ImageChunk * chunk)595 int ReconstructDeflateChunk(ImageChunk* chunk) {
596   if (chunk->type != CHUNK_DEFLATE) {
597     printf("attempt to reconstruct non-deflate chunk\n");
598     return -1;
599   }
600 
601   size_t p = 0;
602   unsigned char* out = reinterpret_cast<unsigned char*>(malloc(BUFFER_SIZE));
603 
604   // We only check two combinations of encoder parameters:  level 6
605   // (the default) and level 9 (the maximum).
606   for (chunk->level = 6; chunk->level <= 9; chunk->level += 3) {
607     chunk->windowBits = -15;  // 32kb window; negative to indicate a raw stream.
608     chunk->memLevel = 8;      // the default value.
609     chunk->method = Z_DEFLATED;
610     chunk->strategy = Z_DEFAULT_STRATEGY;
611 
612     if (TryReconstruction(chunk, out) == 0) {
613       free(out);
614       return 0;
615     }
616   }
617 
618   free(out);
619   return -1;
620 }
621 
622 /*
623  * Given source and target chunks, compute a bsdiff patch between them
624  * by running bsdiff in a subprocess.  Return the patch data, placing
625  * its length in *size.  Return NULL on failure.  We expect the bsdiff
626  * program to be in the path.
627  */
MakePatch(ImageChunk * src,ImageChunk * tgt,size_t * size)628 unsigned char* MakePatch(ImageChunk* src, ImageChunk* tgt, size_t* size) {
629   if (tgt->type == CHUNK_NORMAL) {
630     if (tgt->len <= 160) {
631       tgt->type = CHUNK_RAW;
632       *size = tgt->len;
633       return tgt->data;
634     }
635   }
636 
637   char ptemp[] = "/tmp/imgdiff-patch-XXXXXX";
638   int fd = mkstemp(ptemp);
639 
640   if (fd == -1) {
641     printf("MakePatch failed to create a temporary file: %s\n",
642            strerror(errno));
643     return NULL;
644   }
645   close(fd); // temporary file is created and we don't need its file
646              // descriptor
647 
648   int r = bsdiff(src->data, src->len, &(src->I), tgt->data, tgt->len, ptemp);
649   if (r != 0) {
650     printf("bsdiff() failed: %d\n", r);
651     return NULL;
652   }
653 
654   struct stat st;
655   if (stat(ptemp, &st) != 0) {
656     printf("failed to stat patch file %s: %s\n",
657             ptemp, strerror(errno));
658     return NULL;
659   }
660 
661   size_t sz = static_cast<size_t>(st.st_size);
662   // TODO: Memory leak on error return.
663   unsigned char* data = reinterpret_cast<unsigned char*>(malloc(sz));
664 
665   if (tgt->type == CHUNK_NORMAL && tgt->len <= sz) {
666     unlink(ptemp);
667 
668     tgt->type = CHUNK_RAW;
669     *size = tgt->len;
670     return tgt->data;
671   }
672 
673   *size = sz;
674 
675   FILE* f = fopen(ptemp, "rb");
676   if (f == NULL) {
677     printf("failed to open patch %s: %s\n", ptemp, strerror(errno));
678     return NULL;
679   }
680   if (fread(data, 1, sz, f) != sz) {
681     printf("failed to read patch %s: %s\n", ptemp, strerror(errno));
682     return NULL;
683   }
684   fclose(f);
685 
686   unlink(ptemp);
687 
688   tgt->source_start = src->start;
689   switch (tgt->type) {
690     case CHUNK_NORMAL:
691       tgt->source_len = src->len;
692       break;
693     case CHUNK_DEFLATE:
694       tgt->source_len = src->deflate_len;
695       tgt->source_uncompressed_len = src->len;
696       break;
697   }
698 
699   return data;
700 }
701 
702 /*
703  * Cause a gzip chunk to be treated as a normal chunk (ie, as a blob
704  * of uninterpreted data).  The resulting patch will likely be about
705  * as big as the target file, but it lets us handle the case of images
706  * where some gzip chunks are reconstructible but others aren't (by
707  * treating the ones that aren't as normal chunks).
708  */
ChangeDeflateChunkToNormal(ImageChunk * ch)709 void ChangeDeflateChunkToNormal(ImageChunk* ch) {
710   if (ch->type != CHUNK_DEFLATE) return;
711   ch->type = CHUNK_NORMAL;
712   free(ch->data);
713   ch->data = ch->deflate_data;
714   ch->len = ch->deflate_len;
715 }
716 
717 /*
718  * Return true if the data in the chunk is identical (including the
719  * compressed representation, for gzip chunks).
720  */
AreChunksEqual(ImageChunk * a,ImageChunk * b)721 int AreChunksEqual(ImageChunk* a, ImageChunk* b) {
722     if (a->type != b->type) return 0;
723 
724     switch (a->type) {
725         case CHUNK_NORMAL:
726             return a->len == b->len && memcmp(a->data, b->data, a->len) == 0;
727 
728         case CHUNK_DEFLATE:
729             return a->deflate_len == b->deflate_len &&
730                 memcmp(a->deflate_data, b->deflate_data, a->deflate_len) == 0;
731 
732         default:
733             printf("unknown chunk type %d\n", a->type);
734             return 0;
735     }
736 }
737 
738 /*
739  * Look for runs of adjacent normal chunks and compress them down into
740  * a single chunk.  (Such runs can be produced when deflate chunks are
741  * changed to normal chunks.)
742  */
MergeAdjacentNormalChunks(ImageChunk * chunks,int * num_chunks)743 void MergeAdjacentNormalChunks(ImageChunk* chunks, int* num_chunks) {
744   int out = 0;
745   int in_start = 0, in_end;
746   while (in_start < *num_chunks) {
747     if (chunks[in_start].type != CHUNK_NORMAL) {
748       in_end = in_start+1;
749     } else {
750       // in_start is a normal chunk.  Look for a run of normal chunks
751       // that constitute a solid block of data (ie, each chunk begins
752       // where the previous one ended).
753       for (in_end = in_start+1;
754            in_end < *num_chunks && chunks[in_end].type == CHUNK_NORMAL &&
755              (chunks[in_end].start ==
756               chunks[in_end-1].start + chunks[in_end-1].len &&
757               chunks[in_end].data ==
758               chunks[in_end-1].data + chunks[in_end-1].len);
759            ++in_end);
760     }
761 
762     if (in_end == in_start+1) {
763 #if 0
764       printf("chunk %d is now %d\n", in_start, out);
765 #endif
766       if (out != in_start) {
767         memcpy(chunks+out, chunks+in_start, sizeof(ImageChunk));
768       }
769     } else {
770 #if 0
771       printf("collapse normal chunks %d-%d into %d\n", in_start, in_end-1, out);
772 #endif
773 
774       // Merge chunks [in_start, in_end-1] into one chunk.  Since the
775       // data member of each chunk is just a pointer into an in-memory
776       // copy of the file, this can be done without recopying (the
777       // output chunk has the first chunk's start location and data
778       // pointer, and length equal to the sum of the input chunk
779       // lengths).
780       chunks[out].type = CHUNK_NORMAL;
781       chunks[out].start = chunks[in_start].start;
782       chunks[out].data = chunks[in_start].data;
783       chunks[out].len = chunks[in_end-1].len +
784         (chunks[in_end-1].start - chunks[in_start].start);
785     }
786 
787     ++out;
788     in_start = in_end;
789   }
790   *num_chunks = out;
791 }
792 
FindChunkByName(const char * name,ImageChunk * chunks,int num_chunks)793 ImageChunk* FindChunkByName(const char* name,
794                             ImageChunk* chunks, int num_chunks) {
795   int i;
796   for (i = 0; i < num_chunks; ++i) {
797     if (chunks[i].type == CHUNK_DEFLATE && chunks[i].filename &&
798         strcmp(name, chunks[i].filename) == 0) {
799       return chunks+i;
800     }
801   }
802   return NULL;
803 }
804 
DumpChunks(ImageChunk * chunks,int num_chunks)805 void DumpChunks(ImageChunk* chunks, int num_chunks) {
806     for (int i = 0; i < num_chunks; ++i) {
807         printf("chunk %d: type %d start %zu len %zu\n",
808                i, chunks[i].type, chunks[i].start, chunks[i].len);
809     }
810 }
811 
main(int argc,char ** argv)812 int main(int argc, char** argv) {
813   int zip_mode = 0;
814 
815   if (argc >= 2 && strcmp(argv[1], "-z") == 0) {
816     zip_mode = 1;
817     --argc;
818     ++argv;
819   }
820 
821   size_t bonus_size = 0;
822   unsigned char* bonus_data = NULL;
823   if (argc >= 3 && strcmp(argv[1], "-b") == 0) {
824     struct stat st;
825     if (stat(argv[2], &st) != 0) {
826       printf("failed to stat bonus file %s: %s\n", argv[2], strerror(errno));
827       return 1;
828     }
829     bonus_size = st.st_size;
830     bonus_data = reinterpret_cast<unsigned char*>(malloc(bonus_size));
831     FILE* f = fopen(argv[2], "rb");
832     if (f == NULL) {
833       printf("failed to open bonus file %s: %s\n", argv[2], strerror(errno));
834       return 1;
835     }
836     if (fread(bonus_data, 1, bonus_size, f) != bonus_size) {
837       printf("failed to read bonus file %s: %s\n", argv[2], strerror(errno));
838       return 1;
839     }
840     fclose(f);
841 
842     argc -= 2;
843     argv += 2;
844   }
845 
846   if (argc != 4) {
847     usage:
848     printf("usage: %s [-z] [-b <bonus-file>] <src-img> <tgt-img> <patch-file>\n",
849             argv[0]);
850     return 2;
851   }
852 
853   int num_src_chunks;
854   ImageChunk* src_chunks;
855   int num_tgt_chunks;
856   ImageChunk* tgt_chunks;
857   int i;
858 
859   if (zip_mode) {
860     if (ReadZip(argv[1], &num_src_chunks, &src_chunks, 1) == NULL) {
861       printf("failed to break apart source zip file\n");
862       return 1;
863     }
864     if (ReadZip(argv[2], &num_tgt_chunks, &tgt_chunks, 0) == NULL) {
865       printf("failed to break apart target zip file\n");
866       return 1;
867     }
868   } else {
869     if (ReadImage(argv[1], &num_src_chunks, &src_chunks) == NULL) {
870       printf("failed to break apart source image\n");
871       return 1;
872     }
873     if (ReadImage(argv[2], &num_tgt_chunks, &tgt_chunks) == NULL) {
874       printf("failed to break apart target image\n");
875       return 1;
876     }
877 
878     // Verify that the source and target images have the same chunk
879     // structure (ie, the same sequence of deflate and normal chunks).
880 
881     if (!zip_mode) {
882         // Merge the gzip header and footer in with any adjacent
883         // normal chunks.
884         MergeAdjacentNormalChunks(tgt_chunks, &num_tgt_chunks);
885         MergeAdjacentNormalChunks(src_chunks, &num_src_chunks);
886     }
887 
888     if (num_src_chunks != num_tgt_chunks) {
889       printf("source and target don't have same number of chunks!\n");
890       printf("source chunks:\n");
891       DumpChunks(src_chunks, num_src_chunks);
892       printf("target chunks:\n");
893       DumpChunks(tgt_chunks, num_tgt_chunks);
894       return 1;
895     }
896     for (i = 0; i < num_src_chunks; ++i) {
897       if (src_chunks[i].type != tgt_chunks[i].type) {
898         printf("source and target don't have same chunk "
899                 "structure! (chunk %d)\n", i);
900         printf("source chunks:\n");
901         DumpChunks(src_chunks, num_src_chunks);
902         printf("target chunks:\n");
903         DumpChunks(tgt_chunks, num_tgt_chunks);
904         return 1;
905       }
906     }
907   }
908 
909   for (i = 0; i < num_tgt_chunks; ++i) {
910     if (tgt_chunks[i].type == CHUNK_DEFLATE) {
911       // Confirm that given the uncompressed chunk data in the target, we
912       // can recompress it and get exactly the same bits as are in the
913       // input target image.  If this fails, treat the chunk as a normal
914       // non-deflated chunk.
915       if (ReconstructDeflateChunk(tgt_chunks+i) < 0) {
916         printf("failed to reconstruct target deflate chunk %d [%s]; "
917                "treating as normal\n", i, tgt_chunks[i].filename);
918         ChangeDeflateChunkToNormal(tgt_chunks+i);
919         if (zip_mode) {
920           ImageChunk* src = FindChunkByName(tgt_chunks[i].filename, src_chunks, num_src_chunks);
921           if (src) {
922             ChangeDeflateChunkToNormal(src);
923           }
924         } else {
925           ChangeDeflateChunkToNormal(src_chunks+i);
926         }
927         continue;
928       }
929 
930       // If two deflate chunks are identical (eg, the kernel has not
931       // changed between two builds), treat them as normal chunks.
932       // This makes applypatch much faster -- it can apply a trivial
933       // patch to the compressed data, rather than uncompressing and
934       // recompressing to apply the trivial patch to the uncompressed
935       // data.
936       ImageChunk* src;
937       if (zip_mode) {
938         src = FindChunkByName(tgt_chunks[i].filename, src_chunks, num_src_chunks);
939       } else {
940         src = src_chunks+i;
941       }
942 
943       if (src == NULL || AreChunksEqual(tgt_chunks+i, src)) {
944         ChangeDeflateChunkToNormal(tgt_chunks+i);
945         if (src) {
946           ChangeDeflateChunkToNormal(src);
947         }
948       }
949     }
950   }
951 
952   // Merging neighboring normal chunks.
953   if (zip_mode) {
954     // For zips, we only need to do this to the target:  deflated
955     // chunks are matched via filename, and normal chunks are patched
956     // using the entire source file as the source.
957     MergeAdjacentNormalChunks(tgt_chunks, &num_tgt_chunks);
958   } else {
959     // For images, we need to maintain the parallel structure of the
960     // chunk lists, so do the merging in both the source and target
961     // lists.
962     MergeAdjacentNormalChunks(tgt_chunks, &num_tgt_chunks);
963     MergeAdjacentNormalChunks(src_chunks, &num_src_chunks);
964     if (num_src_chunks != num_tgt_chunks) {
965       // This shouldn't happen.
966       printf("merging normal chunks went awry\n");
967       return 1;
968     }
969   }
970 
971   // Compute bsdiff patches for each chunk's data (the uncompressed
972   // data, in the case of deflate chunks).
973 
974   DumpChunks(src_chunks, num_src_chunks);
975 
976   printf("Construct patches for %d chunks...\n", num_tgt_chunks);
977   unsigned char** patch_data = reinterpret_cast<unsigned char**>(malloc(
978       num_tgt_chunks * sizeof(unsigned char*)));
979   size_t* patch_size = reinterpret_cast<size_t*>(malloc(num_tgt_chunks * sizeof(size_t)));
980   for (i = 0; i < num_tgt_chunks; ++i) {
981     if (zip_mode) {
982       ImageChunk* src;
983       if (tgt_chunks[i].type == CHUNK_DEFLATE &&
984           (src = FindChunkByName(tgt_chunks[i].filename, src_chunks,
985                                  num_src_chunks))) {
986         patch_data[i] = MakePatch(src, tgt_chunks+i, patch_size+i);
987       } else {
988         patch_data[i] = MakePatch(src_chunks, tgt_chunks+i, patch_size+i);
989       }
990     } else {
991       if (i == 1 && bonus_data) {
992         printf("  using %zu bytes of bonus data for chunk %d\n", bonus_size, i);
993         src_chunks[i].data = reinterpret_cast<unsigned char*>(realloc(src_chunks[i].data,
994             src_chunks[i].len + bonus_size));
995         memcpy(src_chunks[i].data+src_chunks[i].len, bonus_data, bonus_size);
996         src_chunks[i].len += bonus_size;
997      }
998 
999       patch_data[i] = MakePatch(src_chunks+i, tgt_chunks+i, patch_size+i);
1000     }
1001     printf("patch %3d is %zu bytes (of %zu)\n",
1002            i, patch_size[i], tgt_chunks[i].source_len);
1003   }
1004 
1005   // Figure out how big the imgdiff file header is going to be, so
1006   // that we can correctly compute the offset of each bsdiff patch
1007   // within the file.
1008 
1009   size_t total_header_size = 12;
1010   for (i = 0; i < num_tgt_chunks; ++i) {
1011     total_header_size += 4;
1012     switch (tgt_chunks[i].type) {
1013       case CHUNK_NORMAL:
1014         total_header_size += 8*3;
1015         break;
1016       case CHUNK_DEFLATE:
1017         total_header_size += 8*5 + 4*5;
1018         break;
1019       case CHUNK_RAW:
1020         total_header_size += 4 + patch_size[i];
1021         break;
1022     }
1023   }
1024 
1025   size_t offset = total_header_size;
1026 
1027   FILE* f = fopen(argv[3], "wb");
1028 
1029   // Write out the headers.
1030 
1031   fwrite("IMGDIFF2", 1, 8, f);
1032   Write4(num_tgt_chunks, f);
1033   for (i = 0; i < num_tgt_chunks; ++i) {
1034     Write4(tgt_chunks[i].type, f);
1035 
1036     switch (tgt_chunks[i].type) {
1037       case CHUNK_NORMAL:
1038         printf("chunk %3d: normal   (%10zu, %10zu)  %10zu\n", i,
1039                tgt_chunks[i].start, tgt_chunks[i].len, patch_size[i]);
1040         Write8(tgt_chunks[i].source_start, f);
1041         Write8(tgt_chunks[i].source_len, f);
1042         Write8(offset, f);
1043         offset += patch_size[i];
1044         break;
1045 
1046       case CHUNK_DEFLATE:
1047         printf("chunk %3d: deflate  (%10zu, %10zu)  %10zu  %s\n", i,
1048                tgt_chunks[i].start, tgt_chunks[i].deflate_len, patch_size[i],
1049                tgt_chunks[i].filename);
1050         Write8(tgt_chunks[i].source_start, f);
1051         Write8(tgt_chunks[i].source_len, f);
1052         Write8(offset, f);
1053         Write8(tgt_chunks[i].source_uncompressed_len, f);
1054         Write8(tgt_chunks[i].len, f);
1055         Write4(tgt_chunks[i].level, f);
1056         Write4(tgt_chunks[i].method, f);
1057         Write4(tgt_chunks[i].windowBits, f);
1058         Write4(tgt_chunks[i].memLevel, f);
1059         Write4(tgt_chunks[i].strategy, f);
1060         offset += patch_size[i];
1061         break;
1062 
1063       case CHUNK_RAW:
1064         printf("chunk %3d: raw      (%10zu, %10zu)\n", i,
1065                tgt_chunks[i].start, tgt_chunks[i].len);
1066         Write4(patch_size[i], f);
1067         fwrite(patch_data[i], 1, patch_size[i], f);
1068         break;
1069     }
1070   }
1071 
1072   // Append each chunk's bsdiff patch, in order.
1073 
1074   for (i = 0; i < num_tgt_chunks; ++i) {
1075     if (tgt_chunks[i].type != CHUNK_RAW) {
1076       fwrite(patch_data[i], 1, patch_size[i], f);
1077     }
1078   }
1079 
1080   fclose(f);
1081 
1082   return 0;
1083 }
1084