#! /usr/bin/perl

use strict;
use warnings;

use Socket;
use Time::HiRes qw(gettimeofday tv_interval);

my $verbose = undef;
# The number of connections seems to affect the chance of hitting
# the race. So might need tuning on different machines.
my $num_connected = 20;
my $port = 2345;
my $proto = getprotobyname("tcp");
my $remote = "127.0.0.1";
my $iaddr   = inet_aton($remote)       || die "no host: $remote";
my $paddr   = sockaddr_in($port, $iaddr);

print <<END;
This script tries to demonstrate a race/bug between TCP connect() in a client
process and SIGKILL on the server process. When the timing is just right, the
connect() can succeed while the killed process is closing existing connections,
but the resulting extra connection does not get any notification that it has
been closed on the server, so poll()/select() on the client hangs indefinitely.

Occasionally we also observe that the extra connect() completes only after
(almost) exactly 1 second, after which it returns failure to connect.

END
for (my $i = 1; ; ++$i) {
  print "Iteration: $i\n"
      if (($i % 100) == 0);
  test(1 + int(rand($num_connected)), $i);
  select(undef, undef, undef, 0.01);
}


sub test {
  my ($N, $iter) = @_;

  # Spawn a server process.
  my $pid = open(my $child, '-|');
  die unless defined($pid);
  if (!$pid) {
    server($N);
    exit(0);  # NOTREACHED
  }

  # Wait for the server process to setup the listening socket.
  sysread($child, my $dummy, 1);

  my $socks = [];
  my $rin = my $win = my $ein = '';
  for my $i (1..$N) {
    socket($socks->[$i], PF_INET, SOCK_STREAM, $proto)
        or die "socket: $!";
    connect($socks->[$i], $paddr)
        or die "connect: $!";
    vec($rin, fileno($socks->[$i]), 1) = 1;
    sysread($socks->[$i], my $dummy2, 1);
    #print "Connected! sock=", fileno($socks->[$i]), "\n";
  }
  $ein = $rin;
  print "select() vector: ", unpack('H*', $rin), "\n"
      if $verbose;

  socket(my $extra_sock, PF_INET, SOCK_STREAM, $proto) or die "extra_sock: $!";

  # Wait for the server to die, which causes connection sockets to close and notify us.
  my $t0 = [gettimeofday];
  my $res = select(my $rout = $rin, my $wout = $win, my $eout = $ein, undef);
  # Try to get in an extra connection while the server process is still getting
  # killed. This is the connection that can show the race/bug, where the
  # connection gets established, but there is never any notification that
  # the connection was closed at the other end.
  my $t1 = [gettimeofday];
  my $res2 = connect($extra_sock, $paddr);
  my $t2 = [gettimeofday];
  print "select() returns: $res out=", unpack('H*', $rout), " err=", unpack('H*', $eout), "\n"
      if $verbose;
  my $elapsed1 = tv_interval($t0, $t1);
  my $elapsed2 = tv_interval($t1, $t2);
  print "Oops, select() took $elapsed1 seconds! (found=$res)\n"
      if $elapsed1 > 0.5;
  # Curiously, the extra connection sometimes takes almost exactly 1 second
  # to fail.
  print "Oops, connect() took $elapsed2 seconds! (connect=", ($res2 ? 'Yes' : 'No'), ")\n"
      if $elapsed2 > 0.5;
  if ($res2) {
    print "Oh! Got an extra connection: fd=", fileno($extra_sock), "\n";
    my $rin2 = my $win2 = my $ein2 = '';
    vec($rin2, fileno($extra_sock), 1) = 1;
    $ein2 = $rin2;
    my $res3 = select(my $rout2 = $rin2, my $wout2 = $win2, my $eout2 = $ein2, 10);
    print "extra select() returns: $res3 out=",
            unpack('H*', $rout2), " err=", unpack('H*', $eout2), "\n"
        if $verbose;
    if (!$res3) {
      print STDERR "AHA! select() on extra connection timed out on iteration $iter!\n";
      my ($extra_port, $extra_host) = sockaddr_in(getsockname($extra_sock));
      print STDERR "Extra connection fd=", fileno($extra_sock),
          " port=", $extra_port, "\n";
      exit(1);
    }
  } else {
    print "Bummer, no extra connection this time...\n"
        if $verbose;
  }

  close $child;
  close $socks->[$_] for (1..$N);
  close $extra_sock;
}

# Server.
sub server {
  my ($N) = @_;
  socket(my $server, PF_INET, SOCK_STREAM, $proto) || die "socket: $!";
  setsockopt($server, SOL_SOCKET, SO_REUSEADDR, pack("l", 1))
                                                || die "setsockopt: $!";
  bind($server, sockaddr_in($port, INADDR_ANY)) || die "bind: $!";
  listen($server, SOMAXCONN)                    || die "listen: $!";

  # Signal parent that socket is ready!
  syswrite(STDOUT, '.', 1);

  # First accept a bunch of connections.
  my $clis = [];
  for my $i (1..$N) {
    my $res = accept($clis->[$i], $server);
    die unless $res;
    syswrite($clis->[$i], '!', 1) or die;
  }
  # Then SIGKILL ourselves, see if we can trigger a client managing to get a
  # connection in during the close-of-sockets.
  select(undef, undef, undef, 0.02);
  print STDERR "SIGKILL self\n"
      if $verbose;
  kill 9, $$;
  sleep(10) while (1);
}