libsession 0.1.3

//! Swarm ID computation and closest-swarm lookup.
//!
//! Maps public keys into swarm space via XOR-folding and finds the closest
//! swarm on the ID ring using wrapping distance.

use std::collections::HashMap;

use crate::network::key_types::X25519Pubkey;
use crate::network::service_node::ServiceNode;

/// Swarm identifier type (u64).
pub type SwarmId = u64;
/// Sentinel value indicating an invalid or unassigned swarm ID.
pub const INVALID_SWARM_ID: SwarmId = u64::MAX;

/// Maps a public key into swarm ID space by XOR-folding the 32 bytes into 8 bytes,
/// then interpreting as big-endian u64.
///
/// Mirrors `session::network::swarm::pubkey_to_swarm_space` from the C++ code.
pub fn pubkey_to_swarm_space(pk: &X25519Pubkey) -> SwarmId {
    let bytes = pk.as_bytes();
    let mut result: u64 = 0;

    // XOR four 8-byte chunks together
    for i in 0..4 {
        let mut buf = [0u8; 8];
        buf.copy_from_slice(&bytes[i * 8..(i + 1) * 8]);
        // In the C++ code, the bytes are memcpy'd into a u64 (native endian),
        // then after XOR the result is big_to_host. Since we read as native-endian
        // and then convert, we replicate by reading as little-endian on all platforms
        // and then at the end treating the result as big-endian.
        result ^= u64::from_ne_bytes(buf);
    }

    // The C++ does big_to_host_inplace on the XOR result
    u64::from_be(result)
}

/// Groups service nodes by their swarm_id and returns a sorted list of
/// (swarm_id, nodes) pairs.
pub fn generate_swarms(nodes: &[ServiceNode]) -> Vec<(SwarmId, Vec<ServiceNode>)> {
    let mut grouped: HashMap<SwarmId, Vec<ServiceNode>> = HashMap::new();

    for node in nodes {
        grouped.entry(node.swarm_id).or_default().push(node.clone());
    }

    let mut result: Vec<(SwarmId, Vec<ServiceNode>)> = grouped.into_iter().collect();
    result.sort_by_key(|(id, _)| *id);
    result
}

/// Finds the closest swarm to the given pubkey using the swarm ID ring.
///
/// Uses wrapping arithmetic distance, matching the C++ implementation.
pub fn get_swarm(
    swarm_pubkey: &X25519Pubkey,
    all_swarms: &[(SwarmId, Vec<ServiceNode>)],
) -> Option<(SwarmId, Vec<ServiceNode>)> {
    if all_swarms.is_empty() {
        return None;
    }

    if all_swarms.len() == 1 {
        return Some(all_swarms[0].clone());
    }

    let swarm_id = pubkey_to_swarm_space(swarm_pubkey);

    // Find the first swarm with id >= swarm_id (right boundary)
    let right_idx = all_swarms
        .iter()
        .position(|(id, _)| *id >= swarm_id);

    let right_idx = right_idx.unwrap_or(0); // wrap around if beyond end

    // Left is the one just before right (with wraparound)
    let left_idx = if right_idx == 0 {
        all_swarms.len() - 1
    } else {
        right_idx - 1
    };

    // Use wrapping distance
    let d_right = all_swarms[right_idx].0.wrapping_sub(swarm_id);
    let d_left = swarm_id.wrapping_sub(all_swarms[left_idx].0);

    let chosen = if d_right < d_left {
        right_idx
    } else {
        left_idx
    };

    Some(all_swarms[chosen].clone())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::network::key_types::Ed25519Pubkey;

    fn make_node(swarm_id: SwarmId, hex_suffix: u8) -> ServiceNode {
        let mut pk_bytes = [0u8; 32];
        pk_bytes[31] = hex_suffix;
        ServiceNode {
            ed25519_pubkey: Ed25519Pubkey(pk_bytes),
            ip: [1, 2, 3, 4],
            https_port: 443,
            omq_port: 22000,
            storage_server_version: [2, 11, 0],
            swarm_id,
            requested_unlock_height: 0,
        }
    }

    #[test]
    fn test_generate_swarms() {
        let nodes = vec![
            make_node(100, 1),
            make_node(200, 2),
            make_node(100, 3),
            make_node(300, 4),
        ];

        let swarms = generate_swarms(&nodes);
        assert_eq!(swarms.len(), 3);
        // Should be sorted by swarm_id
        assert_eq!(swarms[0].0, 100);
        assert_eq!(swarms[0].1.len(), 2);
        assert_eq!(swarms[1].0, 200);
        assert_eq!(swarms[1].1.len(), 1);
        assert_eq!(swarms[2].0, 300);
        assert_eq!(swarms[2].1.len(), 1);
    }

    #[test]
    fn test_get_swarm_single() {
        let swarms = vec![(100, vec![make_node(100, 1)])];
        let pk = X25519Pubkey([0u8; 32]);
        let result = get_swarm(&pk, &swarms).unwrap();
        assert_eq!(result.0, 100);
    }

    #[test]
    fn test_get_swarm_empty() {
        let pk = X25519Pubkey([0u8; 32]);
        assert!(get_swarm(&pk, &[]).is_none());
    }
}