diff --git a/shell_wrapper/status_macros.h b/shell_wrapper/status_macros.h index a3ec64a..fb63d66 100644 --- a/shell_wrapper/status_macros.h +++ b/shell_wrapper/status_macros.h @@ -49,4 +49,9 @@ return status; \ } +// Internal helper to handle results from Rust FFI calls, which return a +// secure_aggregation::FfiStatus instead of a absl::Status. +#define SECAGG_RETURN_IF_FFI_ERROR(expr) \ + SECAGG_RETURN_IF_ERROR(secure_aggregation::UnwrapFfiStatus(expr)) + #endif // SECURE_AGGREGATION_SHELL_WRAPPER_STATUS_MACROS_H_ diff --git a/willow/src/shell/BUILD b/willow/src/shell/BUILD index 8efd481..29309ce 100644 --- a/willow/src/shell/BUILD +++ b/willow/src/shell/BUILD @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +load("@cxx.rs//tools/bazel:rust_cxx_bridge.bzl", "rust_cxx_bridge") +load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") load("@rules_rust//rust:defs.bzl", "rust_library", "rust_test") package( @@ -113,10 +115,49 @@ rust_library( crate_root = "parameters_utils.rs", deps = [ ":kahe_shell", + ":parameters_shell", "@protobuf//rust:protobuf", + "@cxx.rs//:cxx", "//shell_wrapper:kahe", "//shell_wrapper:status", "//willow/proto/shell:shell_parameters_rust_proto", + "//willow/proto/willow:aggregation_config_rust_proto", + "//willow/src/api:aggregation_config", + "//willow/src/traits:proto_serialization_traits", + ], +) + +rust_cxx_bridge( + name = "shell_parameters_utils_cxx", + src = "parameters_utils.rs", + deps = [ + ":shell_parameters_utils", + "//shell_wrapper:status_cxx", + ], +) + +cc_library( + name = "shell_parameters_utils_cc", + srcs = ["parameters_utils.cc"], + hdrs = ["parameters_utils.h"], + deps = [ + ":shell_parameters_utils_cxx", + "@abseil-cpp//absl/status", + "@abseil-cpp//absl/status:statusor", + "@cxx.rs//:core", + "//shell_wrapper:status_cc", + "//shell_wrapper:status_macros", + "//willow/proto/willow:aggregation_config_cc_proto", + ], +) + +cc_test( + name = "parameters_utils_test", + srcs = ["parameters_utils_test.cc"], + deps = [ + ":shell_parameters_utils_cc", + "@googletest//:gtest_main", + "//willow/proto/willow:aggregation_config_cc_proto", ], ) diff --git a/willow/src/shell/parameters_utils.cc b/willow/src/shell/parameters_utils.cc new file mode 100644 index 0000000..f91b911 --- /dev/null +++ b/willow/src/shell/parameters_utils.cc @@ -0,0 +1,48 @@ +/* + * 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/src/shell/parameters_utils.h" + +#include +#include +#include +#include + +#include "absl/status/statusor.h" +#include "include/cxx.h" +#include "shell_wrapper/status_macros.h" +#include "willow/proto/willow/aggregation_config.pb.h" +#include "willow/src/shell/parameters_utils.rs.h" + +namespace secure_aggregation { +namespace willow { + +absl::StatusOr CreateHumanReadableShellConfig( + const AggregationConfigProto& config) { + std::string serialized_config = config.SerializeAsString(); + rust::Vec result; + + SECAGG_RETURN_IF_FFI_ERROR( + secure_aggregation::create_human_readable_shell_config( + std::make_unique(std::move(serialized_config)), + &result)); + + return std::string(reinterpret_cast(result.data()), + result.size()); +} + +} // namespace willow +} // namespace secure_aggregation diff --git a/willow/src/shell/parameters_utils.h b/willow/src/shell/parameters_utils.h new file mode 100644 index 0000000..020791d --- /dev/null +++ b/willow/src/shell/parameters_utils.h @@ -0,0 +1,36 @@ +/* + * 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_SRC_SHELL_PARAMETER_UTILS_H_ +#define SECURE_AGGREGATION_WILLOW_SRC_SHELL_PARAMETER_UTILS_H_ + +#include + +#include "absl/status/statusor.h" +#include "willow/proto/willow/aggregation_config.pb.h" + +namespace secure_aggregation { +namespace willow { + +// Returns the ShellKaheConfig and ShellAheConfig as a human-readable string, +// for the given AggregationConfigProto. +absl::StatusOr CreateHumanReadableShellConfig( + const AggregationConfigProto& config); + +} // namespace willow +} // namespace secure_aggregation + +#endif // SECURE_AGGREGATION_WILLOW_SRC_SHELL_PARAMETER_UTILS_H_ diff --git a/willow/src/shell/parameters_utils.rs b/willow/src/shell/parameters_utils.rs index cf3bd11..76a9d95 100644 --- a/willow/src/shell/parameters_utils.rs +++ b/willow/src/shell/parameters_utils.rs @@ -12,17 +12,67 @@ // See the License for the specific language governing permissions and // limitations under the License. +//! This file contains some utility functions for working with Willow parameters: +//! - Conversions between Rust structs and their corresponding protos. +//! - FFI bridge to generate and print Shell parameters from C++. + +use aggregation_config::AggregationConfig; +use aggregation_config_rust_proto::AggregationConfigProto; +use cxx::{CxxString, UniquePtr}; use kahe::PackedVectorConfig; use kahe_shell::ShellKaheConfig; -use protobuf::{proto, ProtoStr}; +use parameters_shell::create_shell_configs; +use proto_serialization_traits::FromProto; +use protobuf::{proto, Parse, ProtoStr}; use shell_parameters_rust_proto::{ PackedVectorConfigProto, PackedVectorConfigProtoView, ShellKaheConfigProto, ShellKaheConfigProtoView, }; use std::collections::BTreeMap; -/// This file contains some utility functions for working with Willow parameters: -/// - Conversions between Rust structs and their corresponding protos. +#[cxx::bridge] +pub mod ffi { + // Re-define FfiStatus since CXX requires shared structs to be defined in the same module + // (https://github.com/dtolnay/cxx/issues/297#issuecomment-727042059) + unsafe extern "C++" { + include!("shell_wrapper/status.rs.h"); + type FfiStatus = status::ffi::FfiStatus; + } + + #[namespace = "secure_aggregation"] + extern "Rust" { + unsafe fn create_human_readable_shell_config( + aggregation_config_proto: UniquePtr, + out: *mut Vec, + ) -> FfiStatus; + } +} + +fn create_human_readable_shell_config_impl( + aggregation_config_proto: UniquePtr, +) -> Result, status::StatusError> { + let config_proto = + AggregationConfigProto::parse(aggregation_config_proto.as_bytes()).map_err(|e| { + status::invalid_argument(format!("Failed to parse AggregationConfigProto: {}", e)) + })?; + let config = AggregationConfig::from_proto(config_proto, ())?; + let (kahe_config, ahe_config) = create_shell_configs(&config)?; + let kahe_config_string = format!("{:#?}", kahe_config); + let ahe_config_string = format!("{:#?}", ahe_config); + let result = + format!("ShellKaheConfig: {}\nShellAheConfig: {}", kahe_config_string, ahe_config_string); + Ok(result.into_bytes()) +} + +/// SAFETY: `out` must not be null. +unsafe fn create_human_readable_shell_config( + aggregation_config_proto: UniquePtr, + out: *mut Vec, +) -> ffi::FfiStatus { + create_human_readable_shell_config_impl(aggregation_config_proto) + .map(|result| *out = result) + .into() +} /// Convert a rust struct `PackedVectorConfig` to the corresponding proto. pub fn packed_vector_config_to_proto(config: &PackedVectorConfig) -> PackedVectorConfigProto { diff --git a/willow/src/shell/parameters_utils_test.cc b/willow/src/shell/parameters_utils_test.cc new file mode 100644 index 0000000..fd78607 --- /dev/null +++ b/willow/src/shell/parameters_utils_test.cc @@ -0,0 +1,46 @@ +/* + * 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/src/shell/parameters_utils.h" + +#include "gmock/gmock.h" +#include "gtest/gtest.h" +#include "willow/proto/willow/aggregation_config.pb.h" + +namespace secure_aggregation { +namespace willow { +namespace { + +TEST(ParameterUtilsTest, CreateHumanReadableShellConfigTest) { + 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_session_id("test_session"); + + auto result = CreateHumanReadableShellConfig(config); + + ASSERT_TRUE(result.ok()); + EXPECT_THAT(*result, ::testing::HasSubstr("ShellKaheConfig")); + EXPECT_THAT(*result, ::testing::HasSubstr("ShellAheConfig")); +} + +} // namespace +} // namespace willow +} // namespace secure_aggregation