willow_v1_shell.rs

// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use ahe_traits::AheBase;
use client_traits::SecureAggregationClient;
use decryptor_traits::SecureAggregationDecryptor;
use googletest::prelude::container_eq;
use googletest::{gtest, verify_eq, verify_that};
use kahe_traits::KaheBase;
use messages::{
    CiphertextContribution, ClientMessage, DecryptionRequestContribution, DecryptorPublicKeyShare,
    PartialDecryptionRequest, PartialDecryptionResponse,
};
use proto_serialization_traits::{FromProto, ToProto};
use server_traits::SecureAggregationServer;
use shell_kahe::ShellKahe;
use shell_parameters::{create_shell_ahe_config, create_shell_kahe_config};
use shell_vahe::ShellVahe;
use status::StatusErrorCode;
use status_matchers_rs::status_is;
use std::collections::HashMap;
use std::rc::Rc;
use testing_utils::{
    generate_aggregation_config, generate_random_nonce, generate_random_unsigned_vector,
};
use verifier_traits::SecureAggregationVerifier;
use willow_v1_client::WillowV1Client;
use willow_v1_decryptor::{DecryptorState, WillowV1Decryptor};
use willow_v1_server::{ServerState, WillowV1Server};
use willow_v1_verifier::{VerifierState, WillowV1Verifier};

const CONTEXT_STRING: &[u8] = b"testing_context_string";

/// Encrypt and decrypt with a single decryptor and single client.
#[gtest]
fn encrypt_decrypt_one() -> googletest::Result<()> {
    let default_id = String::from("default");
    let aggregation_config = generate_aggregation_config(default_id.clone(), 16, 10, 1, 1);
    let max_number_of_decryptors = aggregation_config.max_number_of_decryptors;

    // Create common KAHE/VAHE instances.
    let vahe = Rc::new(
        ShellVahe::new(create_shell_ahe_config(max_number_of_decryptors).unwrap(), CONTEXT_STRING)
            .unwrap(),
    );
    let kahe = Rc::new(
        ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
            .unwrap(),
    );

    // Create client.
    let client =
        WillowV1Client::new_with_randomly_generated_seed(Rc::clone(&kahe), Rc::clone(&vahe))
            .unwrap();

    let mut decryptor_state = DecryptorState::default();
    let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(Rc::clone(&vahe)).unwrap();

    // Create server.
    let server = WillowV1Server { kahe: Rc::clone(&kahe), vahe: Rc::clone(&vahe) };
    let mut server_state = ServerState::default();

    // Create verifier.
    let verifier = WillowV1Verifier { vahe: Rc::clone(&vahe) };
    let mut verifier_state = VerifierState::default();

    // Decryptor generates public key share.
    let public_key_share = decryptor.create_public_key_share(&mut decryptor_state).unwrap();

    // Server handles the public key share.
    server
        .handle_decryptor_public_key_share(public_key_share, "Decryptor 0", &mut server_state)
        .unwrap();

    // Server creates the public key.
    let public_key = server.create_decryptor_public_key(&server_state).unwrap();

    // Client encrypts.
    let input_values = vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1];
    let client_plaintext = HashMap::from([(default_id.as_str(), input_values.as_slice())]);
    let nonce = generate_random_nonce();
    let client_message =
        client.create_client_message(&client_plaintext, &public_key, &nonce).unwrap();

    // The client message is split and handled by the server and verifier.
    let (ciphertext_contribution, decryption_request_contribution) =
        server.split_client_message(client_message).unwrap();
    verifier.verify_and_include(decryption_request_contribution, &mut verifier_state).unwrap();
    server.handle_ciphertext_contribution(ciphertext_contribution, &mut server_state).unwrap();

    // Verifier creates the partial decryption request.
    let pd_ct = verifier.create_partial_decryption_request(verifier_state).unwrap();

    // Decryptor creates partial decryption.
    let pd = decryptor.handle_partial_decryption_request(pd_ct, &decryptor_state).unwrap();

    // Server handles the partial decryption.
    server.handle_partial_decryption(pd, &mut server_state).unwrap();

    // Server recovers the aggregation result.
    let aggregation_result = server.recover_aggregation_result(&server_state).unwrap();

    // Check that the (padded) result matches the client plaintext.
    verify_that!(aggregation_result.keys().collect::<Vec<_>>(), container_eq([&default_id]))?;
    let client_plaintext_length = client_plaintext.get(default_id.as_str()).unwrap().len();
    verify_eq!(
        aggregation_result.get(default_id.as_str()).unwrap()[..client_plaintext_length],
        client_plaintext.get(default_id.as_str()).unwrap()[..]
    )
}

/// Encrypt and decrypt with a single decryptor and single client, using serialization.
#[gtest]
fn encrypt_decrypt_one_serialized() -> googletest::Result<()> {
    let default_id = String::from("default");
    let aggregation_config = generate_aggregation_config(default_id.clone(), 16, 10, 1, 1);
    let max_number_of_decryptors = aggregation_config.max_number_of_decryptors;

    // Create common KAHE/VAHE instances.
    let kahe = Rc::new(
        ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
            .unwrap(),
    );
    let vahe = Rc::new(
        ShellVahe::new(create_shell_ahe_config(max_number_of_decryptors).unwrap(), CONTEXT_STRING)
            .unwrap(),
    );

    // Create client.
    let client =
        WillowV1Client::new_with_randomly_generated_seed(Rc::clone(&kahe), Rc::clone(&vahe))
            .unwrap();

    // Create decryptor.
    let mut decryptor_state = DecryptorState::default();
    let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(Rc::clone(&vahe)).unwrap();

    // Create server.
    let server = WillowV1Server { kahe: Rc::clone(&kahe), vahe: Rc::clone(&vahe) };
    let mut server_state = ServerState::default();

    // Create verifier.
    let verifier = WillowV1Verifier { vahe: Rc::clone(&vahe) };
    let mut verifier_state = VerifierState::default();

    // Decryptor generates public key share.
    let public_key_share = decryptor.create_public_key_share(&mut decryptor_state).unwrap();

    // Serialize and deserialize the public key share.
    let public_key_share_proto = public_key_share.to_proto(decryptor.vahe.as_ref())?;
    let public_key_share: DecryptorPublicKeyShare<ShellVahe> =
        DecryptorPublicKeyShare::<ShellVahe>::from_proto(
            public_key_share_proto,
            server.vahe.as_ref(),
        )?;

    // Server handles the public key share.
    server
        .handle_decryptor_public_key_share(public_key_share, "Decryptor 0", &mut server_state)
        .unwrap();

    // Server creates the public key.
    let public_key = server.create_decryptor_public_key(&server_state).unwrap();

    // Serialize and deserialize the public key.
    let public_key_proto = public_key.to_proto(server.vahe.as_ref())?;
    let public_key = messages::DecryptorPublicKey::<ShellVahe>::from_proto(
        public_key_proto,
        client.vahe.as_ref(),
    )?;

    // Client encrypts.
    let client_plaintext =
        HashMap::from([(default_id.clone(), vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1])]);
    let nonce = generate_random_nonce();
    let client_message = client
        .create_client_message(
            &ShellKahe::plaintext_as_slice(&client_plaintext),
            &public_key,
            &nonce,
        )
        .unwrap();

    // Serialize and deserialize the client message.
    let client_message_proto = client_message.to_proto(&client)?;
    let client_message: ClientMessage<ShellKahe, ShellVahe> =
        ClientMessage::from_proto(client_message_proto, &server)?;

    // The client message is split and handled by the server and verifier.
    let (ciphertext_contribution, decryption_request_contribution) =
        server.split_client_message(client_message).unwrap();

    // Serialize and deserialize the contributions.
    let ciphertext_contribution_proto = ciphertext_contribution.to_proto(&server)?;
    let ciphertext_contribution: CiphertextContribution<ShellKahe, ShellVahe> =
        CiphertextContribution::from_proto(ciphertext_contribution_proto, &server)?;

    let decryption_request_contribution_proto =
        decryption_request_contribution.to_proto(&server)?;
    let decryption_request_contribution: DecryptionRequestContribution<ShellVahe> =
        DecryptionRequestContribution::from_proto(
            decryption_request_contribution_proto,
            &verifier,
        )?;

    verifier.verify_and_include(decryption_request_contribution, &mut verifier_state).unwrap();
    server.handle_ciphertext_contribution(ciphertext_contribution, &mut server_state).unwrap();

    // Verifier creates the partial decryption request.
    let pd_ct = verifier.create_partial_decryption_request(verifier_state).unwrap();

    // Serialize and deserialize the partial decryption request.
    let pd_ct_proto = pd_ct.to_proto(&verifier)?;
    let pd_ct: PartialDecryptionRequest<ShellVahe> =
        PartialDecryptionRequest::from_proto(pd_ct_proto, &decryptor)?;

    // Decryptor creates partial decryption.
    let pd = decryptor.handle_partial_decryption_request(pd_ct, &decryptor_state).unwrap();

    // Serialize and deserialize the partial decryption.
    let pd_proto = pd.to_proto(&decryptor)?;
    let pd: PartialDecryptionResponse<ShellVahe> =
        PartialDecryptionResponse::from_proto(pd_proto, &server)?;

    // Server handles the partial decryption.
    server.handle_partial_decryption(pd, &mut server_state).unwrap();

    // Server recovers the aggregation result.
    let aggregation_result = server.recover_aggregation_result(&server_state).unwrap();

    // Check that the (padded) result matches the client plaintext.
    verify_that!(aggregation_result.keys().collect::<Vec<_>>(), container_eq([&default_id]))?;
    let client_plaintext_length = client_plaintext.get(&default_id).unwrap().len();
    verify_eq!(
        aggregation_result.get(&default_id).unwrap()[..client_plaintext_length],
        client_plaintext.get(&default_id).unwrap()[..]
    )
}

// Encrypt and decrypt with multiple clients and a single decryptor.
#[gtest]
fn encrypt_decrypt_multiple_clients() -> googletest::Result<()> {
    const NUM_CLIENTS: i64 = 10;
    let default_id = String::from("default");
    let aggregation_config =
        generate_aggregation_config(default_id.clone(), 16, 10, 1, NUM_CLIENTS);
    let max_number_of_decryptors = aggregation_config.max_number_of_decryptors;

    // Create common KAHE/VAHE instances.
    let vahe = Rc::new(
        ShellVahe::new(create_shell_ahe_config(max_number_of_decryptors).unwrap(), CONTEXT_STRING)
            .unwrap(),
    );
    let kahe = Rc::new(
        ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
            .unwrap(),
    );

    // Create clients.
    let mut clients = vec![];
    for _ in 0..NUM_CLIENTS {
        let client =
            WillowV1Client::new_with_randomly_generated_seed(Rc::clone(&kahe), Rc::clone(&vahe))
                .unwrap();
        clients.push(client);
    }

    // Create decryptor.
    let mut decryptor_state = DecryptorState::default();
    let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(Rc::clone(&vahe)).unwrap();

    // Create server.
    let server = WillowV1Server { kahe: Rc::clone(&kahe), vahe: Rc::clone(&vahe) };
    let mut server_state = ServerState::default();

    // Create verifier.
    let verifier = WillowV1Verifier { vahe: Rc::clone(&vahe) };
    let mut verifier_state = VerifierState::default();

    // Decryptor generates public key share.
    let public_key_share = decryptor.create_public_key_share(&mut decryptor_state).unwrap();

    // Server handles the public key share.
    server
        .handle_decryptor_public_key_share(public_key_share, "Decryptor 0", &mut server_state)
        .unwrap();

    // Server creates the public key.
    let public_key = server.create_decryptor_public_key(&server_state).unwrap();

    // Clients encrypt.
    let mut expected_output = vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
    let mut client_messages = vec![];
    for client in &mut clients {
        let client_input_values = vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1];
        for i in 0..expected_output.len() {
            expected_output[i] += client_input_values[i];
        }
        let client_plaintext =
            HashMap::from([(default_id.as_str(), client_input_values.as_slice())]);
        let nonce = generate_random_nonce();
        let client_message =
            client.create_client_message(&client_plaintext, &public_key, &nonce).unwrap();
        client_messages.push(client_message);
    }

    // Sort client messages by nonce.
    client_messages.sort_by(|a, b| a.nonce.cmp(&b.nonce));

    // Handle client messages.
    for client_message in client_messages.clone() {
        // The client message is split and handled by the server and verifier.
        let (ciphertext_contribution, decryption_request_contribution) =
            server.split_client_message(client_message).unwrap();
        verifier.verify_and_include(decryption_request_contribution, &mut verifier_state).unwrap();
        server.handle_ciphertext_contribution(ciphertext_contribution, &mut server_state).unwrap();
    }

    // Verify again using two states and merge the states to check that merge works.
    let mut verifier_state_1 = VerifierState::default();
    let mut verifier_state_2 = VerifierState::default();
    let half = client_messages.len() / 2;
    for (i, client_message) in client_messages.into_iter().enumerate() {
        let (_, decryption_request_contribution) =
            server.split_client_message(client_message).unwrap();
        let mut verifier_state =
            if i < half { &mut verifier_state_1 } else { &mut verifier_state_2 };
        verifier.verify_and_include(decryption_request_contribution, &mut verifier_state).unwrap();
    }
    let verifier_state_merged = verifier.merge_states(verifier_state_1, verifier_state_2).unwrap();

    // Run the rest of the protocol twice, once with each of the the two copies of the verifier state.
    for (mut server_state, verifier_state) in
        [(server_state.clone(), verifier_state), (server_state, verifier_state_merged)]
    {
        // Verifier creates the partial decryption request.
        let pd_ct = verifier.create_partial_decryption_request(verifier_state).unwrap();

        // Decryptor creates partial decryption.
        let pd = decryptor.handle_partial_decryption_request(pd_ct, &decryptor_state).unwrap();

        // Server handles the partial decryption.
        server.handle_partial_decryption(pd, &mut server_state).unwrap();

        // Server recovers the aggregation result.
        let aggregation_result = server.recover_aggregation_result(&server_state).unwrap();

        // Check that the (padded) result matches the client plaintext.
        verify_that!(aggregation_result.keys().collect::<Vec<_>>(), container_eq([&default_id]))?;
        verify_eq!(
            aggregation_result.get(default_id.as_str()).unwrap()[..expected_output.len()],
            expected_output
        )?;
    }
    Ok(())
}

// Encrypt and decrypt with multiple clients including invalid client proofs and a single decryptor.
#[gtest]
fn encrypt_decrypt_multiple_clients_including_invalid_proofs() -> googletest::Result<()> {
    const NUM_MAX_CLIENTS: i64 = 10;
    const NUM_GOOD_CLIENTS: i64 = 10;
    const NUM_BAD_CLIENTS: i64 = 5;
    let default_id = String::from("default");
    let aggregation_config =
        generate_aggregation_config(default_id.clone(), 16, 10, 1, NUM_MAX_CLIENTS);
    let max_number_of_decryptors = aggregation_config.max_number_of_decryptors;

    // Create common KAHE/VAHE instances.
    let vahe = Rc::new(
        ShellVahe::new(create_shell_ahe_config(max_number_of_decryptors).unwrap(), CONTEXT_STRING)
            .unwrap(),
    );
    let kahe = Rc::new(
        ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
            .unwrap(),
    );

    // Create clients.
    let mut good_clients = vec![];
    for _ in 0..NUM_GOOD_CLIENTS {
        let client =
            WillowV1Client::new_with_randomly_generated_seed(Rc::clone(&kahe), Rc::clone(&vahe))
                .unwrap();
        good_clients.push(client);
    }

    // Create bad clients.
    let mut bad_clients = vec![];
    for _ in 0..NUM_BAD_CLIENTS {
        let client =
            WillowV1Client::new_with_randomly_generated_seed(Rc::clone(&kahe), Rc::clone(&vahe))
                .unwrap();
        bad_clients.push(client);
    }

    // Create decryptor.
    let mut decryptor_state = DecryptorState::default();
    let decryptor = WillowV1Decryptor::new_with_randomly_generated_seed(Rc::clone(&vahe)).unwrap();

    // Create server.
    let server = WillowV1Server { kahe: Rc::clone(&kahe), vahe: Rc::clone(&vahe) };
    let mut server_state = ServerState::default();

    // Create verifier.
    let verifier = WillowV1Verifier { vahe: Rc::clone(&vahe) };
    let mut verifier_state = VerifierState::default();

    // Decryptor generates public key share.
    let public_key_share = decryptor.create_public_key_share(&mut decryptor_state).unwrap();

    // Server handles the public key share.
    server
        .handle_decryptor_public_key_share(public_key_share, "Decryptor 0", &mut server_state)
        .unwrap();

    // Server creates the public key.
    let public_key = server.create_decryptor_public_key(&server_state).unwrap();

    // Good Clients encrypt and should be included in the aggregation.
    let mut expected_output = vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
    let mut client_messages = vec![];
    for client in &mut good_clients {
        let client_input_values = vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1];
        for i in 0..expected_output.len() {
            expected_output[i] += client_input_values[i];
        }
        let client_plaintext =
            HashMap::from([(default_id.as_str(), client_input_values.as_slice())]);
        let nonce = generate_random_nonce();
        let client_message =
            client.create_client_message(&client_plaintext, &public_key, &nonce).unwrap();
        client_messages.push(client_message);
    }

    // Sort client messages by nonce.
    client_messages.sort_by(|a, b| a.nonce.cmp(&b.nonce));

    // Handle client messages.
    for client_message in client_messages {
        // The client message is split and handled by the server and verifier.
        let (ciphertext_contribution, decryption_request_contribution) =
            server.split_client_message(client_message).unwrap();
        verifier.verify_and_include(decryption_request_contribution, &mut verifier_state).unwrap();
        server.handle_ciphertext_contribution(ciphertext_contribution, &mut server_state).unwrap();
    }

    // Use first bad client to create a proof object that the others will use.
    let bad_proof;
    {
        let client = &mut bad_clients[0];
        let client_input_values = vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1];
        let client_plaintext =
            HashMap::from([(default_id.as_str(), client_input_values.as_slice())]);
        let nonce = generate_random_nonce();
        let client_message =
            client.create_client_message(&client_plaintext, &public_key, &nonce).unwrap();
        bad_proof = client_message.proof;
    }
    // Bad Clients encrypt and should not be included in the aggregation.
    let mut client_messages = vec![];
    for i in 1..bad_clients.len() {
        let client = &mut bad_clients[i];
        let client_input_values = vec![8, 7, 6, 5, 4, 3, 2, 1, 2, 3, 4, 5, 6, 7, 8];
        let client_plaintext =
            HashMap::from([(default_id.as_str(), client_input_values.as_slice())]);
        let nonce = generate_random_nonce();
        let mut client_message =
            client.create_client_message(&client_plaintext, &public_key, &nonce).unwrap();
        client_message.proof = bad_proof.clone();
        client_messages.push(client_message);
    }
    client_messages.sort_by(|a, b| a.nonce.cmp(&b.nonce));
    for client_message in client_messages {
        // The client message is split and handled by the server and verifier.
        let (_ciphertext_contribution, decryption_request_contribution) =
            server.split_client_message(client_message).unwrap();
        verify_that!(
            verifier.verify_and_include(decryption_request_contribution, &mut verifier_state),
            status_is(StatusErrorCode::PermissionDenied)
        )?;
    }

    // Verifier creates the partial decryption request.
    let pd_ct = verifier.create_partial_decryption_request(verifier_state).unwrap();

    // Decryptor creates partial decryption.
    let pd = decryptor.handle_partial_decryption_request(pd_ct, &decryptor_state).unwrap();

    // Server handles the partial decryption.
    server.handle_partial_decryption(pd, &mut server_state).unwrap();

    // Server recovers the aggregation result.
    let aggregation_result = server.recover_aggregation_result(&server_state).unwrap();

    // Check that the (padded) result matches the client plaintext.
    verify_that!(aggregation_result.keys().collect::<Vec<_>>(), container_eq([&default_id]))?;
    verify_eq!(
        aggregation_result.get(default_id.as_str()).unwrap()[..expected_output.len()],
        expected_output
    )
}

/// Encrypt and decrypt with multiple clients and multiple decryptors.
/// Note: This test uses RLWE parameters for production use.
#[gtest]
fn encrypt_decrypt_many_clients_decryptors() -> googletest::Result<()> {
    const INPUT_LENGTH: isize = 100_000; // 100K
    const INPUT_DOMAIN: i64 = 1i64 << 32;
    const MAX_NUM_CLIENTS: i64 = 10_000_000; // used to generate parameters.
    const MAX_NUM_DECRYPTORS: i64 = 100; // used to generate parameters.
    const NUM_CLIENTS: usize = 3; // Actual number of clients to create.
    const NUM_DECRYPTORS: usize = 3; // Actual number of decryptors to create.

    let default_id = String::from("default");
    let aggregation_config = generate_aggregation_config(
        default_id.clone(),
        INPUT_LENGTH,
        INPUT_DOMAIN,
        MAX_NUM_DECRYPTORS,
        MAX_NUM_CLIENTS,
    );
    // The parameters used by `create_shell_ahe_config` below is secure for only a single
    // decryptor when assuming all clients are corrupted. As we are testing functionality
    // here anyway, we set `max_number_of_decryptors` to 1 just to obtain the AHE parameters.
    let max_number_of_decryptors = 1;

    // Create common KAHE/VAHE instances.
    let vahe = Rc::new(
        ShellVahe::new(create_shell_ahe_config(max_number_of_decryptors).unwrap(), CONTEXT_STRING)
            .unwrap(),
    );
    let kahe = Rc::new(
        ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
            .unwrap(),
    );

    // Create server.
    let server = WillowV1Server { kahe: Rc::clone(&kahe), vahe: Rc::clone(&vahe) };
    let mut server_state = ServerState::default();

    // Create verifier.
    let verifier = WillowV1Verifier { vahe: Rc::clone(&vahe) };
    let mut verifier_state = VerifierState::default();

    // Create decryptors.
    let mut decryptors = vec![];
    let mut decryptor_states = vec![];
    for i in 0..NUM_DECRYPTORS {
        let mut decryptor_state = DecryptorState::default();
        let decryptor =
            WillowV1Decryptor::new_with_randomly_generated_seed(Rc::clone(&vahe)).unwrap();

        // Decryptor generates public key share.
        let public_key_share = decryptor.create_public_key_share(&mut decryptor_state).unwrap();

        // Server handles the public key share.
        server
            .handle_decryptor_public_key_share(
                public_key_share,
                format!("Decryptor {i}").as_str(),
                &mut server_state,
            )
            .unwrap();

        decryptors.push(decryptor);
        decryptor_states.push(decryptor_state);
    }

    // Server creates the public key.
    let public_key = server.create_decryptor_public_key(&server_state).unwrap();

    // Create clients, and each client generates their messages.
    let mut expected_output = vec![0; INPUT_LENGTH as usize];
    let mut client_messages = vec![];
    for _ in 0..NUM_CLIENTS {
        let client =
            WillowV1Client::new_with_randomly_generated_seed(Rc::clone(&kahe), Rc::clone(&vahe))
                .unwrap();

        let client_input_values =
            generate_random_unsigned_vector(INPUT_LENGTH as usize, INPUT_DOMAIN as u64);
        for i in 0..expected_output.len() {
            expected_output[i] += client_input_values[i];
        }
        let client_plaintext =
            HashMap::from([(default_id.as_str(), client_input_values.as_slice())]);
        let nonce = generate_random_nonce();
        let client_message =
            client.create_client_message(&client_plaintext, &public_key, &nonce).unwrap();
        client_messages.push(client_message);
    }

    // Sort client messages by nonce.
    client_messages.sort_by(|a, b| a.nonce.cmp(&b.nonce));

    // Handle client messages.
    for client_message in client_messages {
        // The client message is split and handled by the server and verifier.
        let (ciphertext_contribution, decryption_request_contribution) =
            server.split_client_message(client_message).unwrap();
        verifier.verify_and_include(decryption_request_contribution, &mut verifier_state).unwrap();
        server.handle_ciphertext_contribution(ciphertext_contribution, &mut server_state).unwrap();
    }

    // Verifier creates the partial decryption request.
    let pd_ct = verifier.create_partial_decryption_request(verifier_state).unwrap();

    // Decryptors perform partial decryption.
    for i in 0..NUM_DECRYPTORS {
        // Each decryptor creates partial decryption.
        let pd = decryptors[i]
            .handle_partial_decryption_request(pd_ct.clone(), &decryptor_states[i])
            .unwrap();

        // Server handles the partial decryption.
        server.handle_partial_decryption(pd, &mut server_state).unwrap();
    }

    // Server recovers the aggregation result.
    let aggregation_result = server.recover_aggregation_result(&server_state).unwrap();

    // Check that the (padded) result matches the client plaintext.
    verify_that!(aggregation_result.keys().collect::<Vec<_>>(), container_eq([&default_id]))?;
    verify_eq!(
        aggregation_result.get(default_id.as_str()).unwrap()[..expected_output.len()],
        expected_output
    )
}

// Encrypt and decrypt with multiple clients and multiple decryptors, but no dropout.
#[gtest]
fn encrypt_decrypt_no_dropout() -> googletest::Result<()> {
    const NUM_CLIENTS: i64 = 10;
    const NUM_DECRYPTORS: i64 = 10;
    let default_id = String::from("default");
    let aggregation_config =
        generate_aggregation_config(default_id.clone(), 16, 10, NUM_DECRYPTORS, NUM_CLIENTS);
    // The parameters used by `create_shell_ahe_config` below is secure for only a single
    // decryptor when assuming all clients are corrupted. As we are testing functionality
    // here anyway, we set `max_number_of_decryptors` to 1 just to obtain the AHE parameters.
    let max_number_of_decryptors = 1;

    // Create common KAHE/VAHE instances.
    let vahe = Rc::new(
        ShellVahe::new(create_shell_ahe_config(max_number_of_decryptors).unwrap(), CONTEXT_STRING)
            .unwrap(),
    );
    let kahe = Rc::new(
        ShellKahe::new(create_shell_kahe_config(&aggregation_config).unwrap(), CONTEXT_STRING)
            .unwrap(),
    );

    // Create clients.
    let mut clients = vec![];
    for _ in 0..NUM_CLIENTS {
        let client =
            WillowV1Client::new_with_randomly_generated_seed(Rc::clone(&kahe), Rc::clone(&vahe))
                .unwrap();
        clients.push(client);
    }

    // Create decryptors.
    let mut decryptor_states = vec![];
    let mut decryptors = vec![];
    for _ in 0..NUM_DECRYPTORS {
        let decryptor_state = DecryptorState::default();
        let decryptor =
            WillowV1Decryptor::new_with_randomly_generated_seed(Rc::clone(&vahe)).unwrap();
        decryptor_states.push(decryptor_state);
        decryptors.push(decryptor);
    }

    // Create server.
    let server = WillowV1Server { kahe: Rc::clone(&kahe), vahe: Rc::clone(&vahe) };
    let mut server_state = ServerState::default();

    // Create verifier.
    let verifier = WillowV1Verifier { vahe: Rc::clone(&vahe) };
    let mut verifier_state = VerifierState::default();

    // Decryptors generate public key shares.
    for i in 0..decryptors.len() {
        let public_key_share =
            decryptors[i].create_public_key_share(&mut decryptor_states[i]).unwrap();
        // Server handles the public key share.
        server
            .handle_decryptor_public_key_share(
                public_key_share,
                format!("Decryptor {i}").as_str(),
                &mut server_state,
            )
            .unwrap();
    }

    // Server creates the public key.
    let public_key = server.create_decryptor_public_key(&server_state).unwrap();

    // Clients encrypt.
    let mut expected_output = vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
    let mut client_messages = vec![];
    for client in &mut clients {
        let client_input_values = vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1];
        for i in 0..expected_output.len() {
            expected_output[i] += client_input_values[i];
        }
        let client_plaintext =
            HashMap::from([(default_id.as_str(), client_input_values.as_slice())]);
        let nonce = generate_random_nonce();
        let client_message =
            client.create_client_message(&client_plaintext, &public_key, &nonce).unwrap();
        client_messages.push(client_message);
    }

    // Sort client messages by nonce.
    client_messages.sort_by(|a, b| a.nonce.cmp(&b.nonce));

    // Handle client messages.
    for client_message in client_messages {
        // The client message is split and handled by the server and verifier.
        let (ciphertext_contribution, decryption_request_contribution) =
            server.split_client_message(client_message).unwrap();
        verifier.verify_and_include(decryption_request_contribution, &mut verifier_state).unwrap();
        server.handle_ciphertext_contribution(ciphertext_contribution, &mut server_state).unwrap();
    }

    // Verifier creates the partial decryption request.
    let pd_ct = verifier.create_partial_decryption_request(verifier_state).unwrap();

    // Decryptors perform partial decryption.
    for i in 0..decryptors.len() {
        let pd = decryptors[i]
            .handle_partial_decryption_request(pd_ct.clone(), &decryptor_states[i])
            .unwrap();
        // Server handles the partial decryption.
        server.handle_partial_decryption(pd, &mut server_state).unwrap();
    }

    // Server recovers the aggregation result.
    let aggregation_result = server.recover_aggregation_result(&server_state).unwrap();

    // Check that the (padded) result matches the client plaintext.
    verify_that!(aggregation_result.keys().collect::<Vec<_>>(), container_eq([&default_id]))?;
    verify_eq!(
        aggregation_result.get(default_id.as_str()).unwrap()[..expected_output.len()],
        expected_output
    )
}