Contractless/src/bin/lookup_online_node_count.rs

use contractless::common::network_startup::get_connections;
use contractless::env;
use contractless::records::memory::network_mapping::structs::{
    NETWORK_SNAPSHOT_HEADER_BYTES, NETWORK_SNAPSHOT_MAGIC, NETWORK_SNAPSHOT_VERSION,
    NODE_DELETED_TIMESTAMP_OFFSET, NODE_MONITOR_COUNT_OFFSET, NODE_RECORD_FIXED_BYTES,
};
use contractless::records::memory::response_channels::generate_uid;
use contractless::standalone_tools::connections::handshake;
use contractless::wallets::structures::Wallet;

fn online_node_count(response: &[u8]) -> Result<usize, String> {
    if response.len() < NETWORK_SNAPSHOT_HEADER_BYTES {
        return Err("network mapping response was shorter than its header".to_string());
    }
    if &response[..NETWORK_SNAPSHOT_MAGIC.len()] != NETWORK_SNAPSHOT_MAGIC {
        return Err("network mapping response used an unknown format".to_string());
    }
    if response[NETWORK_SNAPSHOT_MAGIC.len()] != NETWORK_SNAPSHOT_VERSION {
        return Err("network mapping response used an unsupported version".to_string());
    }

    let mapping_len = u32::from_le_bytes(
        response[5..9]
            .try_into()
            .map_err(|_| "network mapping length was invalid")?,
    ) as usize;
    let monitor_state_len = u32::from_le_bytes(
        response[9..13]
            .try_into()
            .map_err(|_| "network monitor-state length was invalid")?,
    ) as usize;
    let mapping_end = NETWORK_SNAPSHOT_HEADER_BYTES
        .checked_add(mapping_len)
        .ok_or_else(|| "network mapping length overflowed".to_string())?;
    let response_end = mapping_end
        .checked_add(monitor_state_len)
        .ok_or_else(|| "network monitor-state length overflowed".to_string())?;
    if response_end != response.len() {
        return Err("network mapping response length did not match its header".to_string());
    }

    let mut offset = NETWORK_SNAPSHOT_HEADER_BYTES;
    let mut online = 0usize;
    while offset < mapping_end {
        let fixed_end = offset
            .checked_add(NODE_RECORD_FIXED_BYTES)
            .ok_or_else(|| "network node record length overflowed".to_string())?;
        if fixed_end > mapping_end {
            return Err("network mapping ended inside a node record".to_string());
        }

        let deleted_timestamp = u64::from_le_bytes(
            response[offset + NODE_DELETED_TIMESTAMP_OFFSET
                ..offset + NODE_DELETED_TIMESTAMP_OFFSET + 8]
                .try_into()
                .map_err(|_| "node deletion timestamp was invalid")?,
        );
        let monitor_count = u16::from_le_bytes(
            response[offset + NODE_MONITOR_COUNT_OFFSET
                ..offset + NODE_MONITOR_COUNT_OFFSET + 2]
                .try_into()
                .map_err(|_| "node monitor count was invalid")?,
        ) as usize;
        let monitor_bytes = monitor_count
            .checked_mul(Wallet::SHORT_ADDRESS_BYTES_LENGTH)
            .ok_or_else(|| "node monitor list length overflowed".to_string())?;
        let record_end = fixed_end
            .checked_add(monitor_bytes)
            .ok_or_else(|| "network node record length overflowed".to_string())?;
        if record_end > mapping_end {
            return Err("network mapping ended inside a node monitor list".to_string());
        }

        if deleted_timestamp == 0 && monitor_count > 0 {
            online += 1;
        }
        offset = record_end;
    }

    Ok(online)
}

#[tokio::main]
async fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() != 3 {
        eprintln!("Usage: lookup_online_node_count WALLET_PATH WALLET_KEY");
        std::process::exit(2);
    }

    let wallet_path = args[1].clone();
    let encryption_key = args[2].clone();
    let rpc_command = 30;

    for connection in get_connections().await {
        let socket_address = match connection.parse() {
            Ok(address) => address,
            Err(err) => {
                eprintln!("Skipping invalid configured node {connection}: {err}");
                continue;
            }
        };
        let response = handshake::connect_and_handshake(
            socket_address,
            String::new(),
            rpc_command,
            handshake::HandshakeWallet::WalletKey {
                encryption_key: encryption_key.clone(),
                wallet_path: wallet_path.clone(),
            },
            generate_uid(),
        )
        .await;

        match response {
            Ok(bytes) => match online_node_count(&bytes) {
                Ok(count) => {
                    println!("{count}");
                    return;
                }
                Err(err) => eprintln!("Invalid mapping from {connection}: {err}"),
            },
            Err(err) => eprintln!("Failed to query {connection}: {err}"),
        }
    }

    eprintln!("failed to retrieve an online node count");
    std::process::exit(1);
}

#[cfg(test)]
mod tests {
    use super::*;

    fn snapshot_record(deleted_timestamp: u64, monitors: u16) -> Vec<u8> {
        let mut record = vec![0u8; NODE_RECORD_FIXED_BYTES];
        record[NODE_DELETED_TIMESTAMP_OFFSET..NODE_DELETED_TIMESTAMP_OFFSET + 8]
            .copy_from_slice(&deleted_timestamp.to_le_bytes());
        record[NODE_MONITOR_COUNT_OFFSET..NODE_MONITOR_COUNT_OFFSET + 2]
            .copy_from_slice(&monitors.to_le_bytes());
        record.extend(vec![
            0u8;
            monitors as usize * Wallet::SHORT_ADDRESS_BYTES_LENGTH
        ]);
        record
    }

    fn snapshot(records: &[Vec<u8>]) -> Vec<u8> {
        let mapping = records.concat();
        let mut response = Vec::new();
        response.extend_from_slice(NETWORK_SNAPSHOT_MAGIC);
        response.push(NETWORK_SNAPSHOT_VERSION);
        response.extend_from_slice(&(mapping.len() as u32).to_le_bytes());
        response.extend_from_slice(&0u32.to_le_bytes());
        response.extend_from_slice(&mapping);
        response
    }

    #[test]
    fn counts_only_active_monitored_nodes() {
        let response = snapshot(&[
            snapshot_record(0, 2),
            snapshot_record(0, 0),
            snapshot_record(1234, 0),
        ]);
        assert_eq!(online_node_count(&response).unwrap(), 1);
    }

    #[test]
    fn rejects_truncated_monitor_lists() {
        let mut response = snapshot(&[snapshot_record(0, 1)]);
        response.pop();
        assert!(online_node_count(&response).is_err());
    }
}