aria2-rust-pro/crates/aria2-rust-pro-tests/test_support/support.rs

#![doc(hidden)]
#![expect(
    dead_code,
    unreachable_pub,
    reason = "shared integration-test support intentionally exposes a superset of helpers because each integration suite only consumes part of it"
)]

use aria2_rust_pro_cli as _;
use aria2_rust_pro_compat as _;
use aria2_rust_pro_core as _;
use aria2_rust_pro_storage as _;
use aria2_rust_pro_tests as _;
use criterion as _;

use std::{
    collections::BTreeMap,
    io::{Read, Write},
    net::TcpListener,
    path::Path,
    sync::Mutex,
    thread,
};

use aria2_rust_pro_protocol::{
    DhtMessageModel, DhtNodeModel, DhtTransport, PeerWireTransportConnector,
    PeerWireTransportRequest, PeerWireTransportResponse, ReqwestTrackerTransport,
    TorrentMessageModel,
    torrent::{PeerWireHandshakeModel, PeerWireMessageKind},
    transport::{TransportEndpoint, TransportError, TransportScheme},
};
use aria2_rust_pro_rpc::{
    InProcessRpcDispatcher, JsonRpcRequest, JsonRpcResponse, RpcMeta, RpcMethod, RpcValue,
};

pub const BT_TORRENT_FIXTURE_BASE64: &str = "ZDg6YW5ub3VuY2UzNTpodHRwOi8vdHJhY2tlci5leGFtcGxlLm9yZy9hbm5vdW5jZTQ6aW5mb2Q0Om5hbWUxMDp1YnVudHUuaXNvMTI6cGllY2UgbGVuZ3RoaTE2Mzg0ZTY6bGVuZ3RoaTMyNzY4ZTY6cGllY2VzNDA6YWFhYWFhYWFhYWFhYWFhYWFhYWFiYmJiYmJiYmJiYmJiYmJiYmJiYmVl";
pub const BT_TORRENT_FIXTURE_BYTES: &[u8] = b"d8:announce35:http://tracker.example.org/announce4:infod4:name10:ubuntu.iso12:piece lengthi16384e6:lengthi32768e6:pieces40:aaaaaaaaaaaaaaaaaaaabbbbbbbbbbbbbbbbbbbbee";

pub fn write_torrent_fixture(path: &Path) {
    std::fs::write(path, BT_TORRENT_FIXTURE_BYTES).expect("torrent fixture should write");
}

pub fn rpc_request(method: RpcMethod, params: Vec<RpcValue>) -> JsonRpcRequest {
    JsonRpcRequest {
        jsonrpc: Some("2.0".to_owned()),
        id: None,
        method: method.as_str().to_owned(),
        params,
        meta: RpcMeta::default(),
    }
}

pub fn rpc_result_object(response: JsonRpcResponse) -> BTreeMap<String, RpcValue> {
    match response.result {
        Some(RpcValue::Object(payload)) => payload,
        other => panic!("unexpected object rpc result: {other:?}"),
    }
}

pub fn rpc_result_array(response: JsonRpcResponse) -> Vec<RpcValue> {
    match response.result {
        Some(RpcValue::Array(payload)) => payload,
        other => panic!("unexpected array rpc result: {other:?}"),
    }
}

pub fn rpc_string_field(payload: &BTreeMap<String, RpcValue>, field: &str) -> String {
    match payload.get(field) {
        Some(RpcValue::String(value)) => value.clone(),
        other => panic!("unexpected string field {field}: {other:?}"),
    }
}

pub fn rpc_u64_field(payload: &BTreeMap<String, RpcValue>, field: &str) -> u64 {
    match payload.get(field) {
        Some(RpcValue::String(value)) => value
            .parse()
            .unwrap_or_else(|error| panic!("unexpected u64 string field {field}: {error}")),
        Some(RpcValue::Number(value)) => (*value)
            .try_into()
            .unwrap_or_else(|_| panic!("unexpected negative number field {field}: {value}")),
        other => panic!("unexpected u64 field {field}: {other:?}"),
    }
}

pub fn rpc_bool_field(payload: &BTreeMap<String, RpcValue>, field: &str) -> bool {
    match payload.get(field) {
        Some(RpcValue::Bool(value)) => *value,
        Some(RpcValue::String(value)) => value
            .parse()
            .unwrap_or_else(|error| panic!("unexpected bool string field {field}: {error}")),
        other => panic!("unexpected bool field {field}: {other:?}"),
    }
}

pub fn add_torrent(dispatcher: &mut InProcessRpcDispatcher) -> String {
    let response = dispatcher.dispatch_json(rpc_request(
        RpcMethod::Aria2AddTorrent,
        vec![RpcValue::String(BT_TORRENT_FIXTURE_BASE64.to_owned())],
    ));
    match response.result {
        Some(RpcValue::String(gid)) => gid,
        other => panic!("unexpected addTorrent result: {other:?}"),
    }
}

pub fn add_magnet(dispatcher: &mut InProcessRpcDispatcher, magnet: &str) -> String {
    let response = dispatcher.dispatch_json(rpc_request(
        RpcMethod::Aria2AddUri,
        vec![RpcValue::String(magnet.to_owned())],
    ));
    match response.result {
        Some(RpcValue::String(gid)) => gid,
        other => panic!("unexpected addUri magnet result: {other:?}"),
    }
}

pub fn decode_hex_20(raw: &str) -> [u8; 20] {
    let mut chunks = raw.as_bytes().chunks_exact(2);
    let bytes = chunks
        .by_ref()
        .map(|chunk| {
            let pair = std::str::from_utf8(chunk).expect("hex field should stay ascii");
            u8::from_str_radix(pair, 16)
        })
        .collect::<Result<Vec<_>, _>>()
        .expect("hex field should parse");
    assert!(
        chunks.remainder().is_empty(),
        "hex field should contain an even number of digits"
    );
    bytes.try_into().expect("hex field should be 20 bytes")
}

pub fn peer_wire_handshake_and_frames(
    info_hash: [u8; 20],
    peer_id: [u8; 20],
    frames: &[PeerWireMessageKind],
) -> Vec<u8> {
    let mut bytes = PeerWireHandshakeModel::new(info_hash, peer_id).serialize();
    for frame in frames {
        bytes.extend_from_slice(
            &TorrentMessageModel::from_peer_wire_kind(frame.clone())
                .serialize_peer_wire_frame()
                .expect("peer-wire frame should serialize"),
        );
    }
    bytes
}

pub fn compact_peer(address: [u8; 4], port: u16) -> Vec<u8> {
    let mut bytes = address.to_vec();
    bytes.extend_from_slice(&port.to_be_bytes());
    bytes
}

pub fn compact_node(node_id_byte: u8, address: [u8; 4], port: u16) -> Vec<u8> {
    let mut bytes = vec![node_id_byte; 20];
    bytes.extend_from_slice(&address);
    bytes.extend_from_slice(&port.to_be_bytes());
    bytes
}

#[derive(Debug)]
pub struct RecordingDhtTransport {
    response: DhtMessageModel,
    seen: Mutex<Vec<(DhtNodeModel, DhtMessageModel)>>,
}

impl RecordingDhtTransport {
    pub const fn new(response: DhtMessageModel) -> Self {
        Self {
            response,
            seen: Mutex::new(Vec::new()),
        }
    }

    pub fn seen(&self) -> Vec<(DhtNodeModel, DhtMessageModel)> {
        self.seen
            .lock()
            .expect("dht seen mutex should not be poisoned")
            .clone()
    }
}

impl DhtTransport for RecordingDhtTransport {
    fn send_message(
        &self,
        node: &DhtNodeModel,
        message: &DhtMessageModel,
    ) -> Result<DhtMessageModel, TransportError> {
        self.seen
            .lock()
            .expect("dht seen mutex should not be poisoned")
            .push((node.clone(), message.clone()));
        Ok(self.response.clone())
    }
}

#[derive(Debug)]
pub struct FakePeerWireConnector {
    response_payload: Vec<u8>,
    seen: Mutex<Vec<PeerWireTransportRequest>>,
}

impl FakePeerWireConnector {
    pub const fn new(response_payload: Vec<u8>) -> Self {
        Self {
            response_payload,
            seen: Mutex::new(Vec::new()),
        }
    }

    pub fn seen(&self) -> Vec<PeerWireTransportRequest> {
        self.seen
            .lock()
            .expect("peer-wire seen mutex should not be poisoned")
            .clone()
    }
}

impl PeerWireTransportConnector for FakePeerWireConnector {
    fn connect_peer_wire(
        &self,
        request: &PeerWireTransportRequest,
    ) -> Result<PeerWireTransportResponse, TransportError> {
        self.seen
            .lock()
            .expect("peer-wire seen mutex should not be poisoned")
            .push(request.clone());
        Ok(PeerWireTransportResponse {
            endpoint: TransportEndpoint {
                scheme: TransportScheme::BitTorrent,
                address: request.endpoint.address.clone(),
            },
            payload: self.response_payload.clone(),
        })
    }
}

#[derive(Debug)]
pub struct LocalTrackerServer {
    announce_url: String,
    handle: Option<thread::JoinHandle<()>>,
}

impl LocalTrackerServer {
    pub fn spawn(compact_peers: Vec<u8>, interval_secs: u64) -> Self {
        let listener = TcpListener::bind("127.0.0.1:0").expect("tracker listener should bind");
        let addr = listener
            .local_addr()
            .expect("tracker listener should report local addr");
        let announce_url = format!("http://{addr}/announce");
        let handle = thread::spawn(move || {
            let mut payload = format!(
                "d8:intervali{interval_secs}e5:peers{}:",
                compact_peers.len()
            )
            .into_bytes();
            payload.extend_from_slice(&compact_peers);
            payload.extend_from_slice(b"e");

            let (mut stream, _) = listener.accept().expect("tracker client should connect");
            let mut request = [0_u8; 2048];
            let _ = stream
                .read(&mut request)
                .expect("tracker request should read");
            let response = format!(
                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
                payload.len()
            );
            stream
                .write_all(response.as_bytes())
                .expect("tracker response head should write");
            stream
                .write_all(&payload)
                .expect("tracker response body should write");
        });

        Self {
            announce_url,
            handle: Some(handle),
        }
    }

    pub fn announce_url(&self) -> &str {
        &self.announce_url
    }
}

impl Drop for LocalTrackerServer {
    fn drop(&mut self) {
        if let Some(handle) = self.handle.take() {
            handle.join().expect("tracker server thread should join");
        }
    }
}

pub fn tracker_transport() -> ReqwestTrackerTransport {
    ReqwestTrackerTransport::new().expect("reqwest tracker transport should build")
}