Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions shell_wrapper/status_macros.h
Original file line number Diff line number Diff line change
Expand Up @@ -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_
41 changes: 41 additions & 0 deletions willow/src/shell/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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",
],
)

Expand Down
48 changes: 48 additions & 0 deletions willow/src/shell/parameters_utils.cc
Original file line number Diff line number Diff line change
@@ -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
36 changes: 36 additions & 0 deletions willow/src/shell/parameters_utils.h
Original file line number Diff line number Diff line change
@@ -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_
56 changes: 53 additions & 3 deletions willow/src/shell/parameters_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
46 changes: 46 additions & 0 deletions willow/src/shell/parameters_utils_test.cc
Original file line number Diff line number Diff line change
@@ -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