diff --git a/willow/api/BUILD b/willow/api/BUILD index 97a384b..f65a27d 100644 --- a/willow/api/BUILD +++ b/willow/api/BUILD @@ -204,3 +204,72 @@ cc_test( "//willow/testing_utils:testing_utils_cc", ], ) + +rust_library( + name = "coordinator_rust", + srcs = ["coordinator.rs"], + deps = [ + ":aggregation_config", + ":proto_serialization_traits", + "@protobuf//rust:protobuf", + "@cxx.rs//:cxx", + "//ffi_utils:status", + "//willow/crypto:ahe_traits", + "//willow/crypto:kahe_traits", + "//willow/crypto:shell_kahe", + "//willow/crypto:shell_parameters", + "//willow/crypto:shell_vahe", + "//willow/crypto:vahe_traits", + "//willow/proto/shell:shell_ciphertexts_rust_proto", + "//willow/proto/willow:aggregation_config_rust_proto", + "//willow/proto/willow:messages_rust_proto", + "//willow/protocol:decryptor_traits", + "//willow/protocol:messages", + "//willow/protocol:willow_v1_coordinator", + ], +) + +rust_cxx_bridge( + name = "coordinator_cxx", + src = "coordinator.rs", + deps = [ + ":coordinator_rust", + "//ffi_utils:status_cxx", + ], +) + +cc_library( + name = "coordinator", + srcs = ["coordinator.cc"], + hdrs = ["coordinator.h"], + deps = [ + ":coordinator_cxx", + ":coordinator_cxx/include", # fixdeps: keep + ":coordinator_rust", # fixdeps: keep + "@abseil-cpp//absl/status", + "@abseil-cpp//absl/status:statusor", + "@abseil-cpp//absl/types:span", + "@cxx.rs//:core", + "//ffi_utils:cxx_utils", + "//ffi_utils:status_macros", + "//willow/proto/shell:shell_ciphertexts_cc_proto", + "//willow/proto/willow:aggregation_config_cc_proto", + "//willow/proto/willow:messages_cc_proto", + ], +) + +cc_test( + name = "coordinator_test", + srcs = ["coordinator_test.cc"], + deps = [ + ":coordinator", + ":coordinator_cxx", # fixdeps: keep + "@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", + ], +) diff --git a/willow/api/coordinator.cc b/willow/api/coordinator.cc new file mode 100644 index 0000000..d4ad31f --- /dev/null +++ b/willow/api/coordinator.cc @@ -0,0 +1,130 @@ +// Copyright 2026 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. + +#include "willow/api/coordinator.h" + +#include +#include +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/types/span.h" +#include "ffi_utils/cxx_utils.h" +#include "ffi_utils/status_macros.h" +#include "include/cxx.h" +#include "willow/api/coordinator.rs.h" +#include "willow/proto/shell/ciphertexts.pb.h" +#include "willow/proto/willow/aggregation_config.pb.h" +#include "willow/proto/willow/messages.pb.h" + +namespace secure_aggregation { +namespace willow { +namespace { + +// Helper to serialize an array of C++ protobufs and pass them +// across the FFI boundary as a slice of byte views. +template +struct SerializedProtoSpan { + std::vector strs; + std::vector views; + rust::Slice slice; + + explicit SerializedProtoSpan(absl::Span protos) { + strs.reserve(protos.size()); + views.reserve(protos.size()); + for (const auto& proto : protos) { + strs.push_back(proto.SerializeAsString()); + views.push_back(SerializedProtoView{ToRustSlice(strs.back())}); + } + slice = rust::Slice(views.data(), views.size()); + } +}; + +} // namespace + +absl::StatusOr> Coordinator::Create( + const AggregationConfigProto& aggregation_config) { + std::string config_str = aggregation_config.SerializeAsString(); + secure_aggregation::WillowShellCoordinator* coord_ptr = nullptr; + SECAGG_RETURN_IF_FFI_ERROR(NewWillowShellCoordinatorFromSerializedConfig( + ToRustSlice(config_str), &coord_ptr)); + auto coord_box = WillowShellCoordinatorIntoBox(coord_ptr); + return std::unique_ptr(new Coordinator(std::move(coord_box))); +} + +absl::StatusOr +Coordinator::HandleSetupSubmissions( + absl::Span non_reputable_contributions, + absl::Span reputable_contributions) { + SerializedProtoSpan non_reputable_helper( + non_reputable_contributions); + SerializedProtoSpan reputable_helper( + reputable_contributions); + + rust::Vec result_bytes; + // We only pass slices to Rust, which point to views of strings owned by + // SerializedProtoSpan + SECAGG_RETURN_IF_FFI_ERROR(coordinator_->HandleSetupSubmissions( + non_reputable_helper.slice, reputable_helper.slice, result_bytes)); + + VerifyKeyContributionsRequest verify_request; + if (!verify_request.ParseFromArray(result_bytes.data(), + result_bytes.size())) { + return absl::InternalError( + "Failed to parse VerifyKeyContributionsRequest from serialized FFI " + "bytes."); + } + return verify_request; +} + +absl::StatusOr Coordinator::PrepareDecryptionRequest( + const ShellAhePartialDecCiphertext& verifier_ciphertext) { + std::string ct_str = verifier_ciphertext.SerializeAsString(); + rust::Vec result_bytes; + SECAGG_RETURN_IF_FFI_ERROR(coordinator_->PrepareDecryptionRequest( + ToRustSlice(ct_str), result_bytes)); + + PartialDecryptionRequest dec_request; + if (!dec_request.ParseFromArray(result_bytes.data(), result_bytes.size())) { + return absl::InternalError( + "Failed to parse PartialDecryptionRequest from serialized FFI bytes."); + } + return dec_request; +} + +absl::StatusOr +Coordinator::AggregateAndFinalizePartialDecryptions( + absl::Span partial_responses) { + SerializedProtoSpan responses_helper( + partial_responses); + + rust::Vec result_bytes; + SECAGG_RETURN_IF_FFI_ERROR( + coordinator_->AggregateAndFinalizePartialDecryptions( + responses_helper.slice, result_bytes)); + + FinalizedPartialDecryption finalized_pd; + if (!finalized_pd.ParseFromArray(result_bytes.data(), result_bytes.size())) { + return absl::InternalError( + "Failed to parse FinalizedPartialDecryption from serialized FFI " + "bytes."); + } + return finalized_pd; +} + +} // namespace willow +} // namespace secure_aggregation diff --git a/willow/api/coordinator.h b/willow/api/coordinator.h new file mode 100644 index 0000000..ae36906 --- /dev/null +++ b/willow/api/coordinator.h @@ -0,0 +1,69 @@ +// Copyright 2026 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. + +#ifndef SECURE_AGGREGATION_WILLOW_API_COORDINATOR_H_ +#define SECURE_AGGREGATION_WILLOW_API_COORDINATOR_H_ + +#include +#include + +#include "absl/status/statusor.h" +#include "absl/types/span.h" +#include "include/cxx.h" +#include "willow/api/coordinator.rs.h" +#include "willow/proto/shell/ciphertexts.pb.h" +#include "willow/proto/willow/aggregation_config.pb.h" +#include "willow/proto/willow/messages.pb.h" + +namespace secure_aggregation { +namespace willow { + +// Implements a C++ wrapper around the untrusted multi-decryptor Coordinator. +// Manages protocol flow and aggregates messages from all decryptors. +class Coordinator { + public: + // Creates a new coordinator with the given aggregation_config. + static absl::StatusOr> Create( + const AggregationConfigProto& aggregation_config); + + // 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. + absl::StatusOr HandleSetupSubmissions( + absl::Span non_reputable_contributions, + absl::Span reputable_contributions); + + // Combines the verifier ciphertext half (obtained from the accumulation + // pipeline) with the AHE DP noise component sum received at setup, to prepare + // a decryption request for decryptors. + absl::StatusOr PrepareDecryptionRequest( + const ShellAhePartialDecCiphertext& verifier_ciphertext); + + // Accumulates partial decryptions and finalizes the partial decryption sum. + absl::StatusOr + AggregateAndFinalizePartialDecryptions( + absl::Span partial_responses); + + private: + explicit Coordinator( + rust::Box coordinator) + : coordinator_(std::move(coordinator)) {} + + rust::Box coordinator_; +}; + +} // namespace willow +} // namespace secure_aggregation + +#endif // SECURE_AGGREGATION_WILLOW_API_COORDINATOR_H_ diff --git a/willow/api/coordinator.rs b/willow/api/coordinator.rs new file mode 100644 index 0000000..070b7aa --- /dev/null +++ b/willow/api/coordinator.rs @@ -0,0 +1,246 @@ +// Copyright 2026 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 aggregation_config::AggregationConfig; +use aggregation_config_rust_proto::AggregationConfigProto; +use ahe_traits::AheBase; +use decryptor_traits::SecureAggregationCoordinator; +use kahe_traits::{HasKahe, KaheBase}; +use messages::CoordinatorState; +use proto_serialization_traits::{FromProto, ToProto}; +use protobuf::prelude::*; +use shell_ciphertexts_rust_proto::ShellAhePartialDecCiphertext as ShellAhePartialDecCiphertextProto; +use shell_kahe::ShellKahe; +use shell_parameters::create_shell_configs; +use shell_vahe::ShellVahe; +use status::StatusError; +use std::rc::Rc; +use vahe_traits::HasVahe; +use willow_v1_coordinator::WillowV1Coordinator; + +#[cxx::bridge(namespace = "secure_aggregation")] +pub mod ffi { + // CXX doesn't support nested types, so we define a struct here so that functions can take + // an array of protobuf views. The slice is over values owned by the C++ + // SerializedProtoSpan object. + struct SerializedProtoView<'a> { + data: &'a [u8], + } + + unsafe extern "C++" { + include!("ffi_utils/status.rs.h"); + type FfiStatus = status::ffi::FfiStatus; + } + + extern "Rust" { + type WillowShellCoordinator; + + #[cxx_name = "NewWillowShellCoordinatorFromSerializedConfig"] + unsafe fn new_willow_shell_coordinator_from_serialized_config( + serialized_aggregation_config: &[u8], + out: *mut *mut WillowShellCoordinator, + ) -> FfiStatus; + + #[cxx_name = "WillowShellCoordinatorIntoBox"] + unsafe fn willow_shell_coordinator_into_box( + ptr: *mut WillowShellCoordinator, + ) -> Box; + + #[cxx_name = "HandleSetupSubmissions"] + fn handle_setup_submissions_ffi( + self: &mut WillowShellCoordinator, + non_reputable_contributions: &[SerializedProtoView], + reputable_contributions: &[SerializedProtoView], + out_verify_request: &mut Vec, + ) -> FfiStatus; + + #[cxx_name = "PrepareDecryptionRequest"] + fn prepare_decryption_request_ffi( + self: &mut WillowShellCoordinator, + verifier_ciphertext: &[u8], + out_decryption_request: &mut Vec, + ) -> FfiStatus; + + #[cxx_name = "AggregateAndFinalizePartialDecryptions"] + fn aggregate_and_finalize_partial_decryptions_ffi( + self: &mut WillowShellCoordinator, + partial_responses: &[SerializedProtoView], + out_finalized_pd: &mut Vec, + ) -> FfiStatus; + } +} + +pub struct WillowShellCoordinator { + kahe: Rc, + coord: WillowV1Coordinator, + coord_state: CoordinatorState, +} + +impl HasKahe for WillowShellCoordinator { + type Kahe = ShellKahe; + fn kahe(&self) -> &Self::Kahe { + &self.kahe + } +} + +impl HasVahe for WillowShellCoordinator { + type Vahe = ShellVahe; + fn vahe(&self) -> &Self::Vahe { + self.coord.vahe() + } +} + +/// Converts a raw pointer to a Box. Ideally we would use `rust::Box::from_raw` +/// (https://cxx.rs/binding/box.html) directly from C++, but that causes linker errors. +/// +/// SAFETY: `ptr` must have been created by Box::into_raw or one of the functions in this module. +unsafe fn willow_shell_coordinator_into_box( + ptr: *mut WillowShellCoordinator, +) -> Box { + unsafe { Box::from_raw(ptr) } +} + +/// SAFETY: `out` must be a valid pointer for writes. +unsafe fn new_willow_shell_coordinator_from_serialized_config( + serialized_aggregation_config: &[u8], + out: *mut *mut WillowShellCoordinator, +) -> ffi::FfiStatus { + WillowShellCoordinator::new_from_serialized_config(serialized_aggregation_config) + .map(|coord| unsafe { *out = Box::into_raw(Box::new(coord)) }) + .into() +} + +impl WillowShellCoordinator { + fn new_from_serialized_config(config: &[u8]) -> Result { + let aggregation_config_proto = AggregationConfigProto::parse(config).map_err(|e| { + 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 = WillowV1Coordinator { vahe }; + let coord_state = CoordinatorState::default(); + Ok(WillowShellCoordinator { kahe, coord, coord_state }) + } + + /// Helper function to parse protobuf views using the coordinator's context. + fn parse_views<'a, T>( + &'a self, + views: &[ffi::SerializedProtoView], + ) -> Result, StatusError> + where + T: FromProto<&'a WillowShellCoordinator>, + T::Proto: protobuf::Message, + { + let mut result = Vec::with_capacity(views.len()); + for view in views { + let proto = ::parse(view.data) + .map_err(|e| status::internal(&format!("Failed to parse proto view: {}", e)))?; + result.push(T::from_proto(proto, self)?); + } + Ok(result) + } + + fn handle_setup_submissions( + &mut self, + non_reputable_views: &[ffi::SerializedProtoView], + reputable_views: &[ffi::SerializedProtoView], + ) -> Result, StatusError> { + let non_reputable = self.parse_views(non_reputable_views)?; + let reputable = self.parse_views(reputable_views)?; + + let verify_request = + self.coord.handle_setup_submissions(non_reputable, reputable, &mut self.coord_state)?; + let proto = verify_request.to_proto(&*self)?; + proto.serialize().map_err(|e| { + status::internal(&format!("Failed to serialize VerifyKeyContributionsRequest: {}", e)) + }) + } + + fn handle_setup_submissions_ffi( + &mut self, + non_reputable_contributions: &[ffi::SerializedProtoView], + reputable_contributions: &[ffi::SerializedProtoView], + out_verify_request: &mut Vec, + ) -> ffi::FfiStatus { + self.handle_setup_submissions(non_reputable_contributions, reputable_contributions) + .map(|serialized| { + *out_verify_request = serialized; + }) + .into() + } + + fn prepare_decryption_request( + &mut self, + verifier_ciphertext: &[u8], + ) -> Result, StatusError> { + let ct_proto = + ShellAhePartialDecCiphertextProto::parse(verifier_ciphertext).map_err(|e| { + status::internal(&format!("Failed to parse ShellAhePartialDecCiphertext: {}", e)) + })?; + let ct = + ::PartialDecCiphertext::from_proto(ct_proto, self.coord.vahe())?; + + let request = self.coord.prepare_decryption_request(&ct, &mut self.coord_state)?; + let proto = request.to_proto(&*self)?; + proto.serialize().map_err(|e| { + status::internal(&format!("Failed to serialize PartialDecryptionRequest: {}", e)) + }) + } + + fn prepare_decryption_request_ffi( + &mut self, + verifier_ciphertext: &[u8], + out_decryption_request: &mut Vec, + ) -> ffi::FfiStatus { + self.prepare_decryption_request(verifier_ciphertext) + .map(|serialized| { + *out_decryption_request = serialized; + }) + .into() + } + + fn aggregate_and_finalize_partial_decryptions( + &mut self, + partial_response_views: &[ffi::SerializedProtoView], + ) -> Result, StatusError> { + let responses = self.parse_views(partial_response_views)?; + + self.coord.aggregate_partial_decryptions( + responses, + Some(self.kahe.as_ref()), + &mut self.coord_state, + )?; + + let finalized_pd = self.coord.finalize_partial_decryption(&mut self.coord_state)?; + let proto = finalized_pd.to_proto(&*self)?; + proto.serialize().map_err(|e| { + status::internal(&format!("Failed to serialize FinalizedPartialDecryption: {}", e)) + }) + } + + fn aggregate_and_finalize_partial_decryptions_ffi( + &mut self, + partial_responses: &[ffi::SerializedProtoView], + out_finalized_pd: &mut Vec, + ) -> ffi::FfiStatus { + self.aggregate_and_finalize_partial_decryptions(partial_responses) + .map(|serialized| { + *out_finalized_pd = serialized; + }) + .into() + } +} diff --git a/willow/api/coordinator_test.cc b/willow/api/coordinator_test.cc new file mode 100644 index 0000000..450a082 --- /dev/null +++ b/willow/api/coordinator_test.cc @@ -0,0 +1,95 @@ +// Copyright 2026 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. + +#include "willow/api/coordinator.h" + +#include +#include + +#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" + +namespace secure_aggregation { +namespace willow { +namespace { + +using ::testing::IsEmpty; + +AggregationConfigProto CreateValidConfig() { + AggregationConfigProto config; + VectorConfig vector_config; + vector_config.set_length(10); + vector_config.set_bound(100); + (*config.mutable_vector_configs())["test_vector"] = vector_config; + config.set_max_number_of_decryptors(1); + config.set_max_number_of_clients(10); + config.set_key_id("test_key"); + return config; +} + +TEST(CoordinatorTest, CreateSucceedsWithValidConfig) { + AggregationConfigProto config = CreateValidConfig(); + auto coordinator_or = Coordinator::Create(config); + ASSERT_TRUE(coordinator_or.ok()) << coordinator_or.status(); + EXPECT_NE(*coordinator_or, nullptr); +} + +TEST(CoordinatorTest, HandleSetupSubmissionsSucceedsAndTransitionsState) { + AggregationConfigProto config = CreateValidConfig(); + SECAGG_ASSERT_OK_AND_ASSIGN(auto coordinator, Coordinator::Create(config)); + + // Handle setup submissions with empty contributions for state testing. + std::vector non_reputable; + std::vector reputable; + SECAGG_ASSERT_OK_AND_ASSIGN( + auto verify_request, + coordinator->HandleSetupSubmissions(non_reputable, reputable)); + EXPECT_THAT(verify_request.key_contributions(), IsEmpty()); + + // 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()); +} + +TEST(CoordinatorTest, PrepareDecryptionRequestFailsWhenNotInCorrectState) { + AggregationConfigProto config = CreateValidConfig(); + SECAGG_ASSERT_OK_AND_ASSIGN(auto coordinator, Coordinator::Create(config)); + + // PrepareDecryptionRequest without calling HandleSetupSubmissions first + // should fail. + ShellAhePartialDecCiphertext dummy_ct; + auto request_or = coordinator->PrepareDecryptionRequest(dummy_ct); + EXPECT_FALSE(request_or.ok()); +} + +TEST(CoordinatorTest, AggregateAndFinalizeFailsWhenNotInCorrectState) { + AggregationConfigProto config = CreateValidConfig(); + SECAGG_ASSERT_OK_AND_ASSIGN(auto coordinator, Coordinator::Create(config)); + + std::vector responses; + auto finalize_or = + coordinator->AggregateAndFinalizePartialDecryptions(responses); + EXPECT_FALSE(finalize_or.ok()); +} + +} // namespace +} // namespace willow +} // namespace secure_aggregation diff --git a/willow/protocol/willow_v1_coordinator.rs b/willow/protocol/willow_v1_coordinator.rs index 70a88bb..b72b05d 100644 --- a/willow/protocol/willow_v1_coordinator.rs +++ b/willow/protocol/willow_v1_coordinator.rs @@ -125,7 +125,6 @@ where &self, _coordinator_state: &mut Self::CoordinatorState, ) -> Result, StatusError> { - // Dropout recovery is not yet implemented. Err(status::unimplemented("Dropout recovery is not yet implemented")) } @@ -134,7 +133,6 @@ where _recovery_responses: Vec, _coordinator_state: &mut Self::CoordinatorState, ) -> Result<(), StatusError> { - // Dropout recovery is not yet implemented. Err(status::unimplemented("Dropout recovery is not yet implemented")) }