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_
8 changes: 3 additions & 5 deletions willow/benches/shell_benchmarks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -131,9 +131,7 @@ fn setup_base(args: &Args) -> BaseInputs {
// Create client.
let kahe = ShellKahe::new(kahe_config.clone(), CONTEXT_STRING).unwrap();
let vahe = ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap();
let seed = SingleThreadHkdfPrng::generate_seed().unwrap();
let prng = SingleThreadHkdfPrng::create(&seed).unwrap();
let client = WillowV1Client { kahe, vahe, prng };
let client = WillowV1Client::new_from_randomly_generated_seed(kahe, vahe).unwrap();

// Create decryptor.
let vahe = ShellVahe::new(ahe_config.clone(), CONTEXT_STRING).unwrap();
Expand Down Expand Up @@ -218,7 +216,7 @@ struct VerifierInputs {
}

fn setup_verifier_verify_client_message(args: &Args) -> VerifierInputs {
let mut inputs = setup_base(args);
let inputs = setup_base(args);
let mut decryption_request_contributions = vec![];
for _ in 0..args.n_iterations {
// Generates a plaintext and encrypts.
Expand Down Expand Up @@ -257,7 +255,7 @@ fn run_verifier_verify_client_message(inputs: &mut VerifierInputs) {
}

fn setup_server_handle_client_message(args: &Args) -> ServerInputs {
let mut inputs = setup_base(args);
let inputs = setup_base(args);
let mut ciphertext_contributions = vec![];
for _ in 0..args.n_iterations {
// Generates a plaintext and encrypts.
Expand Down
4 changes: 1 addition & 3 deletions willow/src/api/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -86,9 +86,7 @@ impl WillowShellClient {
let context_bytes = aggregation_config.compute_context_bytes()?;
let kahe = ShellKahe::new(kahe_config, &context_bytes)?;
let vahe = ShellVahe::new(ahe_config, &context_bytes)?;
let client_seed = SingleThreadHkdfPrng::generate_seed()?;
let prng = SingleThreadHkdfPrng::create(&client_seed)?;
let client = WillowV1Client { kahe, vahe, prng };
let client = WillowV1Client::new_from_randomly_generated_seed(kahe, vahe)?;
Ok(WillowShellClient(client))
}

Expand Down
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
2 changes: 1 addition & 1 deletion willow/src/traits/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ pub trait SecureAggregationClient: HasKahe + HasVahe {
/// Creates a client message to be sent to the Server.
/// nonce is used for the VAHE encryption, has to be unique.
fn create_client_message(
&mut self,
&self,
plaintext: &Self::PlaintextSlice<'_>,
signed_public_key: &DecryptorPublicKey<<Self as HasVahe>::Vahe>,
nonce: &[u8],
Expand Down
Loading