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