Re-seed after fork lazily using pthread_atfork handler

MANIFEST

@@ -11,6 +11,7 @@ README
 standard.h
 t/00-system.t
 t/alphabets.t
+t/fork.t
 t/kernel-seeding.t
 t/no-mod-bias.t
 t/reference.t

Makefile.PL

@@ -12,7 +12,7 @@ my %args = (
     NAME              => 'Session::Token',
     VERSION_FROM      => 'lib/Session/Token.pm',
     PREREQ_PM         => $deps,
-    LIBS              => [''],
+    LIBS              => ['-lpthread'],
     DEFINE            => '',
     INC               => '-I.',
     OBJECT            => 'Token.o randport.o',

Token.xs

@@ -8,13 +8,20 @@
 #include <stdlib.h>
 #include <string.h>
 #include <assert.h>
+#include <signal.h>
+
+#ifndef _WIN32
+#include <pthread.h>
+#endif
 
 
 struct session_token_ctx {
   int mask;
   int count;
   int curr_word;
   int bytes_left_in_curr_word;
+  int reseed_on_fork;
+  int seed_generation;
   struct randctx isaac_ctx;
   char *alphabet;
   size_t alphabet_length;
@@ -23,6 +30,54 @@ struct session_token_ctx {
 
 typedef struct session_token_ctx * Session_Token;
 
+static volatile sig_atomic_t fork_generation = 0;
+
+#ifndef _WIN32
+static void session_token_atfork_child(void) {
+  fork_generation++;
+}
+#endif
+
+static void seed_ctx(struct session_token_ctx *ctx, char *seedp) {
+  memcpy(&ctx->isaac_ctx.randrsl, seedp, 1024);
+  randinit(&ctx->isaac_ctx, TRUE);
+  isaac(&ctx->isaac_ctx);
+
+  ctx->count = 0;
+  ctx->curr_word = 0;
+  ctx->bytes_left_in_curr_word = 0;
+  ctx->seed_generation = fork_generation;
+}
+
+static void reseed_ctx(struct session_token_ctx *ctx) {
+  dSP;
+  SV *seed;
+  char *seedp;
+  size_t len;
+  int count;
+
+  ENTER;
+  SAVETMPS;
+
+  PUSHMARK(SP);
+  count = call_pv("Session::Token::_get_seed", G_SCALAR);
+
+  SPAGAIN;
+
+  if (count != 1) croak("_get_seed didn't return a seed");
+
+  seed = POPs;
+  seedp = SvPV(seed, len);
+
+  if (len != 1024) croak("unexpected seed length: %lu", len);
+
+  seed_ctx(ctx, seedp);
+
+  PUTBACK;
+  FREETMPS;
+  LEAVE;
+}
+
 static inline int get_new_byte(struct session_token_ctx *ctx) {
   int output;
 
@@ -57,12 +112,19 @@ MODULE = Session::Token		PACKAGE = Session::Token
 
 PROTOTYPES: ENABLE
 
+BOOT:
+#ifndef _WIN32
+    if (pthread_atfork(NULL, NULL, session_token_atfork_child) != 0)
+      croak("unable to register pthread_atfork handler");
+#endif
+
 
 Session_Token
-_new_context(seed, alphabet, token_length)
+_new_context(seed, alphabet, token_length, reseed_on_fork)
         SV *seed
         SV *alphabet
         size_t token_length
+        int reseed_on_fork
     CODE:
         struct session_token_ctx *ctx;
         char *seedp;
@@ -80,9 +142,9 @@ _new_context(seed, alphabet, token_length)
         ctx = malloc(sizeof(struct session_token_ctx));
         memset(ctx, '\0', sizeof(struct session_token_ctx));
 
-        memcpy(&ctx->isaac_ctx.randrsl, seedp, 1024);
-        randinit(&ctx->isaac_ctx, TRUE);
-        isaac(&ctx->isaac_ctx);
+        seed_ctx(ctx, seedp);
+
+        ctx->reseed_on_fork = reseed_on_fork;
 
         ctx->alphabet_length = SvCUR(alphabet);
         ctx->alphabet = malloc(ctx->alphabet_length);
@@ -114,6 +176,9 @@ get(ctx)
         char *outputp;
         size_t i, curr;
 
+        if (ctx->reseed_on_fork && ctx->seed_generation != fork_generation)
+          reseed_ctx(ctx);
+
         output = newSVpvn("", 0);
         SvGROW(output, ctx->token_length + 1);
         SvCUR_set(output, ctx->token_length);

lib/Session/Token.pm

@@ -42,25 +42,8 @@ sub new {
     $seed = $args{seed};
   }
 
-  if (!defined $seed) {
-    if ($is_windows) {
-      my $windows_rng_source = Crypt::Random::Source::Strong::Win32->new;
-      $seed = $windows_rng_source->get(1024);
-      die "Win32 RNG source didn't provide 1024 bytes" unless length($seed) == 1024;
-    } else {
-      my ($fh, $err1, $err2);
-
-      open($fh, '<:raw', '/dev/urandom') || ($err1 = $!);
-      open($fh, '<:raw', '/dev/arandom') || ($err2 = $!)
-        unless defined $fh;
-
-      if (!defined $fh) {
-        croak "unable to open /dev/urandom ($err1) or /dev/arandom ($err2)";
-      }
-
-      sysread($fh, $seed, 1024) == 1024 || croak "unable to read from random device: $!";
-    }
-  }
+  $seed = _get_seed()
+    if !defined $seed;
 
 
   ## Init alphabet
@@ -90,7 +73,34 @@ sub new {
   }
 
 
-  return _new_context($seed, $alphabet, $token_length);
+  ## Don't re-seed after fork if using a custom seed
+
+  return _new_context($seed, $alphabet, $token_length, defined $args{seed} ? 0 : 1);
+}
+
+
+sub _get_seed {
+  my $seed;
+
+  if ($is_windows) {
+    my $windows_rng_source = Crypt::Random::Source::Strong::Win32->new;
+    $seed = $windows_rng_source->get(1024);
+    die "Win32 RNG source didn't provide 1024 bytes" unless length($seed) == 1024;
+  } else {
+    my ($fh, $err1, $err2);
+
+    open($fh, '<:raw', '/dev/urandom') || ($err1 = $!);
+    open($fh, '<:raw', '/dev/arandom') || ($err2 = $!)
+      unless defined $fh;
+
+    if (!defined $fh) {
+      croak "unable to open /dev/urandom ($err1) or /dev/arandom ($err2)";
+    }
+
+    sysread($fh, $seed, 1024) == 1024 || croak "unable to read from random device: $!";
+  }
+
+  return $seed;
 }
 
 
@@ -144,9 +154,9 @@ When a Session::Token object is created, 1024 bytes are read from C</dev/urandom
 
 Once a generator is created, you can repeatedly call the C<get> method on the generator object and it will return a new token each time.
 
-B<IMPORTANT>: If your application calls C<fork>, make sure that any generators are re-created in one of the processes after the fork since forking will duplicate the generator state and both parent and child processes will go on to produce identical tokens (just like perl's L<rand> after it is seeded).
+If your application calls C<fork>, generators will detect this and re-seed themselves from the kernel before producing another token, so parent and child processes will not go on to produce identical tokens (as happens with perl's L<rand> after it is seeded).
 
-After the generator context is created, no system calls are used to generate tokens. This is one way that Session::Token helps with efficiency. However, this is only important for certain use cases (generally not web sessions).
+After the generator context is created, no system calls are used to generate tokens (except to re-seed after a fork). This is one way that Session::Token helps with efficiency. However, this is only important for certain use cases (generally not web sessions).
 
 ISAAC is a cryptographically secure PRNG that improves on the well-known RC4 algorithm in some important areas. For instance, it doesn't have short cycles or initial bias like RC4 does. A theoretical shortest possible cycle in ISAAC is C<2**40>, although no cycles this short have ever been found (and probably don't exist at all). On average, ISAAC cycles are C<2**8295>.
 
@@ -376,6 +386,8 @@ This is done in the test-suite to compare against Jenkins' reference ISAAC outpu
 
 One valid reason for manually seeding is if you have some reason to believe that there isn't enough entropy in your kernel's randomness pool and therefore you don't trust C</dev/urandom>. In this case you should acquire your own seed data from somewhere trustworthy (maybe C</dev/random> or a previously stored trusted seed).
 
+Generators created with a custom seed are not re-seeded after a C<fork>.
+
 
 
 
@@ -416,8 +428,6 @@ It supports all the options of this module via command line parameters, and mult
 
 Should check for biased alphabets and print warnings.
 
-Would be cool if it could detect forks and warn or re-seed in the child process (without incurring C<getpid> overhead).
-
 There is currently no way to extract the seed from a Session::Token object. Note when implementing this: The saved seed must either store the current state of the ISAAC round as well as the 1024 byte C<randsl> array or else do some kind of minimum fast forwarding in order to protect against a partially duplicated output-stream bug.
 
 Doesn't work on perl 5.6 and below due to the use of C<:raw> (thanks CPAN testers). It could probably use C<binmode> instead, but meh.

t/fork.t

@@ -0,0 +1,122 @@
+use Session::Token;
+
+## The point of this test is to verify that generators detect forks
+## and re-seed themselves in the child process so that parent and
+## child don't produce identical token streams. Generators with
+## custom seeds are deterministic and must not be re-seeded.
+
+use strict;
+
+use Test::More;
+
+plan skip_all => 'fork test requires a real fork()'
+  if $^O =~ /mswin/i;
+
+plan tests => 6;
+
+
+sub tokens_from_child {
+  my ($gen, $n) = @_;
+
+  pipe(my $rd, my $wr) || die "unable to pipe: $!";
+
+  my $pid = fork;
+  die "unable to fork: $!" if !defined $pid;
+
+  if ($pid == 0) {
+    close($rd);
+    print $wr $gen->get, "\n" for 1 .. $n;
+    close($wr);
+    exit 0;
+  }
+
+  close($wr);
+  my @tokens = <$rd>;
+  close($rd);
+  waitpid($pid, 0);
+
+  chomp @tokens;
+  die "child produced " . scalar(@tokens) . " tokens instead of $n" if @tokens != $n;
+
+  return @tokens;
+}
+
+sub all_unique {
+  my %seen;
+  $seen{$_} = 1 for @_;
+  return scalar(keys %seen) == scalar(@_);
+}
+
+
+{
+  my $gen = Session::Token->new;
+  $gen->get;
+
+  my @child = tokens_from_child($gen, 20);
+  my @parent = map { $gen->get } 1 .. 20;
+
+  ok(all_unique(@child, @parent), "child re-seeds after fork");
+}
+
+
+{
+  my $gen = Session::Token->new;
+
+  my @child = tokens_from_child($gen, 20);
+  my @parent = map { $gen->get } 1 .. 20;
+
+  ok(all_unique(@child, @parent), "child re-seeds when generator was never used before fork");
+}
+
+
+{
+  my $gen = Session::Token->new;
+
+  my @child1 = tokens_from_child($gen, 20);
+  my @child2 = tokens_from_child($gen, 20);
+  my @parent = map { $gen->get } 1 .. 20;
+
+  ok(all_unique(@child1, @child2, @parent), "sequential children re-seed independently");
+}
+
+
+{
+  my $gen = Session::Token->new;
+
+  pipe(my $rd, my $wr) || die "unable to pipe: $!";
+
+  my $pid = fork;
+  die "unable to fork: $!" if !defined $pid;
+
+  if ($pid == 0) {
+    close($rd);
+    my @grandchild = tokens_from_child($gen, 20);
+    print $wr "$_\n" for @grandchild, map { $gen->get } 1 .. 20;
+    close($wr);
+    exit 0;
+  }
+
+  close($wr);
+  chomp(my @lines = <$rd>);
+  close($rd);
+  waitpid($pid, 0);
+
+  is(scalar(@lines), 40, "received all tokens from forked processes");
+
+  my @parent = map { $gen->get } 1 .. 20;
+
+  ok(all_unique(@lines, @parent), "grandchild re-seeds too");
+}
+
+
+{
+  my $gen = Session::Token->new(seed => "\x00" x 1024);
+  my $mirror = Session::Token->new(seed => "\x00" x 1024);
+
+  $gen->get;
+  $mirror->get;
+
+  my ($child_token) = tokens_from_child($gen, 1);
+
+  is($child_token, $mirror->get, "custom seeded generator is not re-seeded after fork");
+}