Contractless/src/bin/lookup_storage.rs

use contractless::blocks::storage_values::{STORAGE_KEY_HASH_BYTES, STORAGE_VALUE_KEY_BYTES};
use contractless::blocks::transfer::{TransferTransaction, UnsignedTransferTransaction};
use contractless::common::cli_prompts::{
    ask_yes_no_question, prompt_hidden_nonempty, prompt_wallet_path,
};
use contractless::common::network_paths_and_settings::block_extension_and_paths;
use contractless::common::network_startup::get_connections;
use contractless::common::types::minimum_transfer_fee;
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;
use contractless::{encode, Utc};

#[derive(Clone)]
struct StorageCostQuote {
    payment_address: String,
    lookup_cost: u64,
    cost_per_byte: u64,
    total_bytes: u32,
}

fn display_base_units(value: u64) -> String {
    format!("{:.8}", value as f64 / 100_000_000.0)
}

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

#[tokio::main]
async fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() != 4 {
        println!("Usage: ./lookup_storage <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 returned"
        );
        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 wallet = match Wallet::try_obtain_wallet(encryption_key, Some(&wallet_path)).await {
        Ok(wallet) => wallet,
        Err(err) => {
            eprintln!("Wallet decryption failed: {err}");
            return;
        }
    };

    let request = format!("{storage_key}|{data_key}|{address}");
    let Some((peer, quote)) = quote_storage_lookup(&request, &wallet).await else {
        eprintln!("failed to connect");
        return;
    };

    let requester_is_node = wallet.saved.short_address == quote.payment_address;
    let txfee = if requester_is_node || quote.lookup_cost == 0 {
        0
    } else {
        minimum_transfer_fee(quote.lookup_cost, true)
    };
    let total_spend = quote.lookup_cost.saturating_add(txfee);

    let quote_json = json!({
        "payment_address": quote.payment_address,
        "lookup_cost": decimal_base_units(quote.lookup_cost),
        "txfee": decimal_base_units(txfee),
        "total_cost": decimal_base_units(total_spend),
        "cost_per_byte": decimal_base_units(quote.cost_per_byte),
        "total_bytes": quote.total_bytes
    });
    println!(
        "{}",
        serde_json::to_string_pretty(&quote_json).unwrap_or_else(|_| quote_json.to_string())
    );

    if total_spend > 0 {
        let (_, base_coin, _, _, _, _, _, _, _) = block_extension_and_paths();
        println!(
            "This lookup will spend {} {} total: {} lookup payment + {} transfer fee.",
            display_base_units(total_spend),
            base_coin.trim(),
            display_base_units(quote.lookup_cost),
            display_base_units(txfee)
        );
        if !ask_yes_no_question("Continue with paid storage lookup?").await {
            println!("Storage lookup cancelled.");
            return;
        }
    }

    let transfer_hex = match build_payment_transfer(&wallet, &quote, quote.lookup_cost, txfee).await
    {
        Ok(hex) => hex,
        Err(err) => {
            eprintln!("{err}");
            return;
        }
    };
    let lookup_request = format!("{request}|{transfer_hex}");

    match call_storage_lookup(peer, lookup_request, &wallet).await {
        Ok(response) => print_storage_response(&response),
        Err(err) => eprintln!("{err}"),
    }
}

async fn quote_storage_lookup(
    request: &str,
    wallet: &Wallet,
) -> Option<(std::net::SocketAddr, StorageCostQuote)> {
    let connections = get_connections().await;
    for conn in connections {
        let Ok(socket_address) = conn.parse() else {
            continue;
        };
        let result = handshake::connect_and_handshake(
            socket_address,
            request.to_string(),
            52,
            handshake::HandshakeWallet::WalletParts {
                public_key: wallet.saved.public_key.clone(),
                private_key: wallet.saved.private_key.clone(),
            },
            generate_uid(),
        )
        .await;
        let Ok(response) = result else {
            continue;
        };
        if let Some(quote) = decode_storage_cost_response(&response) {
            return Some((socket_address, quote));
        }
        let text = String::from_utf8_lossy(&response);
        if !text.trim().is_empty() {
            eprintln!("{}", text.trim());
            return None;
        }
    }
    None
}

async fn call_storage_lookup(
    peer: std::net::SocketAddr,
    request: String,
    wallet: &Wallet,
) -> Result<Vec<u8>, String> {
    handshake::connect_and_handshake(
        peer,
        request,
        53,
        handshake::HandshakeWallet::WalletParts {
            public_key: wallet.saved.public_key.clone(),
            private_key: wallet.saved.private_key.clone(),
        },
        generate_uid(),
    )
    .await
    .map_err(|err| format!("Storage lookup failed: {err}"))
}

async fn build_payment_transfer(
    wallet: &Wallet,
    quote: &StorageCostQuote,
    value: u64,
    txfee: u64,
) -> Result<String, String> {
    let (_, base_coin, _, _, _, _, _, _, _) = block_extension_and_paths();
    let unsigned = UnsignedTransferTransaction::new(
        2,
        Utc::now().timestamp() as u32,
        value,
        &base_coin,
        0,
        &wallet.saved.short_address,
        &quote.payment_address,
        txfee,
    )
    .await;
    let transfer = TransferTransaction::new(unsigned, &wallet.saved.private_key)
        .await
        .map_err(|err| format!("Failed to sign storage lookup payment: {err}"))?;
    let bytes = transfer
        .to_bytes()
        .await
        .map_err(|err| format!("Failed to serialize storage lookup payment: {err}"))?;
    Ok(encode(&bytes))
}

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

    if response.len() != 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 lookup_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..])?;

    Some(StorageCostQuote {
        payment_address,
        lookup_cost,
        cost_per_byte,
        total_bytes,
    })
}

fn print_storage_response(response: &[u8]) {
    let text = String::from_utf8_lossy(response);
    let trimmed = text.trim();
    if trimmed.is_empty() {
        return;
    }
    match serde_json::from_str::<serde_json::Value>(trimmed) {
        Ok(value) => match serde_json::to_string_pretty(&value) {
            Ok(pretty) => println!("{pretty}"),
            Err(_) => println!("{trimmed}"),
        },
        Err(_) => println!("{trimmed}"),
    }
}

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 lookup 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(())
}