Fix write length so it includes the updated header (tag) CRC32

[firmware_header_edit.git] / firmware_header_editor.c

diff --git a/firmware_header_editor.c b/firmware_header_editor.c
index 78e4177..ca0e24b 100644 (file)
--- a/firmware_header_editor.c
+++ b/firmware_header_editor.c
@@ -1,7 +1,7 @@
 static const char *title =\
 "Broadcom Consumer Router Firmware Header Editor"
 ;
-static const float VERSION = 1.2f;
+static const float VERSION = 1.3f;
 
 static const char *copyright = \
 "Copyright 2015-2016 TJ <hacker@iam.tj>\n"
@@ -53,6 +53,8 @@ static const char *help = \
 #include <fcntl.h>
 #include <unistd.h>
 #include <arpa/inet.h>
+#include <stdint.h>
+#include "heap_reap.h"
 
 static unsigned int MESSAGE_SIZE = 1024;
 static unsigned int TAG_VER_LEN = 4;
@@ -137,6 +139,7 @@ pr_error_exit(unsigned int usage, const char *error, ...)
 
  if (usage) pr_usage(usage);
 
+ heap_and_reap(NULL-1, 0, 0);
  exit(EXIT_FAILURE);
 }
 
@@ -175,7 +178,7 @@ main(int argc, char **argv)
 
   fprintf(stderr, "%s\nVersion: %0.2f\n%s\n", title, VERSION, copyright);
 
-  if ((model_id = calloc(MODEL_ID_LEN, 1)) == NULL) {
+  if ((model_id = heap_and_reap(NULL, MODEL_ID_LEN, 1)) == NULL) {
     pr_error_exit(0, "Unable to allocate memory (%d bytes)\n", MODEL_ID_LEN);
   }
   memcpy(model_id, match_model_id, MODEL_ID_LEN);
@@ -214,7 +217,7 @@ main(int argc, char **argv)
             break;
           case 'h': // help
             pr_usage(1);
-            exit(0);
+            goto end;
         }
       } else {
         pr_error_exit(0, "cannot read data from stdin; provide a filename");
@@ -266,7 +269,7 @@ main(int argc, char **argv)
 
   if (filename) {
    if ((fd = open(filename, fd_mode)) > 0) {
-     if ( (buffer = calloc(header_len + sizeof(crc), 1)) != NULL) {
+     if ( (buffer = heap_and_reap(NULL, header_len + sizeof(crc), 1)) != NULL) {
        ssize_t qty;
        if ( (qty = read(fd, buffer, header_len + sizeof(crc))) < header_len) {
          if (!opt_quiet)
@@ -321,8 +324,9 @@ main(int argc, char **argv)
        memcpy(buffer + header_crc_offset, &tmp, sizeof(tmp));
 
        if (!opt_simulate) {
+         ssize_t write_len = header_len > header_crc_offset + sizeof(tmp) ? header_len : header_crc_offset + sizeof(tmp);
          lseek(fd, 0, SEEK_SET);
-         write(fd, buffer, header_len);
+         write(fd, buffer, write_len);
        }
 
        printf( format_spec,
@@ -348,7 +352,7 @@ main(int argc, char **argv)
          );
      }
 
-     free(buffer);
+     heap_and_reap(buffer, 0, 0);
      close(fd);
     } else {
       fprintf(stderr, "Unable to open for %s (%s)\n", opt_write ? "writing" : "reading" , filename );
@@ -357,6 +361,8 @@ main(int argc, char **argv)
     pr_usage(0);
   }
 
+end:
+  heap_and_reap(NULL-1, 0, 0);
   return 0;
 }