Naut/apps/leech/main.c

/* naut_leech — Phase 3 gate: download a torrent from a single peer and write a
 * byte-correct, hash-verified file to disk.
 *
 * Blocking-socket driver around the sans-IO peer codec + download engine. The
 * point of this phase is protocol correctness and interop (it downloads from a
 * libtorrent seed in the integration test), not peak throughput — the io_uring
 * reactor that drives thousands of these comes in Phase 6.
 *
 *   usage: naut_leech [--mse] <file.torrent> <output-dir> <ip> <port>
 */
#include "naut/metainfo.h"
#include "naut/storage.h"
#include "naut/piece.h"
#include "naut/peer.h"
#include "naut/mse.h"
#include "naut/log.h"

#include <arpa/inet.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>

#define PIPELINE_DEPTH 512          /* outstanding requests (~8 MiB in flight) */

static double now(void) {
    struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t);
    return t.tv_sec + t.tv_nsec * 1e-9;
}

static uint8_t *slurp(const char *path, size_t *len) {
    FILE *f = fopen(path, "rb");
    if (!f) { NAUT_ERROR("open %s: %s", path, strerror(errno)); return NULL; }
    fseek(f, 0, SEEK_END); long n = ftell(f); fseek(f, 0, SEEK_SET);
    uint8_t *b = malloc(n);
    if (fread(b, 1, n, f) != (size_t)n) { fclose(f); free(b); return NULL; }
    fclose(f); *len = (size_t)n; return b;
}

static int connect_peer(const char *ip, uint16_t port) {
    int fd = socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) return -1;
    struct sockaddr_in a; memset(&a, 0, sizeof a);
    a.sin_family = AF_INET; a.sin_port = htons(port);
    if (inet_pton(AF_INET, ip, &a.sin_addr) != 1) { close(fd); return -1; }
    if (connect(fd, (struct sockaddr *)&a, sizeof a) != 0) {
        NAUT_ERROR("connect %s:%u: %s", ip, port, strerror(errno));
        close(fd); return -1;
    }
    int one = 1; setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof one);
    return fd;
}

/* send up to PIPELINE_DEPTH outstanding requests */
static bool refill(int fd, naut_mse_stream *mse,
                   naut_download *d, int *outstanding) {
    uint32_t idx, begin, len;
    while (*outstanding < PIPELINE_DEPTH) {
        if (!naut_download_next_request(d, &idx, &begin, &len)) break;
        uint8_t req[17];
        naut_peer_msg_request(req, idx, begin, len);
        if (!naut_mse_send_all(fd, mse, req, sizeof req)) return false;
        (*outstanding)++;
    }
    return true;
}

int main(int argc, char **argv) {
    bool use_mse = argc > 1 && strcmp(argv[1], "--mse") == 0;
    int arg = use_mse ? 2 : 1;
    if (argc - arg != 4) {
        fprintf(stderr, "usage: %s [--mse] <file.torrent> <output-dir> <ip> <port>\n",
                argv[0]);
        return 2;
    }
    naut_log_set_level(NAUT_LOG_INFO);

    size_t tlen;
    uint8_t *tor = slurp(argv[arg], &tlen);
    if (!tor) return 1;
    naut_metainfo mi;
    if (naut_metainfo_parse(tor, tlen, &mi) != NAUT_OK) { NAUT_ERROR("bad torrent"); return 1; }
    free(tor);

    char hex[41]; naut_infohash_v1_hex(&mi, hex);
    NAUT_INFO("torrent '%s': %u pieces, %lld bytes, infohash %s",
              mi.name, mi.num_pieces, (long long)mi.total_length, hex);

    naut_err err;
    naut_storage_opts storage_opts = {
        .direct_io = getenv("NAUT_DIRECT_IO") != NULL,
        .preallocate = true,
    };
    naut_storage *st = naut_storage_open_opts(
        mi.files, mi.num_files, argv[arg + 1], &storage_opts, &err);
    if (!st) { NAUT_ERROR("storage: %s", naut_strerror(err)); return 1; }
    naut_download *d = naut_download_create(&mi, st);
    if (!d) return 1;

    int fd = connect_peer(argv[arg + 2], (uint16_t)atoi(argv[arg + 3]));
    if (fd < 0) return 1;

    /* handshake */
    uint8_t peerid[20]; memcpy(peerid, "-NT0001-", 8);
    for (int i = 8; i < 20; i++) peerid[i] = (uint8_t)(rand() & 0xff);
    uint8_t hs[NAUT_HANDSHAKE_LEN];
    naut_peer_handshake_build(hs, mi.infohash_v1, peerid, 0);
    naut_mse_stream mse = {0};
    uint8_t remote_hs[NAUT_HANDSHAKE_LEN];
    bool hs_done = false;
    if (use_mse) {
        naut_err mse_err = naut_mse_client_handshake(
            fd, mi.infohash_v1, peerid, 0, &mse, remote_hs);
        if (mse_err != NAUT_OK) {
            NAUT_ERROR("MSE handshake failed: %s", naut_strerror(mse_err));
            return 1;
        }
        hs_done = true;
        NAUT_INFO("MSE/RC4 peer transport established");
        uint8_t intr[5];
        naut_peer_msg_simple(intr, NAUT_MSG_INTERESTED);
        if (!naut_mse_send_all(fd, &mse, intr, sizeof intr)) {
            NAUT_ERROR("interested send failed");
            return 1;
        }
    } else if (!naut_mse_send_all(fd, &mse, hs, sizeof hs)) {
        NAUT_ERROR("handshake send failed");
        return 1;
    }

    /* recv buffer */
    size_t cap = 4u << 20, len = 0;
    uint8_t *buf = malloc(cap);
    bool unchoked = false;
    int outstanding = 0;
    double t0 = now();

    while (!naut_download_complete(d)) {
        if (len == cap) { cap *= 2; buf = realloc(buf, cap); }
        ssize_t r = naut_mse_recv(fd, &mse, buf + len, cap - len);
        if (r < 0) { NAUT_ERROR("recv: %s", strerror(errno)); break; }
        if (r == 0) { NAUT_ERROR("peer closed (%.1f%% done)",
                      100.0 * naut_download_pieces_done(d) / mi.num_pieces); break; }
        len += (size_t)r;

        size_t pos = 0;
        if (!hs_done) {
            if (len < NAUT_HANDSHAKE_LEN) continue;
            uint8_t ih[20], pid[20];
            if (!naut_peer_handshake_parse(buf, ih, pid, NULL) ||
                memcmp(ih, mi.infohash_v1, 20) != 0) {
                NAUT_ERROR("handshake mismatch"); break;
            }
            pos = NAUT_HANDSHAKE_LEN;
            hs_done = true;
            uint8_t intr[5]; naut_peer_msg_simple(intr, NAUT_MSG_INTERESTED);
            if (!naut_mse_send_all(fd, &mse, intr, 5)) break;
        }

        /* parse all complete messages */
        for (;;) {
            naut_msg m;
            int c = naut_peer_msg_parse(buf + pos, len - pos, &m);
            if (c == 0) break;
            if (c < 0) { NAUT_ERROR("protocol error"); goto done; }
            pos += (size_t)c;
            switch (m.type) {
            case NAUT_MSG_UNCHOKE: unchoked = true; break;
            case NAUT_MSG_CHOKE:   unchoked = false; break;
            case NAUT_MSG_PIECE: {
                outstanding--;
                bool pdone = false;
                naut_err e = naut_download_on_block(d, m.index, m.begin, m.payload,
                                                    (uint32_t)m.payload_len, &pdone);
                if (e != NAUT_OK) { NAUT_ERROR("block rejected: %s", naut_strerror(e)); goto done; }
                break;
            }
            default: break;  /* bitfield/have/keepalive/port: ignore for a seed */
            }
        }
        /* compact consumed bytes */
        memmove(buf, buf + pos, len - pos);
        len -= pos;

        if (unchoked && !refill(fd, &mse, d, &outstanding)) {
            NAUT_ERROR("request send failed"); break;
        }
    }
done:;
    double dt = now() - t0;
    bool ok = naut_download_complete(d);
    if (ok) {
        double mb = (double)mi.total_length / 1e6;
        NAUT_INFO("COMPLETE: %u/%u pieces, %.1f MB in %.2fs (%.1f MB/s), all SHA-1 verified",
                  naut_download_pieces_done(d), mi.num_pieces, mb, dt, mb / dt);
    } else {
        NAUT_ERROR("INCOMPLETE: %u/%u pieces", naut_download_pieces_done(d), mi.num_pieces);
    }

    naut_storage_sync(st);
    close(fd);
    naut_download_destroy(d);
    naut_storage_close(st);
    naut_metainfo_free(&mi);
    free(buf);
    return ok ? 0 : 1;
}