vbig.cc

   1 #include <config.h>
   2 #include <cstdio>
   3 #include <cstring>
   4 #include <cstdlib>
   5 #include <cerrno>
   6 #include <getopt.h>
   7 #include "Arcfour.h"
   8
   9 const struct option opts[] = {
  10   { "seed", required_argument, 0, 's' },
  11   { "verify", no_argument, 0, 'v' },
  12   { "create", no_argument, 0, 'c' },
  13   { "help", no_argument, 0, 'h' },
  14   { "version", no_argument, 0, 'V' },
  15   { 0, 0, 0, 0 },
  16 };
  17
  18 static void help(void) {
  19   printf("vbig - create or verify a large but pseudo-random file\n"
  20          "\n"
  21          "Usage:\n"
  22          "  vbig [--seed SEED] --verify|--create PATH SIZE\n"
  23          "\n"
  24          "Options:\n"
  25          "  --seed         Specify random seed\n"
  26          "  --verify       Verify that PATH contains the expected contents\n"
  27          "  --create       Create PATH with psuedo-random contents\n"
  28          "  --help         Display usage message\n"
  29          "  --version      Display version string\n");
  30 }
  31
  32 enum mode_type {
  33   NONE,
  34   VERIFY,
  35   CREATE
  36 };
  37
  38 int main(int argc, char **argv) {
  39   const char *seed = "hexapodia as the key insight";
  40   mode_type mode = NONE;
  41   int n;
  42   while((n = getopt_long(argc, argv, "+s:vchV", opts, 0)) >= 0) {
  43     switch(n) {
  44     case 's': seed = optarg; break;
  45     case 'v': mode = VERIFY; break;
  46     case 'c': mode = CREATE; break;
  47     case 'h': help(); exit(0);
  48     case 'V': puts(VERSION); exit(0);
  49     default:
  50       fprintf(stderr, "ERROR: unknown option\n");
  51       exit(1);
  52     }
  53   }
  54   if(mode == NONE) {
  55     fprintf(stderr, "ERROR: must specify one of --verify or --create\n");
  56     exit(1);
  57   }
  58   if(optind + 2 != argc) {
  59     fprintf(stderr, "ERROR: must specify a path and size\n");
  60     exit(1);
  61   }
  62   const char *path = argv[optind];
  63   errno = 0;
  64   char *end;
  65   long long size = strtoll(argv[optind + 1], &end, 10);
  66   if(errno) {
  67     fprintf(stderr, "ERROR: invalid size: %s\n", strerror(errno));
  68     exit(1);
  69   }
  70   if(end == argv[optind + 1]) {
  71     fprintf(stderr, "ERROR: invalid size\n");
  72     exit(1);
  73   }
  74   if(!strcmp(end, "K"))
  75     size *= 1024;
  76   else if(!strcmp(end, "M"))
  77     size *= 1024 * 1024;
  78   else if(!strcmp(end, "G"))
  79     size *= 1024 * 1024 * 1024;
  80   else if(*end) {
  81     fprintf(stderr, "ERROR: invalid size\n");
  82     exit(1);
  83   }
  84   Arcfour rng(seed, strlen(seed));
  85   FILE *fp = fopen(path, mode == VERIFY ? "rb" : "wb");
  86   if(!fp) {
  87     fprintf(stderr, "ERROR: %s: %s\n", path, strerror(errno));
  88     exit(1);
  89   }
  90   char generated[4096], input[4096];
  91   size_t remain = size;
  92   while(remain > 0) {
  93     size_t bytesGenerated = (remain > sizeof generated
  94                              ? sizeof generated
  95                              : remain);
  96     rng.stream(generated, bytesGenerated);
  97     if(mode == CREATE) {
  98       fwrite(generated, 1, bytesGenerated, fp);
  99       if(ferror(fp)) {
 100         fprintf(stderr, "ERROR: %s: %s\n", path, strerror(errno));
 101         exit(1);
 102       }
 103     } else {
 104       size_t bytesRead = fread(input, 1, bytesGenerated, fp);
 105       if(ferror(fp)) {
 106         fprintf(stderr, "ERROR: %s: %s\n", path, strerror(errno));
 107         exit(1);
 108       }
 109       if(bytesRead < bytesGenerated) {
 110         fprintf(stderr, "ERROR: %s: truncated at %lld/%lld bytes\n",
 111                 path, (size - remain + bytesRead), size);
 112         exit(1);
 113       }
 114       if(memcmp(generated, input, bytesGenerated)) {
 115         for(size_t n = 0; n < bytesGenerated; ++n)
 116           if(generated[n] != input[n]){
 117             fprintf(stderr, "ERROR: %s corrupted at %lld/%lld bytes (expected %d got %d)\n",
 118                     path, size - remain + n, size,
 119                     (unsigned char)generated[n], (unsigned char)input[n]);
 120             exit(1);
 121           }
 122       }
 123     }
 124     remain -= bytesGenerated;
 125   }
 126   if(mode == VERIFY && getc(fp) != EOF) {
 127     fprintf(stderr, "ERROR: %s: extended beyond %lld bytes\n",
 128             path, size);
 129     exit(1);
 130   }
 131   if(fclose(fp) < 0) {
 132     fprintf(stderr, "ERROR: %s: %s\n", path, strerror(errno));
 133     exit(1);
 134   }
 135   return 0;
 136 }