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

pub(super) use std::{collections::BTreeMap, sync::Mutex, time::Instant};

pub(super) use aria2_rust_pro_compat::{
    BASELINE_COMMIT, is_required_pro_option, is_required_protocol,
};
pub(super) use aria2_rust_pro_core::{GoalProgress, RuntimeConfig};
pub(super) use aria2_rust_pro_protocol::{
    ChecksumSpec, DhtMessageModel, DhtNodeModel, DhtTransport, HttpCompletionState,
    HttpResponseHeaders, HttpResponseModel, HttpVersion, Protocol, ReqwestTrackerTransport,
    ResponseBody, TorrentPeerModel, TrackerPeerListModel, TrackerRequestModel,
    TrackerResponseModel, TrackerScrapeFileModel, TrackerScrapeModel, TrackerTransport,
    parse_metalink_document,
    torrent::{
        DhtMessageBody, DhtQueryModel, PeerWireBitfieldModel, PeerWireHandshakeModel,
        PeerWireMessageKind, PeerWirePieceBlockModel, TorrentMessageModel,
    },
    transport::{
        PeerWireTransportConnector, PeerWireTransportRequest, PeerWireTransportResponse,
        TransportEndpoint, TransportError, TransportScheme,
    },
};
pub(super) use aria2_rust_pro_rpc::{
    InProcessRpcDispatcher, JsonRpcRequest, JsonRpcResponse, RpcMethod, RpcValue, XmlRpcMethodCall,
    XmlRpcMethodResponse, XmlRpcParam, is_required_rpc_method, jsonrpc_request_from_json,
    rpc_value_to_xmlrpc, xmlrpc_method_call_from_xml, xmlrpc_method_call_to_xml,
    xmlrpc_method_response_from_xml, xmlrpc_method_response_to_xml, xmlrpc_value_to_rpc,
};
pub(super) use aria2_rust_pro_storage::{
    ByteSink, ControlFileVersion, ObservedByteSink, load_session_file,
};

pub(super) const BT_TORRENT_FIXTURE: &str = "ZDg6YW5ub3VuY2UzNTpodHRwOi8vdHJhY2tlci5leGFtcGxlLm9yZy9hbm5vdW5jZTQ6aW5mb2Q0Om5hbWUxMDp1YnVudHUuaXNvMTI6cGllY2UgbGVuZ3RoaTE2Mzg0ZTY6bGVuZ3RoaTMyNzY4ZTY6cGllY2VzNDA6YWFhYWFhYWFhYWFhYWFhYWFhYWFiYmJiYmJiYmJiYmJiYmJiYmJiYmVl";

pub(super) 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: Default::default(),
    }
}

pub(super) fn xmlrpc_request(method_name: &str, params: Vec<RpcValue>) -> XmlRpcMethodCall {
    XmlRpcMethodCall {
        method_name: method_name.to_owned(),
        params: params
            .into_iter()
            .map(|value| XmlRpcParam {
                value: rpc_value_to_xmlrpc(value),
            })
            .collect(),
        meta: Default::default(),
    }
}

pub(super) fn rpc_result(response: JsonRpcResponse) -> RpcValue {
    match response.result {
        Some(result) => result,
        None => panic!("unexpected rpc result envelope without result: {response:?}"),
    }
}

pub(super) fn rpc_result_from_xml(response: XmlRpcMethodResponse) -> RpcValue {
    match response {
        XmlRpcMethodResponse {
            value: Some(value),
            fault: None,
            ..
        } => xmlrpc_value_to_rpc(value),
        other => panic!("unexpected xmlrpc success envelope: {other:?}"),
    }
}

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

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

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

pub(super) 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(super) 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(super) 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:?}"),
    }
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub(super) struct BtStatusProbe {
    pub(super) completed_length: u64,
    pub(super) share_time: u64,
    pub(super) share_ratio: String,
    pub(super) seeder: bool,
    pub(super) connections: u64,
}

pub(super) fn bt_status_probe(dispatcher: &mut InProcessRpcDispatcher, gid: &str) -> BtStatusProbe {
    let status = rpc_object_result(dispatcher.dispatch_json(rpc_request(
        RpcMethod::Aria2TellStatus,
        vec![RpcValue::String(gid.to_owned())],
    )));
    BtStatusProbe {
        completed_length: rpc_u64_field(&status, "completedLength"),
        share_time: rpc_u64_field(&status, "shareTime"),
        share_ratio: rpc_string_field(&status, "shareRatio"),
        seeder: rpc_bool_field(&status, "seeder"),
        connections: rpc_u64_field(&status, "connections"),
    }
}

pub(super) fn add_magnet(dispatcher: &mut InProcessRpcDispatcher, suffix: u64) -> String {
    let magnet = format!(
        "magnet:?xt=urn:btih:{suffix:040x}&dn=fairness-{suffix}&tr=http%3A%2F%2Ftracker.example.org%2Fannounce"
    );
    let response = dispatcher.dispatch_json(rpc_request(
        RpcMethod::Aria2AddUri,
        vec![RpcValue::String(magnet)],
    ));
    match response.result {
        Some(RpcValue::String(gid)) => gid,
        other => panic!("unexpected addUri magnet result in fairness setup: {other:?}"),
    }
}

pub(super) fn decode_hex_20(raw: &str) -> [u8; 20] {
    let bytes = (0..raw.len())
        .step_by(2)
        .map(|offset| u8::from_str_radix(&raw[offset..offset + 2], 16))
        .collect::<Result<Vec<_>, _>>()
        .expect("info hash should be valid hex");
    bytes.try_into().expect("info hash should be 20 bytes")
}

pub(super) 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
}

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

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

    pub(super) fn seen(&self) -> Vec<PeerWireTransportRequest> {
        self.seen
            .lock()
            .expect("peer-wire seen requests 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 requests 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(super) struct FixedPeerWireConnector {
    pub(super) response_payload: Vec<u8>,
}

impl PeerWireTransportConnector for FixedPeerWireConnector {
    fn connect_peer_wire(
        &self,
        request: &PeerWireTransportRequest,
    ) -> Result<PeerWireTransportResponse, TransportError> {
        Ok(PeerWireTransportResponse {
            endpoint: TransportEndpoint {
                scheme: TransportScheme::BitTorrent,
                address: request.endpoint.address.clone(),
            },
            payload: self.response_payload.clone(),
        })
    }
}

pub(super) struct FixedDhtTransport {
    pub(super) response: DhtMessageModel,
}

impl DhtTransport for FixedDhtTransport {
    fn send_message(
        &self,
        _node: &DhtNodeModel,
        _message: &DhtMessageModel,
    ) -> Result<DhtMessageModel, TransportError> {
        Ok(self.response.clone())
    }
}

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

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

    pub(super) fn seen(&self) -> Vec<(String, DhtMessageModel)> {
        self.seen.lock().expect("dht seen mutex").clone()
    }
}

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

pub(super) fn peer_wire_payload(
    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(super) fn info_hash_bytes(input: &str) -> [u8; 20] {
    let encoded = input
        .split("xt=urn:btih:")
        .nth(1)
        .and_then(|rest| rest.split('&').next())
        .unwrap_or(input);
    let mut out = [0_u8; 20];
    for (index, chunk) in encoded.as_bytes().chunks_exact(2).enumerate() {
        let hex = std::str::from_utf8(chunk).expect("info hash should stay utf8 hex");
        out[index] = u8::from_str_radix(hex, 16).expect("info hash should decode from hex");
    }
    out
}

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

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