Expose Coordinator state serialization and reinstantiation in C++ API.

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",
     ],
 )
 

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

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<std::unique_ptr<Coordinator>> Coordinator::Create(
   return std::unique_ptr<Coordinator>(new Coordinator(std::move(coord_box)));
 }
 
+absl::StatusOr<std::unique_ptr<Coordinator>>
+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<std::string> Coordinator::ToSerializedState() const {
+  rust::Vec<uint8_t> serialized_state;
+  SECAGG_RETURN_IF_FFI_ERROR(coordinator_->ToSerializedState(serialized_state));
+  return std::string(reinterpret_cast<const char*>(serialized_state.data()),
+                     serialized_state.size());
+}
+
 absl::StatusOr<VerifyKeyContributionsRequest>
 Coordinator::HandleSetupSubmissions(
     absl::Span<const SetupContribution> non_reputable_contributions,

willow/api/coordinator.h

@@ -16,9 +16,11 @@
 #define SECURE_AGGREGATION_WILLOW_API_COORDINATOR_H_
 
 #include <memory>
+#include <string>
 #include <utility>
 
 #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<std::unique_ptr<Coordinator>> 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<std::unique_ptr<Coordinator>> CreateFromSerializedState(
+      absl::string_view serialized_state);
+
+  // Converts the current state of the coordinator to a serialized
+  // CoordinatorState proto.
+  absl::StatusOr<std::string> 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.

willow/api/coordinator_test.cc

@@ -15,21 +15,29 @@
 #include "willow/api/coordinator.h"
 
 #include <memory>
+#include <string>
+#include <utility>
 #include <vector>
 
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "ffi_utils/status_matchers.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #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<SetupContribution> non_reputable = {std::move(contrib)};
+  std::vector<SetupContribution> 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) {