New Full AU Format w/ support for kernel partition.

src/platform/memento_softwareupdate/split_write.cc

+// Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include <sys/stat.h>
+#include <sys/types.h>
+
+#include <endian.h>
+#include <fcntl.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+
+#include <algorithm>
+
+namespace {
+
+class ScopedFileDescriptorCloser {
+ public:
+  ScopedFileDescriptorCloser(int fd) : fd_(fd) {}
+  ~ScopedFileDescriptorCloser() {
+    close(fd_);
+  }
+ private:
+  const int fd_;
+};
+
+const int kBufSize = 1024 * 1024 * 4;  // 4 MiB
+
+// This program takes two files as args. It will open both and write the
+// first part of stdin into the first file, the second part into the second
+// file. The first 8 bytes contain the unsigned big-endian count of bytes
+// that should go to the first file. Following bytes go to the second file.
+
+// Writes all bytes to fd. Exits on error.
+void write_all(int fd, const void *buf, size_t count) {
+  const char* c_buf = static_cast<const char*>(buf);
+  size_t written = 0;
+  while (written < count) {
+    ssize_t rc = write(fd, c_buf + written, count - written);
+    if (rc < 0) {
+      perror("write");
+      exit(1);
+    }
+    written += static_cast<size_t>(rc);
+  }
+}
+
+// Returns bytes read, which may be short on EOF. Exits on error.
+size_t read_all(int fd, void* buf, size_t count) {
+  char* c_buf = static_cast<char*>(buf);
+  size_t bytes_read = 0;
+  while (bytes_read < count) {
+    ssize_t rc = read(fd, c_buf + bytes_read, count - bytes_read);
+    if (rc == 0) {
+      break;
+    }
+    if (rc < 0) {
+      perror("read");
+      exit(1);
+    }
+    bytes_read += static_cast<size_t>(rc);
+  }
+  return bytes_read;
+}
+
+void usage(char* argv0) {
+  fprintf(stderr, "Usage: %s first_file second_file\n", argv0);
+  exit(1);
+}
+
+// Returns valid fd or exits program.
+int open_file(const char* path) {
+  int fd = open(path, O_WRONLY | O_CREAT | O_TRUNC, 0666);
+  if (fd < 0) {
+    perror("open");
+    fprintf(stderr, "failed to open %s\n", path);
+    exit(1);
+  }
+}
+
+typedef long long int64;
+// Compile assert sizeof(int64) == 8:
+char int64_8_bytes_long[(sizeof(int64) == 8) - 1];
+
+}  // namespace {}
+
+int main(int argc, char** argv) {
+  if (argc != 3) {
+    usage(argv[0]);
+  }
+  const int fd1 = open_file(argv[1]);
+  ScopedFileDescriptorCloser fd1_closer(fd1);
+  const int fd2 = open_file(argv[2]);
+  ScopedFileDescriptorCloser fd2_closer(fd2);
+  const int fd_in = 0;  // stdin
+  char* const buf = static_cast<char*>(malloc(kBufSize));
+  if (buf == NULL) {
+    fprintf(stderr, "malloc on buffer failed.\n");
+    return 1;
+  }
+  int64 first_file_size = 0;
+  size_t bytes_read =
+      read_all(fd_in, &first_file_size, sizeof(first_file_size));
+  if (bytes_read < sizeof(first_file_size)) {
+    fprintf(stderr, "short read on first file size.\n");
+    return 1;
+  }
+  first_file_size = be64toh(first_file_size);
+  int64 first_bytes_written = 0;
+  while (first_bytes_written < first_file_size) {
+    size_t chunk_size = std::min(first_file_size - first_bytes_written,
+                                 static_cast<int64>(kBufSize));
+    chunk_size = read_all(fd_in, buf, chunk_size);
+    if (chunk_size == 0) {
+      // All data went to first partition, none left for second.
+      // This is okay only if the first file size is exactl how much we've
+      // written thus far
+      if (first_file_size == first_bytes_written) {
+        return 0;
+      } else {
+        fprintf(stderr, "file appears truncated.\n");
+        return 1;
+      }
+    }
+    write_all(fd1, buf, chunk_size);
+    first_bytes_written += chunk_size;
+  }
+  // Do the rest on the second file
+  for (;;) {
+    size_t chunk_size = read_all(fd_in, buf, kBufSize);
+    if (chunk_size == 0)
+      break;
+    write_all(fd2, buf, chunk_size);
+  }
+  return 0;
+}