about summary refs log tree commit diff
path: root/src/download-via-ssh/download-via-ssh.cc
blob: 5c4756372afa61c9a3784a615fbbe238320586b6 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
#include "shared.hh"
#include "util.hh"
#include "serialise.hh"
#include "archive.hh"
#include "affinity.hh"
#include "globals.hh"
#include "serve-protocol.hh"
#include "worker-protocol.hh"
#include "store-api.hh"

#include <iostream>
#include <cstdlib>
#include <unistd.h>

using namespace nix;

// !!! TODO:
// * Respect more than the first host
// * use a database
// * show progress


static std::pair<FdSink, FdSource> connect(const string & conn)
{
    Pipe to, from;
    to.create();
    from.create();
    startProcess([&]() {
        if (dup2(to.readSide, STDIN_FILENO) == -1)
            throw SysError("dupping stdin");
        if (dup2(from.writeSide, STDOUT_FILENO) == -1)
            throw SysError("dupping stdout");
        execlp("ssh", "ssh", "-x", "-T", conn.c_str(), "nix-store --serve", NULL);
        throw SysError("executing ssh");
    });
    // If child exits unexpectedly, we'll EPIPE or EOF early.
    // If we exit unexpectedly, child will EPIPE or EOF early.
    // So no need to keep track of it.

    return std::pair<FdSink, FdSource>(to.writeSide.borrow(), from.readSide.borrow());
}


static void substitute(std::pair<FdSink, FdSource> & pipes, Path storePath, Path destPath)
{
    pipes.first << cmdDumpStorePath << storePath;
    pipes.first.flush();
    restorePath(destPath, pipes.second);
    std::cout << std::endl;
}


static void query(std::pair<FdSink, FdSource> & pipes)
{
    for (string line; getline(std::cin, line);) {
        Strings tokenized = tokenizeString<Strings>(line);
        string cmd = tokenized.front();
        tokenized.pop_front();
        if (cmd == "have") {
            pipes.first
                << cmdQueryValidPaths
                << 0 // don't lock
                << 0 // don't substitute
                << tokenized;
            pipes.first.flush();
            PathSet paths = readStrings<PathSet>(pipes.second);
            for (auto & i : paths)
                std::cout << i << std::endl;
        } else if (cmd == "info") {
            pipes.first << cmdQueryPathInfos << tokenized;
            pipes.first.flush();
            while (1) {
                Path path = readString(pipes.second);
                if (path.empty()) break;
                assertStorePath(path);
                std::cout << path << std::endl;
                string deriver = readString(pipes.second);
                if (!deriver.empty()) assertStorePath(deriver);
                std::cout << deriver << std::endl;
                PathSet references = readStorePaths<PathSet>(pipes.second);
                std::cout << references.size() << std::endl;
                for (auto & i : references)
                    std::cout << i << std::endl;
                std::cout << readLongLong(pipes.second) << std::endl;
                std::cout << readLongLong(pipes.second) << std::endl;
            }
        } else
            throw Error(format("unknown substituter query '%1%'") % cmd);
        std::cout << std::endl;
    }
}


int main(int argc, char * * argv)
{
    return handleExceptions(argv[0], [&]() {
        if (argc < 2)
            throw UsageError("download-via-ssh requires an argument");

        initNix();

        settings.update();

        if (settings.sshSubstituterHosts.empty())
            return;

        std::cout << std::endl;

        /* Pass on the location of the daemon client's SSH
           authentication socket. */
        string sshAuthSock = settings.get("ssh-auth-sock", string(""));
        if (sshAuthSock != "") setenv("SSH_AUTH_SOCK", sshAuthSock.c_str(), 1);

        string host = settings.sshSubstituterHosts.front();
        std::pair<FdSink, FdSource> pipes = connect(host);

        /* Exchange the greeting */
        pipes.first << SERVE_MAGIC_1;
        pipes.first.flush();
        unsigned int magic = readInt(pipes.second);
        if (magic != SERVE_MAGIC_2)
            throw Error("protocol mismatch");
        readInt(pipes.second); // Server version, unused for now
        pipes.first << SERVE_PROTOCOL_VERSION;
        pipes.first.flush();

        string arg = argv[1];
        if (arg == "--query")
            query(pipes);
        else if (arg == "--substitute") {
            if (argc != 4)
                throw UsageError("download-via-ssh: --substitute takes exactly two arguments");
            Path storePath = argv[2];
            Path destPath = argv[3];
            printError(format("downloading '%1%' via SSH from '%2%'...") % storePath % host);
            substitute(pipes, storePath, destPath);
        }
        else
            throw UsageError(format("download-via-ssh: unknown command '%1%'") % arg);
    });
}