diff --git a/willow/protocol/BUILD b/willow/protocol/BUILD index 194ce1f..3a4f377 100644 --- a/willow/protocol/BUILD +++ b/willow/protocol/BUILD @@ -263,6 +263,8 @@ rust_test( ":decryptor", ":messages", "@crate_index//:googletest", + "@protobuf//rust:protobuf", + "//willow/api:proto_serialization_traits", "//willow/crypto:ahe_traits", "//willow/crypto:prng_traits", "//willow/crypto:shell_kahe", @@ -270,6 +272,7 @@ rust_test( "//willow/crypto:shell_vahe", "//willow/crypto:single_thread_hkdf", "//willow/crypto:vahe_traits", + "//willow/proto/willow:messages_rust_proto", ], ) diff --git a/willow/protocol/coordinator.rs b/willow/protocol/coordinator.rs index d30510b..398a26a 100644 --- a/willow/protocol/coordinator.rs +++ b/willow/protocol/coordinator.rs @@ -112,7 +112,7 @@ impl Coordinator { let mut partial_responses_iter = partial_responses.into_iter(); let mut sum = partial_responses_iter .next() - .expect("No partial decryptions provided") + .ok_or_else(|| status::invalid_argument("No partial decryptions provided"))? .partial_decryption; for response in partial_responses_iter { self.vahe.add_partial_decryptions_in_place(&response.partial_decryption, &mut sum)?; @@ -163,7 +163,7 @@ impl Coordinator { partial_decryption_sum: coordinator_state .partial_decryption_sum .clone() - .expect("partial_decryption_sum should be set"), + .ok_or_else(|| status::failed_precondition("partial_decryption_sum is not set"))?, }) } } @@ -176,8 +176,11 @@ mod tests { use decryptor::{Decryptor, DecryptorState}; use googletest::gtest; use googletest::prelude::*; - use messages::{CoordinatorState, CoordinatorStatus}; + use messages::{CoordinatorState, CoordinatorStatus, SecretSharingContribution}; + use messages_rust_proto::CoordinatorStatus as CoordinatorStatusProto; use prng_traits::SecurePrng; + use proto_serialization_traits::{FromProto, ToProto}; + use protobuf::prelude::*; use shell_kahe::ShellKahe; use shell_parameters::create_shell_ahe_config; use shell_vahe::ShellVahe; @@ -242,8 +245,29 @@ mod tests { Ok(()) } + #[gtest] + fn aggregate_partial_decryptions_fails_with_no_responses() -> googletest::Result<()> { + let vahe = Rc::new(ShellVahe::new(create_shell_ahe_config(1)?, CONTEXT_STRING)?); + let coordinator = Coordinator { vahe }; + let mut coord_state = CoordinatorState::default(); + coord_state.status = CoordinatorStatus::AwaitingPartialDecryptions; + + let result = + coordinator.aggregate_partial_decryptions::(vec![], None, &mut coord_state); + verify_true!(matches!( + result.map_err(|e| e.code()), + Err(status::StatusErrorCode::INVALID_ARGUMENT) + ))?; + verify_that!(coord_state.status, eq(CoordinatorStatus::AwaitingPartialDecryptions))?; + verify_true!(coord_state.partial_decryption_sum.is_none())?; + + Ok(()) + } + /// End-to-end test: setup -> encryption -> partial decryption -> recovery - /// using the multi-decryptor protocol with a coordinator and reputable decryptor. + /// using the multi-decryptor protocol with a coordinator, one non-reputable decryptor and + /// one reputable decryptor, + /// serializing and deserializing CoordinatorState at every protocol transition. #[gtest] fn end_to_end_multi_decryptor_protocol() -> googletest::Result<()> { let vahe = Rc::new(ShellVahe::new(create_shell_ahe_config(1)?, CONTEXT_STRING)?); @@ -261,15 +285,35 @@ mod tests { let contribution1 = decryptor1.create_setup_contribution(&mut dec_state1)?; let contribution2 = decryptor2.create_setup_contribution(&mut dec_state2)?; - // Coordinator processes setup. + // Attach encrypted randomness shares to non-reputable decryptor1's contribution. + let mut contribution1 = contribution1; + contribution1.encrypted_randomness_shares = Some(vec![SecretSharingContribution { + encrypted_share: b"test_encrypted_share".to_vec(), + }]); + + // Coordinator processes setup (round-trip default PreSetup state first). let coordinator = Coordinator { vahe: vahe.clone() }; - let mut coord_state = CoordinatorState::default(); + let coord_state = CoordinatorState::default(); + let coord_state_proto = coord_state.to_proto(&coordinator)?; + let mut coord_state = CoordinatorState::from_proto(coord_state_proto, &coordinator)?; + verify_that!(coord_state.status, eq(CoordinatorStatus::PreSetup))?; + let verify_request = coordinator.handle_setup_submissions( - vec![], - vec![contribution1, contribution2], + vec![contribution1], + vec![contribution2], &mut coord_state, )?; + // Round-trip CoordinatorState after setup submissions. + let coord_state_proto = coord_state.to_proto(&coordinator)?; + let mut coord_state = CoordinatorState::from_proto(coord_state_proto, &coordinator)?; + verify_that!(coord_state.status, eq(CoordinatorStatus::KeySharesReceived))?; + verify_that!(coord_state.encrypted_randomness_shares.len(), eq(1))?; + verify_that!( + coord_state.encrypted_randomness_shares[0][0].encrypted_share, + eq(b"test_encrypted_share") + )?; + // Reputable decryptor verifies and aggregates the public key. let public_key = decryptor1.verify_and_aggregate_key_contributions(verify_request)?; @@ -290,6 +334,11 @@ mod tests { let pd_request = coordinator.prepare_decryption_request(&partial_dec_ciphertext, &mut coord_state)?; + // Round-trip CoordinatorState after preparing decryption request. + let coord_state_proto = coord_state.to_proto(&coordinator)?; + let mut coord_state = CoordinatorState::from_proto(coord_state_proto, &coordinator)?; + verify_that!(coord_state.status, eq(CoordinatorStatus::AwaitingPartialDecryptions))?; + // Each decryptor computes a partial decryption. let pd_response1: messages::PartialDecryptionResponse = decryptor1 .handle_partial_decryption_request(pd_request.clone(), None, &mut dec_state1)?; @@ -303,7 +352,10 @@ mod tests { &mut coord_state, )?; - verify_true!(coord_state.status == CoordinatorStatus::OutputReady)?; + // Round-trip CoordinatorState after aggregating partial decryptions. + let coord_state_proto = coord_state.to_proto(&coordinator)?; + let mut coord_state = CoordinatorState::from_proto(coord_state_proto, &coordinator)?; + verify_that!(coord_state.status, eq(CoordinatorStatus::OutputReady))?; // Recover the plaintext using the finalized state. let finalized_state = coordinator.finalize_partial_decryption(&mut coord_state)?; @@ -317,4 +369,100 @@ mod tests { Ok(()) } + + #[gtest] + fn coordinator_state_statuses_roundtrip() -> googletest::Result<()> { + let vahe = Rc::new(ShellVahe::new(create_shell_ahe_config(1)?, CONTEXT_STRING)?); + let coordinator = Coordinator { vahe }; + + for status in [ + CoordinatorStatus::PreSetup, + CoordinatorStatus::KeySharesReceived, + CoordinatorStatus::AwaitingContributions, + CoordinatorStatus::AwaitingPartialDecryptions, + CoordinatorStatus::AwaitingRecovery, + CoordinatorStatus::Finished, + ] { + let state = CoordinatorState { status, ..CoordinatorState::default() }; + let proto = state.to_proto(&coordinator)?; + let roundtrip = CoordinatorState::from_proto(proto.as_view(), &coordinator)?; + verify_that!(roundtrip.status, eq(status))?; + verify_that!(roundtrip.to_proto(&coordinator)?.serialize()?, eq(&proto.serialize()?))?; + } + + Ok(()) + } + + #[gtest] + fn coordinator_state_rejects_invalid_statuses_and_invariants() -> googletest::Result<()> { + let vahe = Rc::new(ShellVahe::new(create_shell_ahe_config(1)?, CONTEXT_STRING)?); + let coordinator = Coordinator { vahe }; + + // Unspecified CoordinatorStatusProto must fail deserialization with InvalidArgument. + let err = CoordinatorStatus::from_proto(CoordinatorStatusProto::Unspecified, ()) + .expect_err("expected error for Unspecified CoordinatorStatus"); + verify_that!(err.message(), contains_substring("CoordinatorStatus"))?; + + // OutputReady without partial_decryption_sum must fail deserialization. + let invalid_output_ready = CoordinatorState { + status: CoordinatorStatus::OutputReady, + ..CoordinatorState::default() + }; + let invalid_proto = invalid_output_ready.to_proto(&coordinator)?; + let err = CoordinatorState::from_proto(invalid_proto, &coordinator) + .err() + .expect("expected error for OutputReady without partial_decryption_sum"); + verify_that!(err.message(), contains_substring("partial_decryption_sum"))?; + + // PreSetup with populated encrypted_randomness_shares must fail deserialization. + let invalid_pre_setup = CoordinatorState { + status: CoordinatorStatus::PreSetup, + encrypted_randomness_shares: vec![vec![SecretSharingContribution { + encrypted_share: b"unexpected_share".to_vec(), + }]], + ..CoordinatorState::default() + }; + let invalid_pre_setup_proto = invalid_pre_setup.to_proto(&coordinator)?; + let err = CoordinatorState::from_proto(invalid_pre_setup_proto, &coordinator) + .err() + .expect("expected error for PreSetup with populated fields"); + verify_that!(err.message(), contains_substring("PreSetup"))?; + + Ok(()) + } + + #[gtest] + fn coordinator_state_optional_fields_roundtrip() -> googletest::Result<()> { + let vahe = Rc::new(ShellVahe::new(create_shell_ahe_config(1)?, CONTEXT_STRING)?); + let decryptor = Decryptor::new_with_randomly_generated_seed(vahe.clone())?; + let mut dec_state = DecryptorState::default(); + let contribution = decryptor.create_setup_contribution(&mut dec_state)?; + let coordinator = Coordinator { vahe: vahe.clone() }; + + let public_key = vahe.aggregate_public_key_shares(std::iter::once( + &contribution.key_contribution.public_key_share, + ))?; + let seed = SingleThreadHkdfPrng::generate_seed()?; + let mut prng = SingleThreadHkdfPrng::create(&seed)?; + let (ciphertext, _proof) = + vahe.verifiable_encrypt(&vec![1i64; 8], &public_key, b"0123456789ABCDEF", &mut prng)?; + let dp_noise_ct = vahe.get_partial_dec_ciphertext(&ciphertext)?; + + let state = CoordinatorState { + status: CoordinatorStatus::KeySharesReceived, + encrypted_randomness_shares: vec![vec![SecretSharingContribution { + encrypted_share: b"share_bytes".to_vec(), + }]], + dp_noise_component_sum: Some(dp_noise_ct), + setup_contributions: Some(vec![contribution]), + partial_decryption_sum: None, + }; + let proto = state.to_proto(&coordinator)?; + let roundtrip = CoordinatorState::from_proto(proto.as_view(), &coordinator)?; + verify_true!(roundtrip.dp_noise_component_sum.is_some())?; + verify_that!(roundtrip.setup_contributions.as_ref().map(|v| v.len()), eq(Some(1)))?; + verify_that!(roundtrip.to_proto(&coordinator)?.serialize()?, eq(&proto.serialize()?))?; + + Ok(()) + } } diff --git a/willow/protocol/messages.rs b/willow/protocol/messages.rs index b47a17c..08c911e 100644 --- a/willow/protocol/messages.rs +++ b/willow/protocol/messages.rs @@ -17,8 +17,10 @@ use ahe_traits::AheBase; use kahe_traits::{HasKahe, KaheBase}; use messages_rust_proto::{ CiphertextContribution as CiphertextContributionProto, ClientMessage as ClientMessageProto, + CoordinatorState as CoordinatorStateProto, CoordinatorStatus as CoordinatorStatusProto, DPSetupContribution as DPSetupContributionProto, DecryptionRequestContribution as DecryptionRequestContributionProto, + EncryptedRandomnessShares as EncryptedRandomnessSharesProto, FinalizedPartialDecryption as FinalizedPartialDecryptionProto, KeyContribution as KeyContributionProto, PartialDecryptionRequest as PartialDecryptionRequestProto, @@ -823,3 +825,170 @@ impl Default for CoordinatorState { } } } + +impl ToProto for CoordinatorStatus { + type Proto = CoordinatorStatusProto; + + fn to_proto(&self, _ctx: ()) -> Result { + Ok(match self { + CoordinatorStatus::PreSetup => CoordinatorStatusProto::PreSetup, + CoordinatorStatus::KeySharesReceived => CoordinatorStatusProto::KeySharesReceived, + CoordinatorStatus::AwaitingContributions => { + CoordinatorStatusProto::AwaitingContributions + } + CoordinatorStatus::AwaitingPartialDecryptions => { + CoordinatorStatusProto::AwaitingPartialDecryptions + } + CoordinatorStatus::AwaitingRecovery => CoordinatorStatusProto::AwaitingRecovery, + CoordinatorStatus::OutputReady => CoordinatorStatusProto::OutputReady, + CoordinatorStatus::Finished => CoordinatorStatusProto::Finished, + }) + } +} + +impl FromProto for CoordinatorStatus { + type Proto = CoordinatorStatusProto; + + fn from_proto( + proto: impl AsView, + _ctx: (), + ) -> Result { + match proto.as_view() { + CoordinatorStatusProto::PreSetup => Ok(CoordinatorStatus::PreSetup), + CoordinatorStatusProto::KeySharesReceived => Ok(CoordinatorStatus::KeySharesReceived), + CoordinatorStatusProto::AwaitingContributions => { + Ok(CoordinatorStatus::AwaitingContributions) + } + CoordinatorStatusProto::AwaitingPartialDecryptions => { + Ok(CoordinatorStatus::AwaitingPartialDecryptions) + } + CoordinatorStatusProto::AwaitingRecovery => Ok(CoordinatorStatus::AwaitingRecovery), + CoordinatorStatusProto::OutputReady => Ok(CoordinatorStatus::OutputReady), + CoordinatorStatusProto::Finished => Ok(CoordinatorStatus::Finished), + CoordinatorStatusProto::Unspecified => { + Err(status::invalid_argument("CoordinatorStatus is unspecified")) + } + _ => Err(status::invalid_argument("Unknown CoordinatorStatus")), + } + } +} + +impl<'a, C, Vahe> ToProto<&'a C> for CoordinatorState +where + C: HasVahe, + Vahe: VaheBase + 'a, + Vahe::PublicKeyShare: ToProto<&'a Vahe, Proto = ShellAhePublicKeyShare>, + Vahe::PartialDecCiphertext: ToProto<&'a Vahe, Proto = ShellAhePartialDecCiphertext>, + Vahe::KeyGenProof: ToProto, + Vahe::EncryptionProof: ToProto, + Vahe::PartialDecryption: ToProto<&'a Vahe, Proto = ShellAhePartialDecryption>, +{ + type Proto = CoordinatorStateProto; + + fn to_proto(&self, context: &'a C) -> Result { + let mut proto = proto!(CoordinatorStateProto { status: self.status.to_proto(())? }); + for shares in &self.encrypted_randomness_shares { + let mut shares_proto = proto!(EncryptedRandomnessSharesProto {}); + for share in shares { + shares_proto.shares_mut().push(share.to_proto(())?); + } + proto.encrypted_randomness_shares_mut().push(shares_proto); + } + if let Some(dp_sum) = &self.dp_noise_component_sum { + proto.set_dp_noise_component_sum(dp_sum.to_proto(context.vahe())?); + } + if let Some(contributions) = &self.setup_contributions { + for contribution in contributions { + proto.setup_contributions_mut().push(contribution.to_proto(context)?); + } + } + if let Some(pd_sum) = &self.partial_decryption_sum { + proto.set_partial_decryption_sum(pd_sum.to_proto(context.vahe())?); + } + Ok(proto) + } +} + +impl<'a, C, Vahe> FromProto<&'a C> for CoordinatorState +where + C: HasVahe, + Vahe: VaheBase + 'a, + Vahe::PublicKeyShare: FromProto<&'a Vahe, Proto = ShellAhePublicKeyShare>, + Vahe::PartialDecCiphertext: FromProto<&'a Vahe, Proto = ShellAhePartialDecCiphertext>, + Vahe::KeyGenProof: FromProto, + Vahe::EncryptionProof: FromProto, + Vahe::PartialDecryption: FromProto<&'a Vahe, Proto = ShellAhePartialDecryption>, +{ + type Proto = CoordinatorStateProto; + + fn from_proto( + proto: impl AsView, + context: &'a C, + ) -> Result { + let proto = proto.as_view(); + let status = CoordinatorStatus::from_proto(proto.status(), ())?; + let encrypted_randomness_shares: Result>, StatusError> = proto + .encrypted_randomness_shares() + .iter() + .map(|shares_proto| { + shares_proto + .shares() + .iter() + .map(|s| SecretSharingContribution::from_proto(s, ())) + .collect() + }) + .collect(); + let encrypted_randomness_shares = encrypted_randomness_shares?; + let dp_noise_component_sum = if proto.has_dp_noise_component_sum() { + Some(Vahe::PartialDecCiphertext::from_proto( + proto.dp_noise_component_sum(), + context.vahe(), + )?) + } else { + None + }; + let setup_contributions = { + let contributions: Result, _> = proto + .setup_contributions() + .iter() + .map(|c| SetupContribution::from_proto(c, context)) + .collect(); + let contributions = contributions?; + if contributions.is_empty() { + None + } else { + Some(contributions) + } + }; + let partial_decryption_sum = if proto.has_partial_decryption_sum() { + Some(Vahe::PartialDecryption::from_proto( + proto.partial_decryption_sum(), + context.vahe(), + )?) + } else { + None + }; + if status == CoordinatorStatus::PreSetup + && (!encrypted_randomness_shares.is_empty() + || dp_noise_component_sum.is_some() + || setup_contributions.is_some() + || partial_decryption_sum.is_some()) + { + return Err(status::invalid_argument( + "CoordinatorState in PreSetup status must not have populated protocol fields", + )); + } + if status == CoordinatorStatus::OutputReady && partial_decryption_sum.is_none() { + return Err(status::invalid_argument( + "CoordinatorState in OutputReady status is missing partial_decryption_sum", + )); + } + Ok(CoordinatorState { + status, + encrypted_randomness_shares, + dp_noise_component_sum, + setup_contributions, + partial_decryption_sum, + }) + } +}