vbig.cc

   1 #include <config.h>
   2 #include <cstdio>
   3 #include <cstring>
   4 #include <cstdlib>
   5 #include <cerrno>
   6 #include <cstdarg>
   7 #include <getopt.h>
   8 #include <unistd.h>
   9 #include <fcntl.h>
  10 #include "Arcfour.h"
  11
  12 // Path to magic file to drop filesystem caches
  13 static const char dropCaches[] = "/proc/sys/vm/drop_caches";
  14
  15 // Command line options
  16 const struct option opts[] = {
  17   { "seed", required_argument, 0, 's' },
  18   { "verify", no_argument, 0, 'v' },
  19   { "create", no_argument, 0, 'c' },
  20   { "flush", no_argument, 0, 'f' },
  21   { "help", no_argument, 0, 'h' },
  22   { "version", no_argument, 0, 'V' },
  23   { 0, 0, 0, 0 },
  24 };
  25
  26 // Display help message
  27 static void help(void) {
  28   printf("vbig - create or verify a large but pseudo-random file\n"
  29          "\n"
  30          "Usage:\n"
  31          "  vbig [--seed SEED] --verify|--create PATH SIZE\n"
  32          "\n"
  33          "Options:\n"
  34          "  --seed, -s     Specify random seed\n"
  35          "  --verify, -v   Verify that PATH contains the expected contents\n"
  36          "  --create, -c   Create PATH with psuedo-random contents\n"
  37          "  --flush, -f    Flush cache\n"
  38          "  --help, -h     Display usage message\n"
  39          "  --version, -V  Display version string\n");
  40 }
  41
  42 // Possible modes of operation
  43 enum mode_type {
  44   NONE,
  45   VERIFY,
  46   CREATE
  47 };
  48
  49 // Report an error and exit
  50 static void fatal(int errno_value, const char *fmt, ...) {
  51   va_list ap;
  52   fprintf(stderr, "ERROR: ");
  53   va_start(ap, fmt);
  54   vfprintf(stderr, fmt, ap);
  55   va_end(ap);
  56   if(errno_value)
  57     fprintf(stderr, ": %s", strerror(errno_value));
  58   fputc('\n', stderr);
  59   exit(1);
  60 }
  61
  62 // Evict whatever FP points to from RAM
  63 static void flushCache(FILE *fp) {
  64   // drop_caches only evicts clean pages, so first the target file is
  65   // synced.
  66   if(fsync(fileno(fp)) < 0)
  67     fatal(errno, "fsync");
  68   int fd;
  69   if((fd = open(dropCaches, O_WRONLY, 0)) < 0)
  70     fatal(errno, "%s", dropCaches);
  71   if(write(fd, "3\n", 2) < 0)
  72     fatal(errno, "%s", dropCaches);
  73   close(fd);
  74 }
  75
  76 int main(int argc, char **argv) {
  77   const char *seed = "hexapodia as the key insight";
  78   mode_type mode = NONE;
  79   bool flush = false;
  80   int n;
  81   while((n = getopt_long(argc, argv, "+s:vcfhV", opts, 0)) >= 0) {
  82     switch(n) {
  83     case 's': seed = optarg; break;
  84     case 'v': mode = VERIFY; break;
  85     case 'c': mode = CREATE; break;
  86     case 'f': flush = true; break;
  87     case 'h': help(); exit(0);
  88     case 'V': puts(VERSION); exit(0);
  89     default:
  90       fatal(0, "unknown option");
  91     }
  92   }
  93   if(mode == NONE) {
  94     fatal(0, "must specify one of --verify or --create");
  95     exit(1);
  96   }
  97   if(optind + 2 != argc) {
  98     fatal(0, "must specify a path and size");
  99     exit(1);
 100   }
 101   const char *path = argv[optind];
 102   errno = 0;
 103   char *end;
 104   long long size = strtoll(argv[optind + 1], &end, 10);
 105   if(errno) {
 106     fatal(errno, "invalid size");
 107     exit(1);
 108   }
 109   if(end == argv[optind + 1]) {
 110     fatal(0, "invalid size");
 111     exit(1);
 112   }
 113   if(!strcmp(end, "K"))
 114     size *= 1024;
 115   else if(!strcmp(end, "M"))
 116     size *= 1024 * 1024;
 117   else if(!strcmp(end, "G"))
 118     size *= 1024 * 1024 * 1024;
 119   else if(*end) {
 120     fatal(0, "invalid size");
 121     exit(1);
 122   }
 123   Arcfour rng(seed, strlen(seed));
 124   FILE *fp = fopen(path, mode == VERIFY ? "rb" : "wb");
 125   if(!fp) {
 126     fatal(errno, "%s", path);
 127     exit(1);
 128   }
 129   if(mode == VERIFY && flush)
 130     flushCache(fp);
 131   char generated[4096], input[4096];
 132   long long remain = size;
 133   while(remain > 0) {
 134     size_t bytesGenerated = (remain > (ssize_t)sizeof generated
 135                              ? sizeof generated
 136                              : remain);
 137     rng.stream(generated, bytesGenerated);
 138     if(mode == CREATE) {
 139       fwrite(generated, 1, bytesGenerated, fp);
 140       if(ferror(fp)) {
 141         fatal(errno, "%s", path);
 142         exit(1);
 143       }
 144     } else {
 145       size_t bytesRead = fread(input, 1, bytesGenerated, fp);
 146       if(ferror(fp)) {
 147         fatal(errno, "%s", path);
 148         exit(1);
 149       }
 150       if(bytesRead < bytesGenerated) {
 151         fatal(0, "%s: truncated at %lld/%lld bytes",
 152                 path, (size - remain + bytesRead), size);
 153         exit(1);
 154       }
 155       if(memcmp(generated, input, bytesGenerated)) {
 156         for(size_t n = 0; n < bytesGenerated; ++n)
 157           if(generated[n] != input[n]){
 158             fatal(0, "%s corrupted at %lld/%lld bytes (expected %d got %d)",
 159                     path, size - remain + n, size,
 160                     (unsigned char)generated[n], (unsigned char)input[n]);
 161             exit(1);
 162           }
 163       }
 164     }
 165     remain -= bytesGenerated;
 166   }
 167   if(mode == VERIFY && getc(fp) != EOF) {
 168     fatal(0, "%s: extended beyond %lld bytes",
 169             path, size);
 170     exit(1);
 171   }
 172   if(mode == CREATE && flush) {
 173     if(fflush(fp) < 0)
 174       fatal(errno, "%s", path);
 175     flushCache(fp);
 176   }
 177   if(fclose(fp) < 0) {
 178     fatal(errno, "%s", path);
 179     exit(1);
 180   }
 181   return 0;
 182 }