From ea1fdc6ea252e787bcdcca8a6d1491c44f45b0c1 Mon Sep 17 00:00:00 2001 From: Pierre Tholoniat Date: Tue, 13 Oct 2026 19:08:51 -0700 Subject: [PATCH] Add state serialization to WillowShellCoordinator and CXX FFI bridge. PiperOrigin-RevId: 999041714 --- willow/api/BUILD | 9 ++++ willow/api/coordinator.rs | 87 ++++++++++++++++++++++++++++++++++++++- 2 files changed, 95 insertions(+), 1 deletion(-) diff --git a/willow/api/BUILD b/willow/api/BUILD index b7c10dd..eb314e0 100644 --- a/willow/api/BUILD +++ b/willow/api/BUILD @@ -226,6 +226,15 @@ rust_library( ], ) +rust_test( + name = "coordinator_rust_test", + crate = ":coordinator_rust", + deps = [ + "@crate_index//:googletest", + "//willow/testing_utils", + ], +) + rust_cxx_bridge( name = "coordinator_cxx", src = "coordinator.rs", diff --git a/willow/api/coordinator.rs b/willow/api/coordinator.rs index 9830a2d..2c4b0f4 100644 --- a/willow/api/coordinator.rs +++ b/willow/api/coordinator.rs @@ -18,6 +18,7 @@ use ahe_traits::AheBase; use coordinator::Coordinator; use kahe_traits::{HasKahe, KaheBase}; use messages::CoordinatorState; +use messages_rust_proto::CoordinatorState as CoordinatorStateProto; use proto_serialization_traits::{FromProto, ToProto}; use protobuf::prelude::*; use shell_ciphertexts_rust_proto::ShellAhePartialDecCiphertext as ShellAhePartialDecCiphertextProto; @@ -56,6 +57,18 @@ pub mod ffi { ptr: *mut WillowShellCoordinator, ) -> Box; + #[cxx_name = "ToSerializedState"] + fn to_serialized_state_ffi( + self: &WillowShellCoordinator, + out_serialized_state: &mut Vec, + ) -> FfiStatus; + + #[cxx_name = "RestoreFromSerializedState"] + fn restore_from_serialized_state_ffi( + self: &mut WillowShellCoordinator, + serialized_state: &[u8], + ) -> FfiStatus; + #[cxx_name = "HandleSetupSubmissions"] fn handle_setup_submissions_ffi( self: &mut WillowShellCoordinator, @@ -84,6 +97,7 @@ pub struct WillowShellCoordinator { kahe: Rc, coord: Coordinator, coord_state: CoordinatorState, + aggregation_config: AggregationConfig, } impl HasKahe for WillowShellCoordinator { @@ -126,13 +140,57 @@ impl WillowShellCoordinator { status::internal(&format!("Failed to parse AggregationConfigProto: {}", e)) })?; let aggregation_config = AggregationConfig::from_proto(aggregation_config_proto, ())?; + let (kahe_config, ahe_config) = create_shell_configs(&aggregation_config)?; let context_bytes = &aggregation_config.key_id; let kahe = Rc::new(ShellKahe::new(kahe_config, context_bytes)?); let vahe = Rc::new(ShellVahe::new(ahe_config, context_bytes)?); let coord = Coordinator { vahe }; let coord_state = CoordinatorState::default(); - Ok(WillowShellCoordinator { kahe, coord, coord_state }) + + Ok(WillowShellCoordinator { kahe, coord, coord_state, aggregation_config }) + } + + fn to_serialized_state(&self) -> Result, StatusError> { + let mut proto = self.coord_state.to_proto(self)?; + proto.set_aggregation_config(self.aggregation_config.to_proto(())?); + proto + .serialize() + .map_err(|e| status::internal(&format!("Failed to serialize CoordinatorState: {}", e))) + } + + fn to_serialized_state_ffi(&self, out_serialized_state: &mut Vec) -> ffi::FfiStatus { + self.to_serialized_state() + .map(|serialized| { + *out_serialized_state = serialized; + }) + .into() + } + + /// Restores the protocol state of this coordinator from a serialized CoordinatorState proto. + /// + /// The AggregationConfig embedded in `state` must match the one this coordinator was created + /// with; otherwise an `InvalidArgument` error is returned and the coordinator is unchanged. + fn restore_from_serialized_state(&mut self, state: &[u8]) -> Result<(), StatusError> { + let state_proto = CoordinatorStateProto::parse(state).map_err(|e| { + status::invalid_argument(&format!("Failed to parse CoordinatorState: {}", e)) + })?; + if !state_proto.has_aggregation_config() { + return Err(status::invalid_argument("CoordinatorState is missing aggregation_config")); + } + let embedded_config = AggregationConfig::from_proto(state_proto.aggregation_config(), ())?; + if embedded_config != self.aggregation_config { + return Err(status::invalid_argument( + "CoordinatorState aggregation_config does not match the coordinator's \ + aggregation_config", + )); + } + self.coord_state = CoordinatorState::from_proto(state_proto, self)?; + Ok(()) + } + + fn restore_from_serialized_state_ffi(&mut self, serialized_state: &[u8]) -> ffi::FfiStatus { + self.restore_from_serialized_state(serialized_state).into() } /// Helper function to parse protobuf views using the coordinator's context. @@ -243,3 +301,30 @@ impl WillowShellCoordinator { .into() } } + +#[cfg(test)] +mod tests { + use super::*; + use googletest::prelude::*; + use status::StatusErrorCode; + use testing_utils::generate_aggregation_config; + + fn new_coordinator(config: &AggregationConfig) -> Result { + let serialized_config = config.to_proto(())?.serialize().map_err(|e| { + status::internal(&format!("Failed to serialize AggregationConfigProto: {}", e)) + })?; + WillowShellCoordinator::new_from_serialized_config(&serialized_config) + } + + #[gtest] + fn restore_from_serialized_state_fails_with_mismatched_config() -> googletest::Result<()> { + let config = generate_aggregation_config("default".to_string(), 16, 10, 1, 1); + let serialized_state = new_coordinator(&config)?.to_serialized_state()?; + + let mut other_config = config.clone(); + other_config.key_id = b"other_key_id".to_vec(); + let mut coordinator = new_coordinator(&other_config)?; + let result = coordinator.restore_from_serialized_state(&serialized_state); + verify_true!(matches!(result.map_err(|e| e.code()), Err(StatusErrorCode::INVALID_ARGUMENT))) + } +}