Naut/tests/unit/test_picker.c
2026-06-17 02:23:15 -04:00

196 lines
8 KiB
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);
/* The picker should not open the entire torrent at once. A large swarm can
* keep many requests in flight, but new-piece fanout is bounded so the
* piece map does not show most pieces "downloading" while few verify. */
enum { CAP_NP = 80 };
uint64_t cap_total = (uint64_t)CAP_NP * NAUT_BLOCK;
uint8_t *cap_hashes = malloc(CAP_NP * NAUT_SHA1_LEN);
CHECK(cap_hashes != NULL);
for (int p = 0; p < CAP_NP; p++)
naut_sha1(data, NAUT_BLOCK, cap_hashes + p * NAUT_SHA1_LEN);
naut_file cap_file[1] = { { (char *)"cap.bin", (int64_t)cap_total } };
naut_metainfo cap_mi; memset(&cap_mi, 0, sizeof cap_mi);
cap_mi.has_v1 = true; cap_mi.num_pieces = CAP_NP;
cap_mi.piece_length = NAUT_BLOCK; cap_mi.total_length = cap_total;
cap_mi.piece_hashes = cap_hashes; cap_mi.files = cap_file;
cap_mi.num_files = 1; cap_mi.name = (char *)"cap";
d = naut_download_create(&cap_mi, st);
CHECK(d != NULL);
naut_bitfield all;
naut_bitfield_init(&all, CAP_NP);
for (int p = 0; p < CAP_NP; p++) naut_bitfield_set(&all, p);
naut_download_add_bitfield(d, &all);
int opened = 0;
while (naut_download_pick(d, &all, &idx, &begin, &len))
opened++;
CHECK(opened > 0);
CHECK(opened < CAP_NP);
naut_bitfield_free(&all);
naut_download_destroy(d);
free(cap_hashes);
/* 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();
}