From bf3eb8647fefac07fa7893bc1dee753294d45e43 Mon Sep 17 00:00:00 2001 From: Pierre Tholoniat Date: Wed, 14 Oct 2026 01:37:56 -0700 Subject: [PATCH] Expose Coordinator state serialization and reinstantiation in C++ API. PiperOrigin-RevId: 999203662 --- willow/api/BUILD | 3 ++ willow/api/api_e2e_test.cc | 24 ++++++++++ willow/api/coordinator.cc | 27 +++++++++++ willow/api/coordinator.h | 12 +++++ willow/api/coordinator_test.cc | 88 ++++++++++++++++++++++++++++++++-- 5 files changed, 151 insertions(+), 3 deletions(-) diff --git a/willow/api/BUILD b/willow/api/BUILD index eb314e0..2f65499 100644 --- a/willow/api/BUILD +++ b/willow/api/BUILD @@ -252,6 +252,7 @@ cc_library( ":coordinator_cxx", "@abseil-cpp//absl/status", "@abseil-cpp//absl/status:statusor", + "@abseil-cpp//absl/strings:string_view", "@abseil-cpp//absl/types:span", "@cxx.rs//:core", "//ffi_utils:cxx_utils", @@ -268,11 +269,13 @@ cc_test( deps = [ ":coordinator", "@googletest//:gtest_main", + "@abseil-cpp//absl/status", "@abseil-cpp//absl/status:statusor", "//ffi_utils:status_matchers", "//willow/proto/shell:shell_ciphertexts_cc_proto", "//willow/proto/willow:aggregation_config_cc_proto", "//willow/proto/willow:messages_cc_proto", + "//willow/testing_utils:shell_testing_decryptor_cc", ], ) diff --git a/willow/api/api_e2e_test.cc b/willow/api/api_e2e_test.cc index d698ade..d32d771 100644 --- a/willow/api/api_e2e_test.cc +++ b/willow/api/api_e2e_test.cc @@ -105,6 +105,13 @@ TEST(WillowV1ApiE2ETest, MultiDecryptorCoordinatorFullPipelineWithCodec) { coordinator->HandleSetupSubmissions( non_reputable_contribs, reputable_contribs)); + // Serialize and reinstantiate Coordinator after setup round. + SECAGG_ASSERT_OK_AND_ASSIGN(std::string serialized_coord_state, + coordinator->ToSerializedState()); + SECAGG_ASSERT_OK_AND_ASSIGN( + coordinator, + Coordinator::CreateFromSerializedState(serialized_coord_state)); + // Reputable decryptor verifies contributions and aggregates global public // key. SECAGG_ASSERT_OK_AND_ASSIGN( @@ -187,6 +194,13 @@ TEST(WillowV1ApiE2ETest, MultiDecryptorCoordinatorFullPipelineWithCodec) { coordinator->PrepareDecryptionRequest( acc_dec_req.partial_dec_ciphertext())); + // Serialize and reinstantiate Coordinator while awaiting partial decryptions. + SECAGG_ASSERT_OK_AND_ASSIGN(serialized_coord_state, + coordinator->ToSerializedState()); + SECAGG_ASSERT_OK_AND_ASSIGN( + coordinator, + Coordinator::CreateFromSerializedState(serialized_coord_state)); + // Multi-decryptors each produce a partial decryption response. SECAGG_ASSERT_OK_AND_ASSIGN( PartialDecryptionResponse response_reputable, @@ -208,6 +222,16 @@ TEST(WillowV1ApiE2ETest, MultiDecryptorCoordinatorFullPipelineWithCodec) { FinalizedPartialDecryption finalized_pd, coordinator->AggregateAndFinalizePartialDecryptions(partial_responses)); + // Serialize and reinstantiate Coordinator in OutputReady state. + SECAGG_ASSERT_OK_AND_ASSIGN(serialized_coord_state, + coordinator->ToSerializedState()); + SECAGG_ASSERT_OK_AND_ASSIGN( + coordinator, + Coordinator::CreateFromSerializedState(serialized_coord_state)); + SECAGG_ASSERT_OK_AND_ASSIGN(std::string roundtrip_output_ready_state, + coordinator->ToSerializedState()); + EXPECT_EQ(roundtrip_output_ready_state, serialized_coord_state); + // Recover final result using FinalResultDecryptor. // In production, this code runs within the pipeline, using the serialized // final_result_decryptor_state produced during Finalize() to decrypt the diff --git a/willow/api/coordinator.cc b/willow/api/coordinator.cc index d4ad31f..4c28d1a 100644 --- a/willow/api/coordinator.cc +++ b/willow/api/coordinator.cc @@ -22,6 +22,7 @@ #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/string_view.h" #include "absl/types/span.h" #include "ffi_utils/cxx_utils.h" #include "ffi_utils/status_macros.h" @@ -66,6 +67,32 @@ absl::StatusOr> Coordinator::Create( return std::unique_ptr(new Coordinator(std::move(coord_box))); } +absl::StatusOr> +Coordinator::CreateFromSerializedState(absl::string_view serialized_state) { + CoordinatorState state_proto; + if (!state_proto.ParseFromString(serialized_state)) { + return absl::InvalidArgumentError( + "Failed to parse CoordinatorState from serialized bytes."); + } + if (!state_proto.has_aggregation_config()) { + return absl::InvalidArgumentError( + "CoordinatorState is missing aggregation_config."); + } + SECAGG_ASSIGN_OR_RETURN( + auto coordinator, Coordinator::Create(state_proto.aggregation_config())); + SECAGG_RETURN_IF_FFI_ERROR( + coordinator->coordinator_->RestoreFromSerializedState( + ToRustSlice(serialized_state))); + return coordinator; +} + +absl::StatusOr Coordinator::ToSerializedState() const { + rust::Vec serialized_state; + SECAGG_RETURN_IF_FFI_ERROR(coordinator_->ToSerializedState(serialized_state)); + return std::string(reinterpret_cast(serialized_state.data()), + serialized_state.size()); +} + absl::StatusOr Coordinator::HandleSetupSubmissions( absl::Span non_reputable_contributions, diff --git a/willow/api/coordinator.h b/willow/api/coordinator.h index ae36906..1dfab21 100644 --- a/willow/api/coordinator.h +++ b/willow/api/coordinator.h @@ -16,9 +16,11 @@ #define SECURE_AGGREGATION_WILLOW_API_COORDINATOR_H_ #include +#include #include #include "absl/status/statusor.h" +#include "absl/strings/string_view.h" #include "absl/types/span.h" #include "include/cxx.h" #include "willow/api/coordinator.rs.h" @@ -37,6 +39,16 @@ class Coordinator { static absl::StatusOr> Create( const AggregationConfigProto& aggregation_config); + // Creates a new coordinator from the given serialized state, which must + // correspond to a serialized CoordinatorState proto produced by + // ToSerializedState() (including its embedded AggregationConfig). + static absl::StatusOr> CreateFromSerializedState( + absl::string_view serialized_state); + + // Converts the current state of the coordinator to a serialized + // CoordinatorState proto. + absl::StatusOr ToSerializedState() const; + // Stores setup contributions from all decryptors during key generation and // creates a request for reputable decryptors to verify and aggregate key // contributions into a public key. diff --git a/willow/api/coordinator_test.cc b/willow/api/coordinator_test.cc index 450a082..e798e6f 100644 --- a/willow/api/coordinator_test.cc +++ b/willow/api/coordinator_test.cc @@ -15,8 +15,11 @@ #include "willow/api/coordinator.h" #include +#include +#include #include +#include "absl/status/status.h" #include "absl/status/statusor.h" #include "ffi_utils/status_matchers.h" #include "gmock/gmock.h" @@ -24,12 +27,17 @@ #include "willow/proto/shell/ciphertexts.pb.h" #include "willow/proto/willow/aggregation_config.pb.h" #include "willow/proto/willow/messages.pb.h" +#include "willow/testing_utils/shell_testing_decryptor.h" namespace secure_aggregation { namespace willow { namespace { +using ::secure_aggregation::secagg_internal::StatusIs; +using ::secure_aggregation::testing::ShellTestingDecryptor; +using ::testing::HasSubstr; using ::testing::IsEmpty; +using ::testing::SizeIs; AggregationConfigProto CreateValidConfig() { AggregationConfigProto config; @@ -64,9 +72,83 @@ TEST(CoordinatorTest, HandleSetupSubmissionsSucceedsAndTransitionsState) { // Calling HandleSetupSubmissions a second time should fail because the // coordinator is no longer in PreSetup status. - auto second_attempt = - coordinator->HandleSetupSubmissions(non_reputable, reputable); - EXPECT_FALSE(second_attempt.ok()); + EXPECT_THAT( + coordinator->HandleSetupSubmissions(non_reputable, reputable), + StatusIs(absl::StatusCode::kFailedPrecondition, HasSubstr("PreSetup"))); +} + +TEST(CoordinatorTest, SerializeAndReinstantiatePreservesState) { + AggregationConfigProto config = CreateValidConfig(); + SECAGG_ASSERT_OK_AND_ASSIGN(auto coordinator, Coordinator::Create(config)); + + // Serialize and reinstantiate in PreSetup state. + SECAGG_ASSERT_OK_AND_ASSIGN(std::string pre_setup_state, + coordinator->ToSerializedState()); + SECAGG_ASSERT_OK_AND_ASSIGN( + coordinator, Coordinator::CreateFromSerializedState(pre_setup_state)); + + SECAGG_ASSERT_OK_AND_ASSIGN(auto decryptor, + ShellTestingDecryptor::Create(config)); + SECAGG_ASSERT_OK_AND_ASSIGN(SetupContribution contrib, + decryptor->CreateSetupContribution()); + contrib.add_encrypted_randomness_shares()->set_encrypted_share("test_share"); + + std::vector non_reputable = {std::move(contrib)}; + std::vector reputable; + SECAGG_ASSERT_OK_AND_ASSIGN( + auto verify_request, + coordinator->HandleSetupSubmissions(non_reputable, reputable)); + EXPECT_THAT(verify_request.key_contributions(), SizeIs(1)); + + // Serialize and reinstantiate in KeySharesReceived state with populated + // encrypted_randomness_shares. + SECAGG_ASSERT_OK_AND_ASSIGN(std::string post_setup_state, + coordinator->ToSerializedState()); + SECAGG_ASSERT_OK_AND_ASSIGN( + coordinator, Coordinator::CreateFromSerializedState(post_setup_state)); + SECAGG_ASSERT_OK_AND_ASSIGN(std::string roundtrip_post_setup_state, + coordinator->ToSerializedState()); + EXPECT_EQ(roundtrip_post_setup_state, post_setup_state); + + CoordinatorState parsed_state; + ASSERT_TRUE(parsed_state.ParseFromString(roundtrip_post_setup_state)); + EXPECT_EQ(parsed_state.status(), COORDINATOR_STATUS_KEY_SHARES_RECEIVED); + ASSERT_THAT(parsed_state.encrypted_randomness_shares(), SizeIs(1)); + ASSERT_THAT(parsed_state.encrypted_randomness_shares(0).shares(), SizeIs(1)); + EXPECT_EQ( + parsed_state.encrypted_randomness_shares(0).shares(0).encrypted_share(), + "test_share"); + + // Reinstantiated coordinator must preserve KeySharesReceived status and + // reject a duplicate HandleSetupSubmissions call. + EXPECT_THAT( + coordinator->HandleSetupSubmissions(non_reputable, reputable), + StatusIs(absl::StatusCode::kFailedPrecondition, HasSubstr("PreSetup"))); +} + +TEST(CoordinatorTest, CreateFromSerializedStateFailsWithInvalidBytes) { + EXPECT_THAT(Coordinator::CreateFromSerializedState("not_a_valid_proto"), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Failed to parse CoordinatorState"))); +} + +TEST(CoordinatorTest, CreateFromSerializedStateFailsWithoutAggregationConfig) { + CoordinatorState state_without_config; + state_without_config.set_status(COORDINATOR_STATUS_PRE_SETUP); + EXPECT_THAT(Coordinator::CreateFromSerializedState( + state_without_config.SerializeAsString()), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("missing aggregation_config"))); +} + +TEST(CoordinatorTest, CreateFromSerializedStateFailsWithUnspecifiedStatus) { + CoordinatorState state_with_unspecified_status; + *state_with_unspecified_status.mutable_aggregation_config() = + CreateValidConfig(); + EXPECT_THAT(Coordinator::CreateFromSerializedState( + state_with_unspecified_status.SerializeAsString()), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("CoordinatorStatus"))); } TEST(CoordinatorTest, PrepareDecryptionRequestFailsWhenNotInCorrectState) {