chiark / gitweb /
Use RC4-drop, not RC4
[vbig.git] / vbig.cc
diff --git a/vbig.cc b/vbig.cc
index 3a35728aa82d115f2707b120ec1d2c5a63fa0301..1823dd316aaacba1cc7d1efb7bea0d6d4e300e29 100644 (file)
--- a/vbig.cc
+++ b/vbig.cc
 #include <getopt.h>
 #include <unistd.h>
 #include <fcntl.h>
+#include <limits.h>
+#include <assert.h>
 #include <sys/stat.h>
 #include "Arcfour.h"
 
 // Command line options
 const struct option opts[] = {
   { "seed", required_argument, 0, 's' },
+  { "both", no_argument, 0, 'b' },
   { "verify", no_argument, 0, 'v' },
   { "create", no_argument, 0, 'c' },
   { "flush", no_argument, 0, 'f' },
+  { "entire", no_argument, 0, 'e' },
   { "help", no_argument, 0, 'h' },
   { "version", no_argument, 0, 'V' },
   { 0, 0, 0, 0 },
@@ -50,15 +54,16 @@ static void help(void) {
          "  --verify, -v   Verify that PATH contains the expected contents\n"
          "  --create, -c   Create PATH with psuedo-random contents\n"
          "  --flush, -f    Flush cache\n"
+         "  --entire, -e   Write until full; read until EOF\n"
          "  --help, -h     Display usage message\n"
          "  --version, -V  Display version string\n");
 }
 
 // Possible modes of operation
 enum mode_type {
-  NONE,
   VERIFY,
-  CREATE
+  CREATE,
+  BOTH
 };
 
 // Report an error and exit
@@ -103,16 +108,24 @@ static void flushCache(FILE *fp) {
 #endif
 }
 
+static long long execute(mode_type mode, bool entire, const char *show);
+
+static const char *seed = "hexapodia as the key insight";
+static const char *path;
+static bool entireopt = false;
+static bool flush = false;
+static long long size;
+
 int main(int argc, char **argv) {
-  const char *seed = "hexapodia as the key insight";
-  mode_type mode = NONE;
-  bool flush = false;
+  mode_type mode = BOTH;
   int n;
-  while((n = getopt_long(argc, argv, "+s:vcfhV", opts, 0)) >= 0) {
+  while((n = getopt_long(argc, argv, "+s:vcefhV", opts, 0)) >= 0) {
     switch(n) {
     case 's': seed = optarg; break;
+    case 'b': mode = BOTH; break;
     case 'v': mode = VERIFY; break;
     case 'c': mode = CREATE; break;
+    case 'e': entireopt = true; break;
     case 'f': flush = true; break;
     case 'h': help(); exit(0);
     case 'V': puts(VERSION); exit(0);
@@ -122,14 +135,18 @@ int main(int argc, char **argv) {
   }
   argc -= optind;
   argv += optind;
-  if(mode == NONE)
-    fatal(0, "must specify one of --verify or --create");
   if(argc > 2)
     fatal(0, "excess arguments");
-  if(argc < (mode == VERIFY ? 1 : 2))
-    fatal(0, "insufficient arguments");
-  const char *path = argv[0];
-  long long size;
+  if(argc == 1 && mode == BOTH)
+    entireopt = true;
+  if(entireopt) {
+    if(argc != 1)
+      fatal(0, "with --entire, size should not be specified");
+  } else {
+    if(argc < (mode == VERIFY ? 1 : 2))
+      fatal(0, "insufficient arguments");
+  }
+  path = argv[0];
   if(argc > 1) {
     errno = 0;
     char *end;
@@ -146,47 +163,76 @@ int main(int argc, char **argv) {
       size *= 1024 * 1024 * 1024;
     else if(*end)
       fatal(0, "invalid size");
+  } else if(entireopt) {
+    size = LONG_LONG_MAX;
   } else {
     struct stat sb;
     if(stat(path, &sb) < 0)
       fatal(errno, "stat %s", path);
     size = sb.st_size;
   }
+  const char *show = entireopt ? (mode == CREATE ? "written" : "verified") : 0;
+  if(mode == BOTH) {
+    size = execute(CREATE, entireopt, 0);
+    execute(VERIFY, false, show);
+  } else {
+    execute(mode, entireopt, show);
+  }
+  return 0;
+}
+
+static long long execute(mode_type mode, bool entire, const char *show) {
   Arcfour rng(seed, strlen(seed));
   FILE *fp = fopen(path, mode == VERIFY ? "rb" : "wb");
   if(!fp)
     fatal(errno, "%s", path);
   if(mode == VERIFY && flush)
     flushCache(fp);
+  if(mode == CREATE && entire)
+    setvbuf(fp, 0, _IONBF, 0);
   char generated[4096], input[4096];
   long long remain = size;
+  static const size_t rc4drop = 3072; // en.wikipedia.org/wiki/RC4#Security
+  assert(rc4drop <= sizeof(generated));
+  rng.stream(generated, rc4drop);
   while(remain > 0) {
     size_t bytesGenerated = (remain > (ssize_t)sizeof generated
                              ? sizeof generated
                              : remain);
     rng.stream(generated, bytesGenerated);
     if(mode == CREATE) {
-      fwrite(generated, 1, bytesGenerated, fp);
-      if(ferror(fp))
-        fatal(errno, "%s", path);
+      size_t bytesWritten = fwrite(generated, 1, bytesGenerated, fp);
+      if(ferror(fp)) {
+       if(!entire || errno != ENOSPC)
+         fatal(errno, "%s", path);
+       remain -= bytesWritten;
+       break;
+      }
+      assert(bytesWritten == bytesGenerated);
     } else {
       size_t bytesRead = fread(input, 1, bytesGenerated, fp);
       if(ferror(fp))
         fatal(errno, "%s", path);
-      if(bytesRead < bytesGenerated)
-        fatal(0, "%s: truncated at %lld/%lld bytes",
-                path, (size - remain + bytesRead), size);
-      if(memcmp(generated, input, bytesGenerated)) {
-        for(size_t n = 0; n < bytesGenerated; ++n)
+      if(memcmp(generated, input, bytesRead)) {
+        for(size_t n = 0; n < bytesRead; ++n)
           if(generated[n] != input[n])
             fatal(0, "%s corrupted at %lld/%lld bytes (expected %d got %d)",
                     path, size - remain + n, size,
                     (unsigned char)generated[n], (unsigned char)input[n]);
       }
+      if(bytesRead < bytesGenerated) {
+       if(entire) {
+         assert(feof(fp));
+         remain -= bytesRead;
+         break;
+       }
+        fatal(0, "%s: truncated at %lld/%lld bytes",
+                path, (size - remain + bytesRead), size);
+      }
     }
     remain -= bytesGenerated;
   }
-  if(mode == VERIFY && getc(fp) != EOF)
+  if(mode == VERIFY && !entire && getc(fp) != EOF)
     fatal(0, "%s: extended beyond %lld bytes",
             path, size);
   if(mode == CREATE && flush) {
@@ -196,5 +242,13 @@ int main(int argc, char **argv) {
   }
   if(fclose(fp) < 0)
     fatal(errno, "%s", path);
-  return 0;
+  long long done = size - remain;
+  if(show) {
+    printf("%lld bytes (%lldM, %lldG) %s\n",
+          done, done >> 20, done >> 30,
+          show);
+    if(ferror(stdout) || fflush(stdout))
+      fatal(errno, "flush stdout");
+  }
+  return done;
 }