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