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);
  }
  return fd;
}

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;
}