#!/usr/bin/perl
### blkrrecv.perl  -*- Perl -*-
## Receive a file block-by-block.

### Ivan Shmakov, 2017, 2019, 2020

## To the extent possible under law, the author(s) have dedicated
## all copyright and related and neighboring rights to this software
## to the public domain worldwide.  This software is distributed
## without any warranty.

## You should have received a copy of the CC0 Public Domain Dedication
## along with this software.  If not, see
## <http://creativecommons.org/publicdomain/zero/1.0/>.

### History:

## 0.4  2020-04-16
##      Fixed: use 1 as the exit code when terminated with a (caught)
##      signal; invoke stats () there as well.

## 0.3  2019-11-01
##      Changed the protocol to use an explicit (hello, welcome)
##      handshake.  Implemented --read-blocks (-N).  Do not die
##      when autoflush fails.

## 0.2  2017-06-24 14:07:19Z
##      (sfn.qyCHe92ZrrZ6JrCzy6xjgFXZ6g2e1tirk4TTXlKm6jc.perl)
##      Implemented debug mode toggle via SIGUSR2.  Invoke stats ()
##      also on die ().  Check filesize on receiving eof.

## 0.1  2017-06-23 21:50:21Z
##      (sfn._TFtGj1MMXGuxG5Plapa9yPedIognqAEpxQP4CG2tmA.perl)
##      Initial revision.

### Code:

use common::sense;
use English qw (-no_match_vars);

# require Data::Dump;
require Digest::SHA;
require Getopt::Long;
require IO::File;
# require IO::Handle;

our $dev_fd_prefix
    = ($ENV{"DEVFDPREFIX"} // "/dev/fd/");

sub open_file {
    my ($fn, $mode) = @_;
    ## .
    return ("GLOB" eq ref ($fn) ? IO::File->new_from_fd ($fn->fileno (), $mode)
            : ($fn eq "-") ? (index ($mode, ">") > 0 ? \*STDOUT : \*STDIN)
            : do {
                  $mode
                      //= "<";
                  my $l
                      = length ($dev_fd_prefix);
                  ## .
                  ((substr ($fn, 0, $l) eq $dev_fd_prefix)
                   ? IO::File->new_from_fd (substr ($fn, $l), $mode)
                   : IO::File->new ($fn, $mode));
            });
}

## main

Getopt::Long::Configure (qw (gnu_compat));

our ($debug_p);
my ($dry_run_p, $command_s)
    = (0);
my ($blk_z, $blk_start, $blk_count)
    = ();

my $parsable_p
    = Getopt::Long::GetOptions (q (block-size=i)    => \$blk_z,
                                q (debug!)      => \$debug_p,
                                "n|dry-run!"    => \$dry_run_p,
                                "e|execute=s"   => \$command_s,
                                "N|read-blocks=i"   => \$blk_count,
                                "j|skip-blocks=i"   => \$blk_start)
    or die ("Cannot parse command line arguments");

require Data::Dump
    if ($debug_p);

die ("Exactly one non-option argument expected")
    unless (1 == @ARGV);
my $fn
    = $ARGV[0];
my $fh
    = open_file ($fn, ($dry_run_p ? "r" : "r+"))
    or die ($fn, ": Cannot open destination file: ", $!);
$fh->binmode ()
    or die ($fn, ": Cannot set binary mode: ", $!);

my $d0
    = Digest::SHA->new (256);

sub exit_on_signal {
    warn ("N: Got SIG", $_[0], ", exiting\n");
    stats ();
    exit (1);
}

## FIXME: add handlers, not replace?
@SIG{qw {TERM QUIT INT}}
    = (\&exit_on_signal) x 3;

my ($seen, $requested, $received, $same, $filesize)
    = (0, 0, 0, 0);
sub stats {
    warn ("I: ", $seen, " block digests seen",
          " ", $requested, " blocks requested",
          " ", $received, " blocks received",
          " ", $same, " blocks found to be the same\n");
}
## FIXME: add handlers, not replace?
$SIG{USR1}
    = \&stats;
$SIG{"__DIE__"} = sub {
    stats ();
    die ($PROGRAM_NAME, ": ", @_);
};

my ($x_in, $x_out, $x_pid);
if (defined ($command_s)) {
    require IPC::Open2;
    $x_pid
        = IPC::Open2::open2 ($x_in, $x_out, $command_s);
} else {
    ($x_in, $x_out)
        = (\*STDIN, \*STDOUT);
}
## FIXME: add handlers, not replace?
$SIG{"USR2"} = sub {
    $debug_p
        = ! $debug_p;
    warn ("N: Debug mode is now ", $debug_p ? "on" : "off");
    require Data::Dump
        if ($debug_p);
    $x_out->flush ()
        or warn ("Cannot flush output: ", $!);
};

$x_in->binmode ()
    or die ("Cannot set binary mode: ", $!);
$x_out->autoflush ()
    or warn ("W: Cannot set autoflush mode: ", $!);
$x_out->print  ("hello -",
                (defined ($blk_z) ? (" size ", $blk_z) : ()),
                " digest sha256",
                (defined ($blk_start) ? (" start ", $blk_start) : ()),
                (defined ($blk_count) ? (" count ", $blk_count) : ()),
                "\n")
    or die ("Sending initial handshake failed: ", $!);

my ($expected, $welcome_seen_p, @requested)
    = ($blk_start, 0);
while ((! defined ($filesize) || @requested > 0)) {
    warn ("D: Expected blocks: ",
          join (", ", $expected, @requested), "\n")
        if ($debug_p);
    my $response
        = $x_in->getline ();
    warn ("D: Got: ", scalar (Data::Dump::dump ($response))) if ($debug_p);

    last
        unless (defined ($response));

    ## FIXME: a specific parameter order is currently expected
    my ($w_size, $w_start, $w_count,
        $block, $size, $sha, $last_p, $eof_p)
        =  ($response =~ m {
                ^ welcome \s+ - \s+ size \s+ (\d+)
                  \s+ digest \s+ sha256
                  \s+ start \s+ (\d+)
                  (?: \s+ count \s+ (\d+))?
                  \s* $
                | ^ block \s+ (\d+) \s+ size \s+ (\d+)
                  (?: \s+ sha256 \s+ ([0-9a-zA-Z+/]{43}))?
                  (?: \s+ (last))?
                  \s* $
                | ^ (eof) \s* $
            }xi)
        or die ("Unrecognized response");

    if (defined ($w_size)) {
        die ("Unanticipated second welcome")
            if ($welcome_seen_p);
        $welcome_seen_p
            = 1;
        $blk_z
            //= $w_size;
        $blk_count
            //= $w_count;
        die ("Unexpected block size",
             " (got ", $w_size, ", expected ", $blk_z, ")")
            unless ($w_size == $blk_z);
        die ("Unexpected start block",
             " (got ", $w_start, ", expected at least ", $blk_start, ")")
            unless (! defined ($blk_start) || $w_start >= $blk_start);
        $blk_start
            = $w_start;
        $expected
            = $blk_start
            if ($expected < $blk_start);
        ## FIXME: check also if welcome is missing count given in hello
        die ("Unexpected block count",
             " (got ", $w_count, ", expected up to ", $blk_count, ")")
            unless (! defined ($w_count) || $w_count <= $blk_count);
        $blk_count
            //= $w_count;
        next;
    }

    if ($eof_p) {
        $filesize
            = ($blk_z * $expected);
        warn ("I: Filesize: ", $filesize, "\n");
        ## NB: otherwise we may end up waiting for one another
        $x_out->close ();
        next;
    }

    die ("Unexpected block size",
         " (got ", $size, ", expected up to ", $blk_z, ")")
        unless ($size <= $blk_z);

    my $o
        = ($blk_z * $block);

    if (defined ($sha)) {
        die  ("Unexpected block SHA",
              " (got ", $block, ", expected ", $expected, ")")
            if ($block != $expected || defined ($filesize));
        ## FIXME: silently ignoring
        next
            if (defined ($blk_count) && $block >= $blk_start + $blk_count);
        ++$expected;
        ++$seen;
        if ($last_p) {
            $filesize
                = ($blk_z * $block + $size);
            warn ("I: Filesize: ", $filesize, "\n");
        } elsif ($size < $blk_z) {
            die ("Unexpected block size",
                 " (got ", $size, ", expected ", $blk_z, ")")
        }
        $fh->seek ($o, 0)
            or die ($fn, ": Cannot seek (to ", $o, "): ", $!);
        my $s;
        my $rd
            = $fh->read ($s, $size);
        die ($fn, ": Cannot read (", $size, " at ", $o, "): ", $!)
            unless (defined ($rd) && $rd > 0);
        if ($rd != $size) {
            warn ("W: ", $fn, ": Short read",
                  " (got ", $rd, " of ", $size, " at ", $o, "): ", $!);
        } elsif ($sha eq $d0->clone ()->add ($s)->b64digest ()) {
            ## Digests match; assume so does the payload.
            next;
        }
        push (@requested, $block);
        $x_out->print ("send ", $block, "\n");
        ++$requested;
        ## NB: otherwise we may end up waiting for one another
        $x_out->close ()
            if (defined ($filesize));
        next;
    }

    die  ("Unexpected block data",
          " (got ", $block, ", expected ", $requested[0], ")")
        unless (@requested > 0 && $requested[0] == $block);
    shift (@requested);

    my $recv;
    do {
        my $rd
            = $x_in->read ($recv, $size);
        die ("Cannot read (", $size, " at ", $o, "): ", $!)
            unless (defined ($rd) && $rd > 0);
        die  ("Short read",
              " (got ", $rd, " of ", $size, " at ", $o, "): ", $!)
            unless ($rd == $size);
    };

    $fh->seek ($o, 0)
        or die ($fn, ": Cannot seek (to ", $o, "): ", $!);
    my $s;
    my $rd
        = $fh->read ($s, $size);
    die ($fn, ": Cannot read (", $size, " at ", $o, "): ", $!)
        unless (defined ($rd) && $rd > 0);

    ++$received;
    if ($s eq $recv) {
        warn ("W: ", $fn, ": Got same data from remote",
              " (at ", $o, ")");
        ++$same;
    } elsif ($rd != $size) {
        warn ("W: ", $fn, ": Short read",
              " (got ", $rd, " of ", $size, " at ", $o, "): ", $!);
    }
    $fh->seek ($o, 0)
        or die ($fn, ": Cannot seek (to ", $o, "): ", $!);
    $fh->write ($recv, $size)
        unless ($dry_run_p);
}

stats ();
die ("Transfer incomplete\n")
    unless (defined ($filesize) && @requested < 1);

### Emacs trailer
## Local variables:
## coding: us-ascii
## fill-column: 72
## indent-tabs-mode: nil
## ispell-local-dictionary: "american"
## End:
### blkrrecv.perl ends here