Add C++ utility to generate human-readable Shell parameters from a Willow aggregation config proto; use FfiStatus for interop.

tholop · copybara-github · commit 798d5e4b58cb · 2026-01-22T07:55:34.000-08:00
PiperOrigin-RevId: 857168135
diff --git 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
@@ -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,50 @@ 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",
+        "@abseil-cpp//absl/status",
+        "//willow/proto/willow:aggregation_config_cc_proto",
     ],
 )
 
diff --git a/willow/src/shell/parameters_utils.cc 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 <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+
+#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<std::string> CreateHumanReadableShellConfig(
+    const AggregationConfigProto& config) {
+  std::string serialized_config = config.SerializeAsString();
+  rust::Vec<uint8_t> result;
+
+  SECAGG_RETURN_IF_FFI_ERROR(
+      secure_aggregation::create_human_readable_shell_config(
+          std::make_unique<std::string>(std::move(serialized_config)),
+          &result));
+
+  return std::string(reinterpret_cast<const char*>(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
@@ -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 <string>
+
+#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<std::string> 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
@@ -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<CxxString>,
+            out: *mut Vec<u8>,
+        ) -> FfiStatus;
+    }
+}
+
+fn create_human_readable_shell_config_impl(
+    aggregation_config_proto: UniquePtr<CxxString>,
+) -> Result<Vec<u8>, 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<CxxString>,
+    out: *mut Vec<u8>,
+) -> 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
@@ -0,0 +1,61 @@
+/*
+ * 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 "absl/status/status.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"));
+}
+
+TEST(ParameterUtilsTest, CreateHumanReadableShellConfigInvalidConfigTest) {
+  AggregationConfigProto 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);
+
+  EXPECT_EQ(result.status().code(), absl::StatusCode::kInvalidArgument);
+  EXPECT_THAT(
+      result.status().message(),
+      ::testing::HasSubstr("empty vector configs in aggregation config"));
+}
+
+}  // namespace
+}  // namespace willow
+}  // namespace secure_aggregation

Original file line number	Diff line number	Diff line change
`@@ -49,4 +49,9 @@`
`49`	`49`	`return status; \`
`50`	`50`	`}`
`51`	`51`
	`52`	`+// Internal helper to handle results from Rust FFI calls, which return a`
	`53`	`+// secure_aggregation::FfiStatus instead of a absl::Status.`
	`54`	`+#define SECAGG_RETURN_IF_FFI_ERROR(expr) \`
	`55`	`+ SECAGG_RETURN_IF_ERROR(secure_aggregation::UnwrapFfiStatus(expr))`
	`56`	`+`
`52`	`57`	`#endif // SECURE_AGGREGATION_SHELL_WRAPPER_STATUS_MACROS_H_`