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 }
  24
  25 enum mode_type {
  26   NONE,
  27   VERIFY,
  28   CREATE
  29 };
  30
  31 int main(int argc, char **argv) {
  32   const char *seed = "hexapodia as the key insight";
  33   mode_type mode = NONE;
  34   int n;
  35   while((n = getopt_long(argc, argv, "+s:vchV", opts, 0)) >= 0) {
  36     switch(n) {
  37     case 's': seed = optarg; break;
  38     case 'v': mode = VERIFY; break;
  39     case 'c': mode = CREATE; break;
  40     case 'h': help(); exit(0);
  41     case 'V': puts(VERSION); exit(0);
  42     default:
  43       fprintf(stderr, "ERROR: unknown option\n");
  44       exit(1);
  45     }
  46   }
  47   if(mode == NONE) {
  48     fprintf(stderr, "ERROR: must specify one of --verify or --create\n");
  49     exit(1);
  50   }
  51   if(optind + 2 != argc) {
  52     fprintf(stderr, "ERROR: must specify a path and size\n");
  53     exit(1);
  54   }
  55   const char *path = argv[optind];
  56   errno = 0;
  57   char *end;
  58   long long size = strtoll(argv[optind + 1], &end, 10);
  59   if(errno) {
  60     fprintf(stderr, "ERROR: invalid size: %s\n", strerror(errno));
  61     exit(1);
  62   }
  63   if(end == argv[optind + 1]) {
  64     fprintf(stderr, "ERROR: invalid size\n");
  65     exit(1);
  66   }
  67   if(!strcmp(end, "K"))
  68     size *= 1024;
  69   else if(!strcmp(end, "M"))
  70     size *= 1024 * 1024;
  71   else if(!strcmp(end, "G"))
  72     size *= 1024 * 1024 * 1024;
  73   else if(*end) {
  74     fprintf(stderr, "ERROR: invalid size\n");
  75     exit(1);
  76   }
  77   Arcfour rng(seed, strlen(seed));
  78   FILE *fp = fopen(path, mode == VERIFY ? "rb" : "wb");
  79   if(!fp) {
  80     fprintf(stderr, "ERROR: %s: %s\n", path, strerror(errno));
  81     exit(1);
  82   }
  83   char generated[4096], input[4096];
  84   size_t remain = size;
  85   while(remain > 0) {
  86     size_t bytesGenerated = (remain > sizeof generated
  87                              ? sizeof generated
  88                              : remain);
  89     rng.stream(generated, bytesGenerated);
  90     if(mode == CREATE) {
  91       fwrite(generated, 1, bytesGenerated, fp);
  92       if(ferror(fp)) {
  93         fprintf(stderr, "ERROR: %s: %s\n", path, strerror(errno));
  94         exit(1);
  95       }
  96     } else {
  97       size_t bytesRead = fread(input, 1, bytesGenerated, fp);
  98       if(ferror(fp)) {
  99         fprintf(stderr, "ERROR: %s: %s\n", path, strerror(errno));
 100         exit(1);
 101       }
 102       if(bytesRead < bytesGenerated) {
 103         fprintf(stderr, "ERROR: %s: truncated at %lld/%lld bytes\n",
 104                 path, (size - remain + bytesRead), size);
 105         exit(1);
 106       }
 107       if(memcmp(generated, input, bytesGenerated)) {
 108         for(size_t n = 0; n < bytesGenerated; ++n)
 109           if(generated[n] != input[n]){
 110             fprintf(stderr, "ERROR: %s corrupted at %lld/%lld bytes (expected %d got %d)\n",
 111                     path, size - remain + n, size,
 112                     (unsigned char)generated[n], (unsigned char)input[n]);
 113             exit(1);
 114           }
 115       }
 116     }
 117     remain -= bytesGenerated;
 118   }
 119   if(mode == VERIFY && getc(fp) != EOF) {
 120     fprintf(stderr, "ERROR: %s: extended beyond %lld bytes\n",
 121             path, size);
 122     exit(1);
 123   }
 124   if(fclose(fp) < 0) {
 125     fprintf(stderr, "ERROR: %s: %s\n", path, strerror(errno));
 126     exit(1);
 127   }
 128   return 0;
 129 }