lix-releng-staging/scripts/download-from-binary-cache.pl.in

#! @perl@ -w @perlFlags@

use DBI;
use File::Basename;
use IO::Select;
use Nix::Config;
use Nix::Store;
use WWW::Curl::Easy;
use WWW::Curl::Multi;
use strict;


my @binaryCacheUrls = map { s/\/+$//; $_ } split(/ /, ($ENV{"NIX_BINARY_CACHES"} || ""));

my ($dbh, $insertNAR, $queryNAR, $insertNegativeNAR, $queryNegativeNAR);
my %cacheIds;

my $curlm = WWW::Curl::Multi->new;
my $activeRequests = 0;
my $curlIdCount = 1;
my %curlHandles;
my $caBundle = $ENV{"CURL_CA_BUNDLE"} || $ENV{"OPENSSL_X509_CERT_FILE"};


sub addRequest {
    my ($url) = @_;
    
    my $curl = WWW::Curl::Easy->new;
    my $curlId = $curlIdCount++;
    $curlHandles{$curlId} = { handle => $curl, content => "" };

    $curl->setopt(CURLOPT_PRIVATE, $curlId);
    $curl->setopt(CURLOPT_URL, $url);
    $curl->setopt(CURLOPT_WRITEDATA, \$curlHandles{$curlId}->{content});
    $curl->setopt(CURLOPT_FOLLOWLOCATION, 1);
    $curl->setopt(CURLOPT_CAINFO, $caBundle) if defined $caBundle;

    $curlm->add_handle($curl);
    $activeRequests++;

    return $curlHandles{$curlId};
}


sub processRequests {
    while ($activeRequests) {
        my ($rfds, $wfds, $efds) = $curlm->fdset();
        #print STDERR "R = @{$rfds}, W = @{$wfds}, E = @{$efds}\n";

        # Sleep until we can read or write some data.
        if (scalar @{$rfds} + scalar @{$wfds} + scalar @{$efds} > 0) {
            IO::Select->select(IO::Select->new(@{$rfds}), IO::Select->new(@{$wfds}), IO::Select->new(@{$efds}), 0.1);
        }
        
        if ($curlm->perform() != $activeRequests) {
            while (my ($id, $result) = $curlm->info_read) {
                if ($id) {
                    my $handle = $curlHandles{$id}->{handle};
                    $curlHandles{$id}->{result} = $result;
                    $curlHandles{$id}->{httpStatus} = $handle->getinfo(CURLINFO_HTTP_CODE);
                    #print STDERR "\nRequest completed ($id, $result, $curlHandles{$id}->{httpStatus})\n";
                    $activeRequests--;
                    delete $curlHandles{$id}->{handle};
                }
            }
        }
    }
}


sub initCache {
    my $dbPath = "$Nix::Config::stateDir/binary-cache-v1.sqlite";

    # Open/create the database.
    $dbh = DBI->connect("dbi:SQLite:dbname=$dbPath", "", "")
        or die "cannot open database `$dbPath'";
    $dbh->{RaiseError} = 1;
    $dbh->{PrintError} = 0;

    $dbh->do("pragma synchronous = off"); # we can always reproduce the cache
    $dbh->do("pragma journal_mode = truncate");

    # Initialise the database schema, if necessary.
    $dbh->do(<<EOF);
        create table if not exists BinaryCaches (
            id        integer primary key autoincrement not null,
            url       text unique not null
        );
EOF
    
    $dbh->do(<<EOF);
        create table if not exists NARs (
            cache            integer not null,
            storePath        text not null,
            url              text not null,
            compression      text not null,
            fileHash         text,
            fileSize         integer,
            narHash          text,
            narSize          integer,
            refs             text,
            deriver          text,
            system           text,
            timestamp        integer not null,
            primary key (cache, storePath),
            foreign key (cache) references BinaryCaches(id) on delete cascade
        );
EOF

    $dbh->do(<<EOF);
        create table if not exists NegativeNARs (
            cache            integer not null,
            storePath        text not null,
            timestamp        integer not null,
            primary key (cache, storePath),
            foreign key (cache) references BinaryCaches(id) on delete cascade
        );
EOF

    $insertNAR = $dbh->prepare(
        "insert or replace into NARs(cache, storePath, url, compression, fileHash, fileSize, narHash, " .
        "narSize, refs, deriver, system, timestamp) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)") or die;

    $queryNAR = $dbh->prepare("select * from NARs where cache = ? and storePath = ?") or die;

    $insertNegativeNAR = $dbh->prepare(
        "insert or replace into NegativeNARs(cache, storePath, timestamp) values (?, ?, ?)") or die;

    $queryNegativeNAR = $dbh->prepare("select 1 from NegativeNARs where cache = ? and storePath = ?") or die;
}


sub getInfoFrom {
    my ($storePath, $pathHash, $binaryCacheUrl) = @_;

    my $cacheId = getCacheId($binaryCacheUrl);

    # Bail out if there is a negative cache entry.
    $queryNegativeNAR->execute($cacheId, basename($storePath));
    return undef if @{$queryNegativeNAR->fetchall_arrayref()} != 0;
    
    my $infoUrl = "$binaryCacheUrl/$pathHash.narinfo";
    print STDERR "checking $infoUrl...\n";
    my $request = addRequest($infoUrl);
    processRequests;

    if ($request->{result} != 0 || $request->{httpStatus} != 200) {
        if ($request->{httpStatus} != 404) {
            print STDERR "could not download ‘$infoUrl’ (" .
                ($request->{result} != 0 ? "Curl error $request->{result}" : "HTTP status $request->{httpStatus}") . ")\n";
        } else {
            $insertNegativeNAR->execute($cacheId, basename($storePath), time());
        }
        return undef;
    }
    
    my ($storePath2, $url, $fileHash, $fileSize, $narHash, $narSize, $deriver, $system);
    my $compression = "bzip2";
    my @refs;
    foreach my $line (split "\n", $request->{content}) {
        $line =~ /^(.*): (.*)$/ or return undef;
        if ($1 eq "StorePath") { $storePath2 = $2; }
        elsif ($1 eq "URL") { $url = $2; }
        elsif ($1 eq "Compression") { $compression = $2; }
        elsif ($1 eq "FileHash") { $fileHash = $2; }
        elsif ($1 eq "FileSize") { $fileSize = int($2); }
        elsif ($1 eq "NarHash") { $narHash = $2; }
        elsif ($1 eq "NarSize") { $narSize = int($2); }
        elsif ($1 eq "References") { @refs = split / /, $2; }
        elsif ($1 eq "Deriver") { $deriver = $2; }
        elsif ($1 eq "System") { $system = $2; }
    }
    return undef if $storePath ne $storePath2;
    if ($storePath ne $storePath2 || !defined $url || !defined $narHash) {
        print STDERR "bad NAR info file ‘$infoUrl’\n";
        return undef;
    }
    
    # Cache the result.
    $insertNAR->execute(
        $cacheId, basename($storePath), $url, $compression, $fileHash, $fileSize,
        $narHash, $narSize, join(" ", @refs), $deriver, $system, time());
    
    return
        { url => $url
        , compression => $compression
        , fileHash => $fileHash
        , fileSize => $fileSize
        , narHash => $narHash
        , narSize => $narSize
        , refs => [ @refs ]
        , deriver => $deriver
        , system => $system
        };
}


sub getCacheId {
    my ($binaryCacheUrl) = @_;
    
    my $cacheId = $cacheIds{$binaryCacheUrl};
    return $cacheId if defined $cacheId;
    
    # FIXME: not atomic.
    my @res = @{$dbh->selectcol_arrayref("select id from BinaryCaches where url = ?", {}, $binaryCacheUrl)};
    if (scalar @res == 1) {
        $cacheId = $res[0];
    } else {
        $dbh->do("insert into BinaryCaches(url) values (?)",
                 {}, $binaryCacheUrl);
        $cacheId = $dbh->last_insert_id("", "", "", "");
    }

    $cacheIds{$binaryCacheUrl} = $cacheId;
    return $cacheId;
}


sub cachedGetInfoFrom {
    my ($storePath, $pathHash, $binaryCacheUrl) = @_;

    $queryNAR->execute(getCacheId($binaryCacheUrl), basename($storePath));
    my $res = $queryNAR->fetchrow_hashref();
    return undef unless defined $res;
    
    return 
        { url => $res->{url}
        , compression => $res->{compression}
        , fileHash => $res->{fileHash}
        , fileSize => $res->{fileSize}
        , narHash => $res->{narHash}
        , narSize => $res->{narSize}
        , refs => [ split " ", $res->{refs} ]
        , deriver => $res->{deriver}
        } if defined $res;
}


sub getInfo {
    my ($storePath) = @_;

    my $pathHash = substr(basename($storePath), 0, 32);

    # First look if we have cached info for one of the URLs.
    foreach my $binaryCacheUrl (@binaryCacheUrls) {
        my $info = cachedGetInfoFrom($storePath, $pathHash, $binaryCacheUrl);
        return $info if defined $info;
    }

    # No, so do an HTTP request until we get a hit.
    foreach my $binaryCacheUrl (@binaryCacheUrls) {
        my $info = getInfoFrom($storePath, $pathHash, $binaryCacheUrl);
        return $info if defined $info;
    }

    return undef;
}


sub downloadBinary {
    my ($storePath) = @_;
    
    my $pathHash = substr(basename($storePath), 0, 32);

    cache: foreach my $binaryCacheUrl (@binaryCacheUrls) {
        my $info = cachedGetInfoFrom($storePath, $pathHash, $binaryCacheUrl);
        $info = getInfoFrom($storePath, $pathHash, $binaryCacheUrl) unless defined $info;
        if (defined $info) {
            my $decompressor;
            if ($info->{compression} eq "bzip2") { $decompressor = "$Nix::Config::bzip2 -d"; }
            elsif ($info->{compression} eq "xz") { $decompressor = "$Nix::Config::xz -d"; }
            else {
                print STDERR "unknown compression method ‘$info->{compression}’\n";
                next;
            }
            print STDERR "\n*** Downloading ‘$info->{url}’ into ‘$storePath’...\n";
            if (system("$Nix::Config::curl --fail --location $binaryCacheUrl/$info->{url} | $decompressor | $Nix::Config::binDir/nix-store --restore $storePath") != 0) {
                die "download of `$info->{url}' failed" . ($! ? ": $!" : "") . "\n" unless $? == 0;
                next;
            }
            # The hash in the manifest can be either in base-16 or
            # base-32.  Handle both.
            $info->{narHash} =~ /^sha256:(.*)$/ or die "invalid hash";
            my $hash = $1;
            my $hash2 = hashPath("sha256", 1, $storePath);
            die "hash mismatch in downloaded path ‘$storePath’; expected $hash, got $hash2\n"
                if $hash ne $hash2;
            print STDERR "\n";
            return 1;
        }
    }

    return 0;
}


initCache();


if ($ARGV[0] eq "--query") {

    while (<STDIN>) {
        my $cmd = $_; chomp $cmd;

        if ($cmd eq "have") {
            my $storePath = <STDIN>; chomp $storePath;
            # FIXME: want to give correct info here, but it's too slow.
            #print "0\n";
            my $info = getInfo($storePath);
            if (defined $info) { print "1\n"; } else { print "0\n"; }
        }

        elsif ($cmd eq "info") {
            my $storePath = <STDIN>; chomp $storePath;
            my $info = getInfo($storePath);
            if (defined $info) {
                print "1\n";
                print $info->{deriver} ? "$Nix::Config::storeDir/$info->{deriver}" : "", "\n";
                print scalar @{$info->{refs}}, "\n";
                print "$Nix::Config::storeDir/$_\n" foreach @{$info->{refs}};
                print $info->{fileSize} || 0, "\n";
                print $info->{narSize} || 0, "\n";
            } else {
                print "0\n";
            }
        }

        else { die "unknown command `$cmd'"; }

        flush STDOUT;
    }

}

elsif ($ARGV[0] eq "--substitute") {
    my $storePath = $ARGV[1] or die;
    if (!downloadBinary($storePath)) {
        print STDERR "could not download ‘$storePath’ from any binary cache\n";
    }
}

else {
    die;
}