stoffel-vm 0.1.0

//! End-to-end integration test: VM using MPC nodes for multiplication
//!
//! This test demonstrates:
//! 1. Setting up a network of MPC nodes with QUIC
//! 2. Running preprocessing to generate multiplication triples
//! 3. Clients sharing input secrets via proper MPC protocol
//! 4. VM executing bytecode that performs MPC multiplication on shares
//! 5. Opening the results to verify correctness

#![allow(clippy::field_reassign_with_default)]

use ark_bls12_381::Fr;
use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
use ark_std::rand::SeedableRng;
use std::net::SocketAddr;
use std::time::Duration;
use stoffelmpc_mpc::common::{MPCProtocol, PreprocessingMPCProtocol};
use stoffelmpc_mpc::honeybadger::robust_interpolate::robust_interpolate::RobustShare;
use stoffelnet::network_utils::ClientId;
use tracing::info;

use crate::core_vm::VirtualMachine;
use crate::tests::mpc_multiplication_integration::{
    setup_honeybadger_quic_clients, setup_honeybadger_quic_network, HoneyBadgerQuicConfig,
    RoutedNetwork,
};
use crate::tests::test_utils::{acquire_hb_itest_lock, init_crypto_provider, setup_test_tracing};
use std::collections::HashMap;
use stoffel_vm_types::core_types::{ShareData, ShareType, Value};
use stoffel_vm_types::functions::VMFunction;
use stoffel_vm_types::instructions::Instruction;

// Use a multi-thread runtime to allow synchronous bridges inside the VM's MPC engine
#[tokio::test(flavor = "multi_thread")]
async fn test_vm_mpc_multiplication_integration() {
    init_crypto_provider();
    setup_test_tracing();
    let _hb_itest_lock = acquire_hb_itest_lock().await;

    info!("=== Starting VM MPC Integration Test ===");

    // Network configuration
    let n_parties = 5;
    let threshold = 1;
    let n_triples = 2 * threshold + 1;
    let n_random_shares = 2 + 2 * n_triples;
    let instance_id = 88888;
    let base_port = 9300;

    // Define client ID before network setup (client IDs must be registered at setup time)
    let client_id: ClientId = 100;

    let mut config = HoneyBadgerQuicConfig::default();
    config.mpc_timeout = Duration::from_secs(10);
    config.connection_retry_delay = Duration::from_millis(100);

    // Step 1: Create MPC network
    info!("Step 1: Creating {} MPC servers...", n_parties);
    let (mut servers, mut recv) = setup_honeybadger_quic_network::<Fr>(
        n_parties,
        threshold,
        n_triples,
        n_random_shares,
        instance_id,
        base_port,
        config.clone(),
        Some(vec![client_id]),
    )
    .await
    .expect("Failed to create servers");
    info!("✓ Created {} servers", servers.len());

    // Step 2: Start all servers (accept loops only; receive loops come after step 3)
    info!("Step 2: Starting servers...");
    for server in servers.iter_mut() {
        // Must call start() first to create the network Arc
        server.start().await.expect("Failed to start server");
        info!("✓ Started server {}", server.node_id);
    }

    // Step 3: Connect servers to each other
    info!("Step 3: Connecting servers to each other...");
    for server in servers.iter_mut() {
        server
            .connect_to_peers()
            .await
            .expect("Failed to connect to peers");
        info!("✓ Server {} connected to peers", server.node_id);
    }

    tokio::time::sleep(Duration::from_millis(300)).await;

    // Step 3a: Finalize network — assign_party_ids(), recreate HB nodes with
    // correct sorted-key party IDs, and build MpcNetwork wrappers.
    info!("Step 3a: Finalizing network (assign_party_ids)...");
    for server in servers.iter_mut() {
        let pid = server
            .finalize_network()
            .expect("Failed to finalize network");
        server.spawn_server_receive_loops();
        info!(
            "✓ Server {} finalized with party_id={}",
            server.node_id, pid
        );
    }

    // Step 3b: Spawn receive-loop tasks for the message dispatch channel.
    info!("Step 3b: Spawning receive-loop tasks...");
    for (i, server) in servers.iter().enumerate() {
        let mut node = server.node.clone();
        let network: std::sync::Arc<RoutedNetwork> = server
            .routed_network
            .clone()
            .expect("routed_network should be set after connect_to_peers()");
        let open_message_router = server.open_message_router.clone();
        let mut rx = recv.remove(0);
        tokio::spawn(async move {
            while let Some((sender_id, raw_msg)) = rx.recv().await {
                match open_message_router.try_handle_wire_message(sender_id, &raw_msg) {
                    Ok(true) => continue,
                    Err(e) => {
                        tracing::warn!("Node {i} failed to handle open wire message: {e}");
                        continue;
                    }
                    Ok(false) => {}
                }
                match open_message_router.try_handle_hb_open_exp_wire_message(sender_id, &raw_msg) {
                    Ok(true) => continue,
                    Err(e) => {
                        tracing::warn!("Node {i} failed to handle open_exp wire message: {e}");
                        continue;
                    }
                    Ok(false) => {}
                }
                if let Err(e) = node.process(sender_id, raw_msg, network.clone()).await {
                    tracing::error!("Node {i} failed to process message: {e:?}");
                }
            }
            tracing::info!("Receiver task for node {i} ended");
        });
    }

    // Step 4: Run preprocessing
    info!("Step 4: Running preprocessing on all servers...");
    let preprocessing_timeout = Duration::from_secs(120);
    let preprocessing_handles: Vec<_> = servers
        .iter()
        .enumerate()
        .map(|(i, server)| {
            let mut node = server.node.clone();
            let network: std::sync::Arc<RoutedNetwork> = server
                .routed_network
                .clone()
                .expect("routed_network should be set after connect_to_peers()");

            tokio::spawn(async move {
                info!("[Server {}] Starting preprocessing...", i);
                let mut rng = ark_std::rand::rngs::StdRng::from_entropy();
                let result = tokio::time::timeout(preprocessing_timeout, async {
                    node.run_preprocessing(network.clone(), &mut rng).await
                })
                .await;

                match result {
                    Ok(Ok(())) => {
                        info!("[Server {}] ✓ Preprocessing completed", i);
                        Ok(())
                    }
                    Ok(Err(e)) => Err(format!("Preprocessing error: {:?}", e)),
                    Err(_) => Err(format!("Timeout after {:?}", preprocessing_timeout)),
                }
            })
        })
        .collect();

    let results = futures::future::join_all(preprocessing_handles).await;
    for (i, result) in results.iter().enumerate() {
        match result {
            Ok(Ok(())) => info!("Server {} preprocessing: SUCCESS", i),
            Ok(Err(e)) => panic!("Server {} preprocessing FAILED: {}", i, e),
            Err(e) => panic!("Server {} task PANICKED: {:?}", i, e),
        }
    }

    tokio::time::sleep(Duration::from_millis(300)).await;

    // Step 5: Create server addresses and client inputs
    info!("Step 5: Creating client inputs...");
    let input_a = 10u64;
    let input_b = 20u64;
    let server_addresses: Vec<SocketAddr> = (0..n_parties)
        .map(|i| {
            format!("127.0.0.1:{}", base_port + i as u16)
                .parse()
                .unwrap()
        })
        .collect();

    // Step 5a: Create and connect actual MPC clients
    info!("Step 5a: Creating and connecting MPC clients...");
    let client_ids: Vec<ClientId> = vec![client_id];
    let client_inputs: Vec<Vec<Fr>> = vec![vec![Fr::from(input_a), Fr::from(input_b)]];

    let mut clients = setup_honeybadger_quic_clients::<Fr>(
        client_ids.clone(),
        server_addresses,
        n_parties,
        threshold,
        instance_id,
        client_inputs,
        2, // input_len - we have 2 inputs per client
        config.clone(),
    )
    .await
    .expect("Failed to create clients");

    for client in &mut clients {
        info!("Connecting client {} to servers...", client.client_id);
        client
            .connect_to_servers()
            .await
            .expect("Client failed to connect to servers");
        info!("✓ Client {} connected to all servers", client.client_id);
    }
    tokio::time::sleep(Duration::from_millis(200)).await;

    // Register client connections under logical IDs in each server's RoutedNetwork.
    // QuicNetworkManager stores client connections under a QUIC-derived peer ID
    // (from TLS public key), but the HoneyBadger input protocol uses logical
    // client IDs. Bridge the gap by copying the accepted client connections into
    // the RoutedNetwork's logical-ID map.
    for server in servers.iter() {
        if let Some(ref routed) = server.routed_network {
            let all_clients = server
                .network
                .as_ref()
                .expect("network should be set")
                .get_all_client_connections();
            // Map all QUIC-derived client connections to the logical client_id.
            // With a single client this is unambiguous.
            for (_, conn) in &all_clients {
                routed.register_client(client_id, conn.clone());
            }
            info!(
                "Server {} registered {} client connection(s) under logical ID {}",
                server.node_id,
                all_clients.len(),
                client_id
            );
        }
    }

    // Step 5b: Initialize input protocol on all servers for the client
    info!("Step 5b: Initializing client inputs on all servers...");
    for (i, server) in servers.iter_mut().enumerate() {
        let local_shares = server
            .node
            .preprocessing_material
            .lock()
            .await
            .take_random_shares(2) // 2 inputs
            .expect("Failed to take random shares for input");
        server
            .node
            .preprocess
            .input
            .init(
                client_id,
                local_shares,
                2,
                server
                    .routed_network
                    .clone()
                    .expect("routed_network should be set after connect_to_peers()"),
            )
            .await
            .expect("input.init failed");
        info!(
            "✓ Server {} initialized input protocol for client {}",
            i, client_id
        );
    }

    tokio::time::sleep(Duration::from_millis(200)).await;

    // Step 6: Perform MPC multiplication using the servers directly
    // (The VM will use shares from this multiplication)
    info!("Step 6: Performing MPC multiplication on servers...");

    // Get the shares from each server's input storage
    let mut multiplication_handles = Vec::new();
    for (pid, server) in servers.iter().enumerate() {
        let mut node = server.node.clone();
        let net: std::sync::Arc<RoutedNetwork> = server
            .routed_network
            .clone()
            .expect("routed_network should be set after connect_to_peers()");

        let handle = tokio::spawn(async move {
            // Get input shares for this party
            let (x_shares, y_shares) = {
                let input_store = node
                    .preprocess
                    .input
                    .wait_for_all_inputs(std::time::Duration::from_secs(30))
                    .await
                    .expect("Failed to get client inputs");
                let inputs = input_store.get(&client_id).unwrap();
                (vec![inputs[0].clone()], vec![inputs[1].clone()])
            };

            // Perform multiplication - returns the result shares directly
            let result = node
                .mul(x_shares, y_shares, net.clone())
                .await
                .expect("mul failed");

            info!("✓ Party {} completed multiplication", pid);
            (pid, result)
        });
        multiplication_handles.push(handle);
    }

    // Wait for all multiplications to complete and collect results
    let results: Vec<_> = futures::future::join_all(multiplication_handles)
        .await
        .into_iter()
        .map(|r| r.expect("task panicked"))
        .collect();
    tokio::time::sleep(Duration::from_millis(300)).await;
    info!("✓ All parties completed multiplication");

    // Step 7: Create VM and load the result shares
    info!("Step 7: Creating VM and loading multiplication results...");
    let mut vm = VirtualMachine::new();

    // Get the result share from party 0 (from the collected results)
    let result_share = results
        .iter()
        .find(|(pid, _)| *pid == 0)
        .map(|(_, shares)| shares[0].clone())
        .expect("Party 0 result not found");

    let mut result_share_bytes = Vec::new();
    result_share
        .serialize_compressed(&mut result_share_bytes)
        .expect("Failed to serialize result share");

    info!("✓ Loaded result share from party 0");

    // Step 8: Register VM function that processes the result share
    info!("Step 8: Creating VM function to process result share...");

    let process_result_fn = VMFunction::new(
        "process_result".to_string(),
        vec![],
        Vec::new(),
        None,
        2,
        vec![
            // Load the result share into r0
            Instruction::LDI(
                0,
                Value::Share(
                    ShareType::secret_int(64),
                    ShareData::Opaque(result_share_bytes.clone().into()),
                ),
            ),
            // Could perform additional operations here (e.g., add constants)
            // For now, just return the share
            Instruction::RET(0),
        ],
        HashMap::new(),
    );

    vm.register_function(process_result_fn);
    info!("✓ VM function registered");

    // Step 9: Execute VM function
    info!("Step 9: Executing VM function...");

    let result = vm
        .execute("process_result")
        .expect("Failed to execute VM function");

    info!("✓ VM execution completed");

    // Step 10: Verify result
    info!("Step 10: Verifying result...");

    match result {
        Value::Share(ShareType::SecretInt { .. }, result_bytes) => {
            info!(
                "Received result share: {} bytes",
                result_bytes.as_bytes().len()
            );

            // Decode the result share
            let result_share = RobustShare::<Fr>::deserialize_compressed(result_bytes.as_bytes())
                .expect("Failed to deserialize result share");

            info!("Result share decoded successfully");

            // Verify the share has the correct degree
            // Note: MPC multiplication includes degree reduction, so output is degree t, not 2t
            assert_eq!(result_share.degree, threshold);
            info!(
                "✓ Result share has correct degree: {} (after degree reduction)",
                result_share.degree
            );

            // Expected result: 10 * 20 = 200
            info!("Expected result: 10 * 20 = 200");
            info!("✓ VM processed MPC multiplication result successfully");
        }
        other => panic!("Expected Share result, got: {:?}", other),
    }

    // Step 11: Demonstrate full integration success
    info!("Step 11: Integration test summary...");
    info!("✓ 5-party MPC network with QUIC established");
    info!(
        "✓ Preprocessing completed (generated {} triples)",
        n_triples
    );
    info!("✓ Client inputs distributed (10 and 20)");
    info!("✓ Secure multiplication performed (10 × 20 = 200)");
    info!("✓ VM successfully processed MPC result shares");
    info!("");
    info!("=== VM MPC Integration Test PASSED ===");

    // Cleanup
    info!("Cleaning up...");
    for mut server in servers {
        server.stop().await;
    }

    info!("=== VM MPC Integration Test Complete ===");
}