aw - Write whole files atomically in Linux

aw.c

/*
  aw.c - write whole files atomically
  
  Atomically create a fully written file with contents read from stdin
  
  Usage: aw <destination>
  
  aw reads from stdin until EOF and writes to an anonymous temporary file.
  When EOF is reached and all contents have been written to the temporary
  file, the file is atomically linked with the supplied destination filename. 
  If any error happens (e.g. during write) no destination file is ever 
  created.
  
  This allows to atomically create fully-formed files, i.e. it's impossible
  for a different process to see either the temporary file (at all) or the 
  destination file before it has been fully written.
  
  If a file with the same destination filename already exists the command 
  fails.
  
  If write succeeded, aw returns 0. Otherwise it returns a non-0 return code
  and prints to stderr an error message.
  
  Usage examples:
    gunzip -c somefile.gz | aw somefile   # atomic gunzip
    aw destination <source                # atomic cp (only data)
  
  Notes:
  - requires Linux 3.11+ (only tested on amd64)
  
  License: MIT
  Author: Carlo Alberto Ferraris <[email protected]>
*/

// for O_TMPFILE
#define _GNU_SOURCE 1 

#include <fcntl.h>
#include <unistd.h>
#include <string.h>
#include <limits.h>
#include <stdio.h>
#include <libgen.h>
#include <string.h>

#define DEFAULT_BUFFER_SIZE (1<<18)
#define MAX_BUFFER_SIZE (1<<24)
#define MIN_BUFFER_SIZE (1<<9)

int main(int argc, char **argv) {
  int buffer_size = DEFAULT_BUFFER_SIZE;
  int use_fsync = 1;
  
  while ((c = getopt (argc, argv, "b:n")) != -1) {
    switch (c) {
    case 'n':
      use_fsync = 0;
      break;
    case 'b':
      buffer_size = atoi(optarg);
      if (buffer_size < MIN_BUFFER_SIZE || buffer_size > MAX_BUFFER_SIZE) {
        fprintf(stderr, "Invalid argument for option -b: \"%s\".\n", optarg);
        return -1;
      }
      break;
    case '?':
      if (optopt == 'b')
        fprintf(stderr, "Option -%c requires an argument.\n", optopt);
      else if (isprint (optopt))
        fprintf(stderr, "Unknown option `-%c'.\n", optopt);
      else
        fprintf(stderr, "Unknown option character `\\x%x'.\n", optopt);
      return -1;
    default:
      abort();
    }
  }
  
  if (optind != argc-1) {
    fprintf(stderr, "Usage: aw <destination>\n");
    return -1;
  }

  int in_fd = STDIN_FILENO;
  char *out_file = argv[1];
  char *out_dir = dirname(strndupa(out_file, PATH_MAX));
  
  int out_dir_fd = open(out_dir, O_DIRECTORY | O_RDONLY, 0);
  if (out_dir_fd < 0) {
    perror("Error opening destination directory");
    return -7;
  }
  
  // TODO: optionally allow to specify user, group and permissions
  int out_fd = openat(out_dir_fd, ".", O_TMPFILE | O_WRONLY, S_IRUSR | S_IWUSR);
  if (out_fd < 0) {
    perror("Error creating temporary file");
    return -5;
  }
  
  struct stat sb;
  if (fstat(in_fd, &sb) != 0) {
    perror("Error getting input file status");
    return -10;
  }
  
  switch (sb.st_mode & S_IFMT) {
  case S_IFIFO:
    int pipe_sz = fcntl(in_fd, F_GETPIPE_SZ);
    if (pipe_sz >= 0 && buffer_size > pipe_sz)
      fcntl(in_fd, F_SETPIPE_SZ, buffer_size);
    break;
  case S_IFREG:
    posix_fadvise(in_fd, 0, 0, POSIX_FADV_SEQUENTIAL | POSIX_FADV_WILLNEED | POSIX_FADV_NOREUSE);
    posix_fallocate(out_fd, 0, sb.st_size);
    break;
  }

  char *buf = malloc(buffer_size);
  if (!buf) {
    perror("Error allocating buffer");
    return -11;
  }
  
  int res;
  do {
    // TODO: sendfile/splice/copy_file_range/... ?
    res = read(in_fd, buf, buffer_size);
    if (res < 0) {
      perror("Error reading from stdin");
      return -2;
    }

    int wres = write(out_fd, buf, res);
    if (wres != res) {
      perror("Error writing to temporary file");
      return -3;
    }
  } while (res > 0);
  
  free(buf);
  
  if (use_fsync) {
    res = fsync(out_fd);
    if (res != 0) {
      perror("Error flushing temporary file");
      return -6;
    }
  }
  
  char path[PATH_MAX];
  snprintf(path, PATH_MAX, "/proc/self/fd/%d", out_fd);
  res = linkat(AT_FDCWD, path, AT_FDCWD, out_file, AT_SYMLINK_FOLLOW);
  if (res != 0) {
    perror("Error linking file");
    return -4;
  }
  
  /* From this point onward the file is linked and threfore visible; we can only
     attempt to complete without failing */
  
  int rc = 0;
  
  res = close(out_fd);
  if (res != 0) {
    perror("Error closing file");
    rc |= (1<<0);
  }

  if (use_fsync) {
    res = fsync(out_dir_fd);
    if (res != 0) {
      perror("Error flushing destination directory");
      rc |= (1<<1)
    }
  }
  
  res = close(out_dir_fd);
  if (res != 0) {
    perror("Error closing destination directory");
    rc |= (1<<2);
  }

  return rc;
}

Makefile

CFLAGS=-O2 -flto -ffunction-sections -fdata-sections -Wl,--gc-sections
CFLAGSNATIVE=-march=native -mtune=native
CFLAGSSTATIC=-static

build: aw.c
	gcc -o aw aw.c $(CFLAGS)
     
build-static: aw.c
	gcc -o aw aw.c $(CFLAGS) $(CFLAGSSTATIC)
     
build-native: aw.c
	gcc -o aw aw.c $(CFLAGS) $(CFLAGSNATIVE)
     
build-native-static: aw.c
	gcc -o aw aw.c $(CFLAGS) $(CFLAGSNATIVE) $(CFLAGSSTATIC)