Contractless/src/bin/lookup_storage_cost.rs

use contractless::blocks::storage_values::{STORAGE_KEY_HASH_BYTES, STORAGE_VALUE_KEY_BYTES};
use contractless::common::cli_prompts::{prompt_hidden_nonempty, prompt_wallet_path};
use contractless::common::network_startup::get_connections;
use contractless::decode;
use contractless::env;
use contractless::json;
use contractless::records::memory::response_channels::generate_uid;
use contractless::standalone_tools::connections::handshake;
use contractless::wallets::structures::Wallet;

fn decimal_base_units(value: u64) -> f64 {
    value as f64 / 100_000_000.0
}

#[tokio::main]
async fn main() {
    // Command 52 asks a peer to price a storage lookup before the paid lookup is submitted.
    let hashmap_key = generate_uid();
    let rpc_command = 52;

    let args: Vec<String> = env::args().collect();
    if args.len() != 4 {
        println!(
            "Usage: ./lookup_storage_cost <storage_key_hash> <data_key_or_all> <wallet_address>"
        );
        println!("  storage_key_hash: 64 hex characters from the StorageKey transaction hash");
        println!("  data_key_or_all: a storage field key up to 50 bytes, or all");
        println!(
            "  wallet_address: short, long, or vanity address whose stored data will be priced"
        );
        return;
    }

    let storage_key = args[1].trim().to_string();
    let data_key = args[2].trim().to_string();
    let address = args[3].trim().to_string();

    if let Err(err) = validate_lookup_args(&storage_key, &data_key, &address) {
        eprintln!("{err}");
        return;
    }

    let wallet_path = prompt_wallet_path().await;
    let encryption_key = prompt_hidden_nonempty(
        "What is your wallet decryption key? ",
        "Wallet key cannot be empty. Please try again.",
    )
    .await;

    let request = format!("{storage_key}|{data_key}|{address}");
    let connections = get_connections().await;
    let mut connected = false;

    for conn in connections {
        if connected {
            break;
        }

        let socket_address = conn.parse().expect("Failed to parse the socket address");
        let result = handshake::connect_and_handshake(
            socket_address,
            request.clone(),
            rpc_command,
            handshake::HandshakeWallet::WalletKey {
                encryption_key: encryption_key.clone(),
                wallet_path: wallet_path.clone(),
            },
            hashmap_key,
        )
        .await;

        match result {
            Ok(response) => {
                if let Some(output) = decode_storage_cost_response(&response) {
                    println!("{output}");
                    connected = true;
                } else {
                    let response_text = String::from_utf8_lossy(&response);
                    let trimmed = response_text.trim();
                    if !trimmed.is_empty() {
                        println!("{trimmed}");
                    }
                    connected = true;
                }
            }
            Err(_) => connected = false,
        }
    }

    if !connected {
        eprintln!("failed to connect");
    }
}

fn decode_storage_cost_response(response: &[u8]) -> Option<String> {
    const TOTAL_BYTES_FIELD: usize = 4;
    const COST_PER_BYTE_FIELD: usize = 8;
    const TOTAL_COST_FIELD: usize = 8;
    const STORAGE_COST_RESPONSE_BYTES: usize = TOTAL_BYTES_FIELD
        + COST_PER_BYTE_FIELD
        + TOTAL_COST_FIELD
        + Wallet::SHORT_ADDRESS_BYTES_LENGTH;

    if response.len() != STORAGE_COST_RESPONSE_BYTES {
        return None;
    }

    let mut cursor = 0usize;
    let total_bytes = u32::from_le_bytes(
        response[cursor..cursor + TOTAL_BYTES_FIELD]
            .try_into()
            .ok()?,
    );
    cursor += TOTAL_BYTES_FIELD;

    let cost_per_byte = u64::from_le_bytes(
        response[cursor..cursor + COST_PER_BYTE_FIELD]
            .try_into()
            .ok()?,
    );
    cursor += COST_PER_BYTE_FIELD;

    let total_cost = u64::from_le_bytes(
        response[cursor..cursor + TOTAL_COST_FIELD]
            .try_into()
            .ok()?,
    );
    cursor += TOTAL_COST_FIELD;

    let payment_address = Wallet::bytes_to_short_address(&response[cursor..])?;
    let output = json!({
        "payment_address": payment_address,
        "total_cost": decimal_base_units(total_cost),
        "cost_per_byte": decimal_base_units(cost_per_byte),
        "total_bytes": total_bytes
    });

    serde_json::to_string_pretty(&output).ok()
}

fn validate_lookup_args(storage_key: &str, data_key: &str, address: &str) -> Result<(), String> {
    let storage_key_bytes =
        decode(storage_key).map_err(|err| format!("Invalid storage key hash: {err}"))?;
    if storage_key_bytes.len() != STORAGE_KEY_HASH_BYTES {
        return Err("Storage key hash must decode to 32 bytes.".to_string());
    }

    if data_key.is_empty() {
        return Err("Data key cannot be empty. Use all to price every key.".to_string());
    }
    if !data_key.eq_ignore_ascii_case("all") && data_key.as_bytes().len() > STORAGE_VALUE_KEY_BYTES
    {
        return Err("Data key must be 50 bytes or less.".to_string());
    }

    if Wallet::normalize_to_short_address(address).is_none() {
        return Err("Wallet address must be a valid short, long, or vanity address.".to_string());
    }

    Ok(())
}