|
12 | 12 | // See the License for the specific language governing permissions and |
13 | 13 | // limitations under the License. |
14 | 14 |
|
| 15 | +//! This file contains some utility functions for working with Willow parameters: |
| 16 | +//! - Conversions between Rust structs and their corresponding protos. |
| 17 | +//! - FFI bridge to generate and print Shell parameters from C++. |
| 18 | +
|
| 19 | +use aggregation_config::AggregationConfig; |
| 20 | +use aggregation_config_rust_proto::AggregationConfigProto; |
| 21 | +use cxx::{CxxString, UniquePtr}; |
15 | 22 | use kahe::PackedVectorConfig; |
16 | 23 | use kahe_shell::ShellKaheConfig; |
17 | | -use protobuf::{proto, ProtoStr}; |
| 24 | +use parameters_shell::create_shell_configs; |
| 25 | +use proto_serialization_traits::FromProto; |
| 26 | +use protobuf::{proto, Parse, ProtoStr}; |
18 | 27 | use shell_parameters_rust_proto::{ |
19 | 28 | PackedVectorConfigProto, PackedVectorConfigProtoView, ShellKaheConfigProto, |
20 | 29 | ShellKaheConfigProtoView, |
21 | 30 | }; |
22 | 31 | use std::collections::BTreeMap; |
23 | 32 |
|
24 | | -/// This file contains some utility functions for working with Willow parameters: |
25 | | -/// - Conversions between Rust structs and their corresponding protos. |
| 33 | +#[cxx::bridge] |
| 34 | +pub mod ffi { |
| 35 | + // Re-define FfiStatus since CXX requires shared structs to be defined in the same module |
| 36 | + // (https://github.com/dtolnay/cxx/issues/297#issuecomment-727042059) |
| 37 | + unsafe extern "C++" { |
| 38 | + include!("shell_wrapper/status.rs.h"); |
| 39 | + type FfiStatus = status::ffi::FfiStatus; |
| 40 | + } |
| 41 | + |
| 42 | + #[namespace = "secure_aggregation"] |
| 43 | + extern "Rust" { |
| 44 | + unsafe fn create_human_readable_shell_config( |
| 45 | + aggregation_config_proto: UniquePtr<CxxString>, |
| 46 | + out: *mut Vec<u8>, |
| 47 | + ) -> FfiStatus; |
| 48 | + } |
| 49 | +} |
| 50 | + |
| 51 | +fn create_human_readable_shell_config_impl( |
| 52 | + aggregation_config_proto: UniquePtr<CxxString>, |
| 53 | +) -> Result<Vec<u8>, status::StatusError> { |
| 54 | + let config_proto = |
| 55 | + AggregationConfigProto::parse(aggregation_config_proto.as_bytes()).map_err(|e| { |
| 56 | + status::invalid_argument(format!("Failed to parse AggregationConfigProto: {}", e)) |
| 57 | + })?; |
| 58 | + let config = AggregationConfig::from_proto(config_proto, ())?; |
| 59 | + let (kahe_config, ahe_config) = create_shell_configs(&config)?; |
| 60 | + let kahe_config_string = format!("{:#?}", kahe_config); |
| 61 | + let ahe_config_string = format!("{:#?}", ahe_config); |
| 62 | + let result = |
| 63 | + format!("ShellKaheConfig: {}\nShellAheConfig: {}", kahe_config_string, ahe_config_string); |
| 64 | + Ok(result.into_bytes()) |
| 65 | +} |
| 66 | + |
| 67 | +/// SAFETY: `out` must not be null. |
| 68 | +unsafe fn create_human_readable_shell_config( |
| 69 | + aggregation_config_proto: UniquePtr<CxxString>, |
| 70 | + out: *mut Vec<u8>, |
| 71 | +) -> ffi::FfiStatus { |
| 72 | + create_human_readable_shell_config_impl(aggregation_config_proto) |
| 73 | + .map(|result| *out = result) |
| 74 | + .into() |
| 75 | +} |
26 | 76 |
|
27 | 77 | /// Convert a rust struct `PackedVectorConfig` to the corresponding proto. |
28 | 78 | pub fn packed_vector_config_to_proto(config: &PackedVectorConfig) -> PackedVectorConfigProto { |
|
0 commit comments