Naut/tests/unit/test_picker.c

/* Exercises the multi-peer swarm logic with a synthetic torrent (real SHA-1
 * piece hashes, so verification actually runs): rarest-first selection, two
 * peers collaborating on one piece, duplicate-block dedup, unrequest, endgame,
 * and a full multi-peer completion. */
#include "naut/piece.h"
#include "naut/metainfo.h"
#include "naut/storage.h"
#include "naut/hash.h"
#include "naut/bitfield.h"
#include "test.h"
#include <stdlib.h>
#include <string.h>

#define NP   10
#define PLEN 32768            /* 2 blocks per piece */
#define TOTAL ((uint64_t)NP * PLEN)

typedef struct {
    uint32_t piece[4], begin[4];
    size_t count;
} active_requests;

static bool is_active(void *ctx, uint32_t piece, uint32_t begin) {
    active_requests *a = ctx;
    for (size_t i = 0; i < a->count; i++)
        if (a->piece[i] == piece && a->begin[i] == begin) return true;
    return false;
}

static void add_active(active_requests *a, uint32_t piece, uint32_t begin) {
    a->piece[a->count] = piece;
    a->begin[a->count] = begin;
    a->count++;
}

int main(void) {
    uint8_t *data = malloc(TOTAL);
    for (uint64_t i = 0; i < TOTAL; i++) data[i] = (uint8_t)(i * 1103515245u + 12345u);
    uint8_t hashes[NP * NAUT_SHA1_LEN];
    for (int p = 0; p < NP; p++) naut_sha1(data + (uint64_t)p * PLEN, PLEN, hashes + p * NAUT_SHA1_LEN);

    naut_file files[1] = { { (char *)"data.bin", (int64_t)TOTAL } };
    naut_metainfo mi; memset(&mi, 0, sizeof mi);
    mi.has_v1 = true; mi.num_pieces = NP; mi.piece_length = PLEN; mi.total_length = TOTAL;
    mi.piece_hashes = hashes; mi.files = files; mi.num_files = 1; mi.name = (char *)"t";

    char tmpl[] = "/tmp/naut_pick_XXXXXX"; char *root = mkdtemp(tmpl);
    naut_err err;
    naut_storage *st = naut_storage_open(files, 1, root, &err);
    CHECK(st && err == NAUT_OK);
    naut_download *d = naut_download_create(&mi, st);
    CHECK(d != NULL);

    /* peer A has pieces 0..8, peer B has all 0..9 -> piece 9 is rarest (avail 1) */
    naut_bitfield ha, hb;
    naut_bitfield_init(&ha, NP); naut_bitfield_init(&hb, NP);
    for (int p = 0; p < NP; p++) { naut_bitfield_set(&hb, p); if (p < 9) naut_bitfield_set(&ha, p); }
    naut_download_add_bitfield(d, &ha);
    naut_download_add_bitfield(d, &hb);

    /* rarest-first: B's first pick must be the unique piece 9 */
    uint32_t idx, begin, len;
    CHECK(naut_download_pick(d, &hb, &idx, &begin, &len));
    CHECK_EQ(idx, 9);
    CHECK_EQ(begin, 0);

    /* collaboration: next pick for B finishes piece 9's second block */
    CHECK(naut_download_pick(d, &hb, &idx, &begin, &len));
    CHECK(idx == 9 && begin == NAUT_BLOCK);
    /* both blocks of piece 9 now requested; A (lacks 9) gets a different piece */
    uint32_t ia, ba, la;
    CHECK(naut_download_pick(d, &ha, &ia, &ba, &la));
    CHECK(ia != 9);

    /* unrequest releases a block for re-pick */
    naut_download_unrequest(d, ia, ba);
    uint32_t ia2, ba2, la2;
    CHECK(naut_download_pick(d, &ha, &ia2, &ba2, &la2));
    CHECK(ia2 == ia && ba2 == ba);

    /* duplicate block is ignored */
    bool done = false;
    uint64_t g9 = (uint64_t)9 * PLEN;
    CHECK(naut_download_on_block(d, 9, 0, data + g9, NAUT_BLOCK, &done) == NAUT_OK);
    CHECK(naut_download_on_block(d, 9, 0, data + g9, NAUT_BLOCK, &done) == NAUT_OK); /* dup */
    CHECK(!done);

    /* endgame flag is off this early (20 blocks, only a couple received) */
    CHECK(!naut_download_in_endgame(d));

    naut_download_destroy(d);

    /* full multi-peer download: alternate peers, all pieces verify */
    d = naut_download_create(&mi, st);
    naut_download_add_bitfield(d, &hb);   /* one peer that has everything */
    int guard = 0;
    while (!naut_download_complete(d) && guard++ < 10000) {
        if (!naut_download_pick(d, &hb, &idx, &begin, &len)) break;
        uint64_t g = (uint64_t)idx * PLEN + begin;
        CHECK(naut_download_on_block(d, idx, begin, data + g, len, &done) == NAUT_OK);
    }
    CHECK(naut_download_complete(d));
    CHECK(naut_download_in_endgame(d));   /* must have passed through endgame near the end */
    CHECK_EQ((long long)naut_download_bytes_done(d), (long long)TOTAL);

    naut_download_destroy(d);

    /* Endgame races each block on at most two distinct peers. Releasing one
     * peer's copy must not erase the other peer's outstanding request. */
    uint8_t small_hash[NAUT_SHA1_LEN];
    naut_sha1(data, PLEN, small_hash);
    naut_file small_file[1] = { { (char *)"small.bin", PLEN } };
    naut_metainfo small; memset(&small, 0, sizeof small);
    small.has_v1 = true; small.num_pieces = 1; small.piece_length = PLEN;
    small.total_length = PLEN; small.piece_hashes = small_hash;
    small.files = small_file; small.num_files = 1; small.name = (char *)"small";
    d = naut_download_create(&small, st);
    CHECK(d != NULL);
    naut_bitfield one;
    naut_bitfield_init(&one, 1);
    naut_bitfield_set(&one, 0);
    naut_download_add_bitfield(d, &one);
    naut_download_add_bitfield(d, &one);

    active_requests pa = {0}, pb = {0};
    CHECK(naut_download_pick_for_peer(d, &one, is_active, &pa,
                                      &idx, &begin, &len));
    add_active(&pa, idx, begin);
    CHECK(naut_download_pick_for_peer(d, &one, is_active, &pa,
                                      &idx, &begin, &len));
    add_active(&pa, idx, begin);
    CHECK(!naut_download_pick_for_peer(d, &one, is_active, &pa,
                                       &idx, &begin, &len));

    CHECK(naut_download_pick_for_peer(d, &one, is_active, &pb,
                                      &idx, &begin, &len));
    add_active(&pb, idx, begin);
    CHECK(naut_download_pick_for_peer(d, &one, is_active, &pb,
                                      &idx, &begin, &len));
    add_active(&pb, idx, begin);
    CHECK(!naut_download_pick_for_peer(d, &one, is_active, &pb,
                                       &idx, &begin, &len));

    naut_download_unrequest(d, pa.piece[0], pa.begin[0]);
    CHECK(!naut_download_pick_for_peer(d, &one, is_active, &pb,
                                       &idx, &begin, &len));
    pa.piece[0] = pa.piece[1]; pa.begin[0] = pa.begin[1]; pa.count = 1;
    CHECK(naut_download_pick_for_peer(d, &one, is_active, &pa,
                                      &idx, &begin, &len));
    CHECK_EQ(begin, 0);

    naut_bitfield_free(&one);
    naut_storage_sync(st);
    /* on-disk bytes match the source */
    uint8_t *rb = malloc(TOTAL);
    CHECK(naut_storage_read(st, 0, rb, TOTAL) == NAUT_OK);
    CHECK(memcmp(rb, data, TOTAL) == 0);
    free(rb);

    naut_download_destroy(d);
    naut_storage_close(st);
    naut_bitfield_free(&ha); naut_bitfield_free(&hb);
    free(data);
    char cmd[256]; snprintf(cmd, sizeof cmd, "rm -rf '%s'", root); if (system(cmd)) {}
    TEST_MAIN_END();
}