Fixed up base64_dump() users and interfaces to fix memory leaks
[jigit.git] / parse_template.c
1 #include <stdio.h>
2 #include <stdlib.h>
3 #include <errno.h>
4 #include <string.h>
5 #include <fcntl.h>
6 #include <sys/types.h>
7 #include <sys/stat.h>
8 #include <limits.h>
9 #include <zlib.h>
10 #ifdef BZ2_SUPPORT
11 #   include <bzlib.h>
12 #endif
13 #include "md5.h"
14 #include "endian.h"
15 #include "jigdb.h"
16 #include "jte.h"
17
18 struct
19 {
20     char   *data_buf;
21     ctype_e algorithm;
22     INT64   buf_size;
23     INT64   offset_in_curr_buf;
24     INT64   total_offset;
25 } zip_state;
26
27 /* Grab the file component from a full path */
28 static char *file_base_name(char *path)
29 {
30     char *endptr = path;
31     char *ptr = path;
32     
33     while (*ptr != '\0')
34     {
35         if ('/' == *ptr)
36             endptr = ++ptr;
37         else
38             ++ptr;
39     }
40     return endptr;
41 }
42
43 static int read_data_block(FILE *template_file)
44 {
45     char inbuf[1024];
46     INT64 i = 0;
47     static INT64 template_offset = -1;
48     INT64 compressed_len = 0;
49     INT64 uncompressed_len = 0;
50     char *comp_buf = NULL;
51     int read_num = 0;
52     int error = 0;
53
54     if (-1 == template_offset)
55     {
56         fseek(template_file, 0, SEEK_SET);
57         fread(inbuf, sizeof(inbuf), 1, template_file);
58         for (i = 0; i < sizeof(inbuf); i++)
59         {
60             if (!strncmp(&inbuf[i], "DATA", 4))
61             {
62                 zip_state.algorithm = CT_GZIP;
63                 template_offset = i;
64                 break;
65             }
66             if (!strncmp(&inbuf[i], "BZIP", 4))
67             {
68                 zip_state.algorithm = CT_BZIP2;
69                 template_offset = i;
70                 break;
71             }
72         }
73         if (-1 == template_offset)
74         {
75             fprintf(G_logfile, "Unable to locate DATA block in template (offset %lld)\n",
76                     template_offset);
77             return EINVAL;
78         }    
79     }
80     
81     fseek(template_file, template_offset, SEEK_SET);
82     fread(inbuf, 16, 1, template_file);
83     if (strncmp(inbuf, "DATA", 4) && strncmp(inbuf, "BZIP", 4))
84     {
85         fprintf(G_logfile, "Unable to locate DATA block in template (offset %lld)\n",
86                 template_offset);
87         return EINVAL;
88     }    
89     
90     compressed_len = read_le48((unsigned char *)&inbuf[4]);
91     uncompressed_len = read_le48((unsigned char *)&inbuf[10]);
92
93     comp_buf = calloc(1, compressed_len);
94     if (!comp_buf)
95     {
96         fprintf(G_logfile, "Unable to locate DATA block in template (offset %lld)\n",
97                 template_offset);
98         return ENOMEM;
99     }
100     
101     zip_state.data_buf = calloc(1, uncompressed_len);
102     if (!zip_state.data_buf)
103     {
104         fprintf(G_logfile, "Unable to allocate %lld bytes for decompression\n",
105                 uncompressed_len);
106         return ENOMEM;
107     }
108
109     read_num = fread(comp_buf, compressed_len, 1, template_file);
110     if (0 == read_num)
111     {
112         fprintf(G_logfile, "Unable to read %lld bytes for decompression\n",
113                 uncompressed_len);
114         return EIO;
115     }
116
117     error = decompress_data_block(comp_buf, compressed_len,
118                                   zip_state.data_buf, uncompressed_len, zip_state.algorithm);
119     if (error)
120     {
121         fprintf(G_logfile, "Unable to decompress data block, error %d\n", error);
122         return error;
123     }
124         
125     template_offset += compressed_len;
126     zip_state.buf_size = uncompressed_len;
127     zip_state.offset_in_curr_buf = 0;
128     free (comp_buf);
129     return 0;
130 }
131
132 static int skip_data_block(INT64 data_size, FILE *template_file)
133 {
134     int error = 0;
135     INT64 remaining = data_size;
136     INT64 size = 0;
137
138     /* If we're coming in in the middle of the image, we'll need to
139        skip through some compressed data */
140     while (remaining)
141     {
142         if (!zip_state.data_buf)
143         {
144             error = read_data_block(template_file);
145             if (error)
146             {
147                 fprintf(G_logfile, "Unable to decompress template data, error %d\n",
148                         error);
149                 return error;
150             }
151         }
152         size = MIN((zip_state.buf_size - zip_state.offset_in_curr_buf), remaining);
153         zip_state.offset_in_curr_buf += size;
154         remaining -= size;
155         
156         if (zip_state.offset_in_curr_buf == zip_state.buf_size)
157         {
158             free(zip_state.data_buf);
159             zip_state.data_buf = NULL;
160         }
161     }
162     
163     fprintf(G_logfile, "skip_data_block: skipped %lld bytes of unmatched data\n", data_size);
164     return error;
165 }
166
167 static int parse_data_block(INT64 data_size, FILE *template_file, struct mk_MD5Context *context, FILE *outfile)
168 {
169     int error = 0;
170     INT64 remaining = data_size;
171     INT64 size = 0;
172     int write_size = 0;
173
174     while (remaining)
175     {
176         if (!zip_state.data_buf)
177         {
178             error = read_data_block(template_file);
179             if (error)
180             {
181                 fprintf(G_logfile, "Unable to decompress template data, error %d\n",
182                         error);
183                 return error;
184             }
185         }
186         size = MIN((zip_state.buf_size - zip_state.offset_in_curr_buf), remaining);
187         write_size = fwrite(&zip_state.data_buf[zip_state.offset_in_curr_buf], size, 1, outfile);
188         if (!write_size)
189         {
190             fprintf(G_logfile, "parse_data_block: fwrite %lld failed with error %d; aborting\n", size, ferror(outfile));
191             return ferror(outfile);
192         }
193
194         if (G_verbose)
195             display_progress(outfile, "template data");
196
197         if (!G_quick)
198             mk_MD5Update(context,
199                                                  (unsigned char *)&zip_state.data_buf[zip_state.offset_in_curr_buf],
200                                                  size);
201         zip_state.offset_in_curr_buf += size;
202         remaining -= size;
203         
204         if (zip_state.offset_in_curr_buf == zip_state.buf_size)
205         {
206             free(zip_state.data_buf);
207             zip_state.data_buf = NULL;
208         }
209     }
210     if (G_verbose > 1)
211         fprintf(G_logfile, "parse_data_block: wrote %lld bytes of unmatched data\n", data_size);
212     return error;
213 }
214
215 static int read_file_data(char *filename, char *missing, INT64 offset, INT64 data_size,
216                           struct mk_MD5Context *file_context, struct mk_MD5Context *image_context,
217                           FILE *outfile)
218 {
219     FILE *input_file = NULL;
220     INT64 remaining = data_size;
221     char buf[BUF_SIZE];
222     int num_read = 0;
223     int write_size = 0;
224
225     input_file = fopen(filename, "rb");
226     if (!input_file)
227     {
228         fprintf(G_logfile, "Unable to open mirror file %s, error %d\n",
229                 filename, errno);
230         return errno;
231     }
232
233     if (missing)
234     {
235         fclose(input_file);
236         return 0;
237     }
238     
239     fseek(input_file, offset, SEEK_SET);
240     while (remaining)
241     {
242         int size = MIN(BUF_SIZE, remaining);
243         memset(buf, 0, BUF_SIZE);
244         
245         num_read = fread(buf, size, 1, input_file);
246         if (!num_read)
247         {
248             fprintf(G_logfile, "Unable to read from mirror file %s, error %d (offset %ld, length %d)\n",
249                     filename, errno, ftell(input_file), size);
250             fclose(input_file);
251             return errno;
252         }
253         if (file_context)
254         {
255             mk_MD5Update(image_context, (unsigned char *)buf, size);
256             mk_MD5Update(file_context, (unsigned char *)buf, size);
257         }
258         
259         write_size = fwrite(buf, size, 1, outfile);
260         if (!write_size)
261         {
262             fprintf(G_logfile, "read_file_data: fwrite %d failed with error %d; aborting\n", size, ferror(outfile));
263             return ferror(outfile);
264         }
265         
266         if (G_verbose)
267             display_progress(outfile, file_base_name(filename));
268         
269         remaining -= size;
270     }
271     if (G_verbose > 1)
272         fprintf(G_logfile, "read_file_data: wrote %lld bytes of data from %s\n",
273                 data_size, filename);
274     fclose(input_file);
275     return 0;
276 }
277
278 static int parse_file_block(INT64 offset, INT64 data_size, INT64 file_size, FILE *outfile,
279                             JIGDB *dbp, unsigned char *md5, struct mk_MD5Context *image_context,
280                             char *missing)
281 {
282     char *base64_md5 = base64_dump(md5, 16);
283     struct mk_MD5Context file_context;
284     struct mk_MD5Context *use_context = NULL;
285     unsigned char file_md5[16];
286     md5_list_t *md5_list_entry = NULL;
287     db_file_entry_t *db_entry = NULL;
288     int error = 0;
289     char *filename = NULL;
290
291     if (!G_quick)
292     {
293         use_context = &file_context;
294         mk_MD5Init(use_context);
295     }
296
297     /* Try the DB first if we have one */
298     if (dbp)
299     {
300         error = db_lookup_file_by_md5(dbp, base64_md5, &db_entry);
301         if (!error)
302             filename = db_entry->filename;
303     }
304
305     /* No joy; fall back to the MD5 list */
306     if (!filename)
307     {
308         md5_list_entry = find_file_in_md5_list(base64_md5);
309         if (md5_list_entry && file_size == md5_list_entry->file_size)
310             filename = md5_list_entry->full_path;
311     }
312
313     if (filename)
314     {
315         error = read_file_data(filename, missing, offset, data_size,
316                                use_context, image_context, outfile);
317         
318         if (error && (ENOENT != error))
319         {
320             fprintf(G_logfile, "Failed to read file %s, error %d\n", filename, error);
321             return error;
322         }
323         
324         if (!G_quick)
325         {
326             mk_MD5Final(file_md5, &file_context);
327             
328             if (memcmp(file_md5, md5, 16))
329             {
330                 fprintf(G_logfile, "MD5 MISMATCH for file %s\n", filename);
331                 fprintf(G_logfile, "    template looking for %s\n", base64_dump(md5, 16));
332                 fprintf(G_logfile, "    file %s is    %s\n", filename, base64_dump(file_md5, 16));
333                 return EINVAL;
334             }
335         }
336         return 0;
337     }
338     
339     /* No file found. Add it to the list of missing files, or complain */
340     if ( missing &&
341          md5_list_entry &&
342          (MISSING == md5_list_entry->file_size) &&
343          (!memcmp(md5_list_entry->md5, base64_md5, 16) ) )
344     {
345         file_missing(missing, md5_list_entry->full_path);
346         return 0;
347     }
348     /* else */
349     return ENOENT;
350 }
351
352 int parse_template_file(char *filename, int sizeonly, char *missing,
353                         FILE *outfile, char *output_name, JIGDB *dbp)
354 {
355     INT64 template_offset = 0;
356     INT64 bytes = 0;
357     char *buf = NULL;
358     FILE *file = NULL;
359     INT64 file_size = 0;
360     INT64 desc_start = 0;
361     INT64 written_length = 0;
362     INT64 output_offset = 0;
363     int i = 0;
364     int error = 0;
365     struct mk_MD5Context template_context;
366     unsigned char image_md5sum[16];
367
368     bzero(&zip_state, sizeof(zip_state));
369     
370     file = fopen(filename, "rb");
371     if (!file)
372     {
373         fprintf(G_logfile, "Failed to open template file %s, error %d!\n", filename, errno);
374         return errno;
375     }
376
377     buf = malloc(BUF_SIZE);
378     if (!buf)
379     {
380         fprintf(G_logfile, "Failed to malloc %d bytes. Abort!\n", BUF_SIZE);
381         fclose(file);
382         return ENOMEM;
383     }
384
385     /* Find the beginning of the desc block */
386     file_size = get_file_size(filename);
387     fseek(file, file_size - 6, SEEK_SET);
388     fread(buf, 6, 1, file);
389     desc_start = file_size - read_le48((unsigned char *)buf);
390
391     /* Now seek back to the beginning image desc block to grab the MD5
392        and image length */
393     fseek(file, file_size - 33, SEEK_SET);
394     fread(buf, BUF_SIZE, 1, file);
395     if (buf[0] != 5) /* image data */
396     {
397         fprintf(G_logfile, "Failed to find image desc in the template file\n");
398         fclose(file);
399         return EINVAL;
400     }
401
402     if (sizeonly)
403     {
404         fclose(file);
405         printf("%lld\n", read_le48((unsigned char *)&buf[1]));
406         return 0;
407     }
408
409     if (G_verbose)
410     {
411         fprintf(G_logfile, "Image MD5 should be    ");
412         for (i = 0; i < 16; i++)
413             fprintf(G_logfile, "%2.2x", (unsigned char)buf[i+7]);
414         fprintf(G_logfile, "\n");
415         fprintf(G_logfile, "Image size should be   %lld bytes\n", read_le48((unsigned char *)&buf[1]));
416     }
417
418     G_out_size = read_le48((unsigned char *)&buf[1]);
419     
420     /* Now seek back to the start of the desc block */
421     fseek(file, desc_start, SEEK_SET);
422     fread(buf, 10, 1, file);
423     if (strncmp(buf, "DESC", 4))
424     {
425         fprintf(G_logfile, "Failed to find desc start in the template file\n");
426         fclose(file);
427         return EINVAL;
428     }
429     if ((file_size - desc_start) != read_le48((unsigned char *)&buf[4]))
430     {
431         fprintf(G_logfile, "Inconsistent desc length in the template file!\n");
432         fprintf(G_logfile, "Final chunk says %lld, first chunk says %lld\n",
433                 file_size - desc_start, read_le48((unsigned char *)&buf[4]));
434         fclose(file);
435         return EINVAL;
436     }
437
438     if (!G_quick)
439         mk_MD5Init(&template_context);
440     template_offset = desc_start + 10;
441
442     if (1 == G_verbose)
443         fprintf(G_logfile, "Creating ISO image %s\n", output_name);
444
445     /* Main loop - walk through the template file and expand each entry we find */
446     while (1)
447     {
448         INT64 extent_size;
449         INT64 skip = 0;
450         INT64 read_length = 0;
451
452         if (template_offset >= (file_size - 33))
453         {
454             if (G_verbose > 1)
455                 fprintf(G_logfile, "Reached end of template file\n");
456             break; /* Finished! */
457         }
458         
459         if (output_offset > G_end_offset) /* Past the range we were asked for */
460         {
461             fprintf(G_logfile, "Reached end of range requested\n");            
462             break;
463         }
464         
465         fseek(file, template_offset, SEEK_SET);
466         bytes = fread(buf, (MIN (BUF_SIZE, file_size - template_offset)), 1, file);
467         if (1 != bytes)
468         {
469             fprintf(G_logfile, "Failed to read template file!\n");
470             fclose(file);
471             return EINVAL;
472         }
473         
474         extent_size = read_le48((unsigned char *)&buf[1]);
475         read_length = extent_size;
476         
477         if (G_start_offset > output_offset)
478             skip = G_start_offset - output_offset;
479         if ((output_offset + extent_size) > G_end_offset)
480             read_length -= (output_offset + extent_size - G_end_offset - 1);
481         read_length -= skip;
482         
483         switch (buf[0])
484         {
485             
486             case 2: /* unmatched data */
487                 template_offset += 7;
488                 if (missing)
489                     break;
490                 if ((output_offset + extent_size) >= G_start_offset)
491                 {
492                     if (skip)
493                         error = skip_data_block(skip, file);
494                     if (error)
495                     {
496                         fprintf(G_logfile, "Unable to read data block to skip, error %d\n", error);
497                         fclose(file);
498                         return error;
499                     }
500                     error = parse_data_block(read_length, file, &template_context, outfile);
501                     if (error)
502                     {
503                         fprintf(G_logfile, "Unable to read data block, error %d\n", error);
504                         fclose(file);
505                         return error;
506                     }
507                     written_length += read_length;
508                 }
509                 else
510                     error = skip_data_block(extent_size, file);
511                 break;
512             case 6:
513                 template_offset += 31;
514                 if ((output_offset + extent_size) >= G_start_offset)
515                 {
516                     error = parse_file_block(skip, read_length, extent_size, outfile, dbp,
517                                              (unsigned char *)&buf[15], &template_context, missing);
518                     if (error)
519                     {
520                         fprintf(G_logfile, "Unable to read file block, error %d\n", error);
521                         fclose(file);
522                         return error;
523                     }
524                     written_length += read_length;
525                 }
526                 break;
527             default:
528                 fprintf(G_logfile, "Unknown block type %d!\n", buf[0]);
529                 fclose(file);
530                 return EINVAL;
531         }
532         output_offset += extent_size;
533     }
534
535     if (missing && G_missing_file)
536         return ENOENT;
537     
538     fclose(file);
539     if (G_verbose)
540     {
541         fprintf(G_logfile, "\n");
542         if (!G_quick)
543         {
544             mk_MD5Final (image_md5sum, &template_context);
545             fprintf(G_logfile, "Output image MD5 is    ");
546             for (i = 0; i < 16; i++)
547                 fprintf(G_logfile, "%2.2x", image_md5sum[i]);
548             fprintf(G_logfile, "\n");
549         }
550         fprintf(G_logfile, "Output image length is %lld bytes\n", written_length);
551     }
552     
553     return 0;
554 }
555