Make FfiStatus Clone

ffi_utils/BUILD

@@ -44,7 +44,10 @@ rust_library(
 rust_cxx_bridge(
     name = "status_cxx",
     src = "status.rs",
-    deps = [":status_cc"],
+    deps = [
+        ":status_cc",
+        "@abseil-cpp//absl/status",
+    ],
 )
 
 rust_test(
@@ -60,6 +63,7 @@ cc_library(
     srcs = ["status.cc"],
     hdrs = ["status.h"],
     deps = [
+        ":cxx_utils",
         ":status_cxx/include",
         "@abseil-cpp//absl/status",
         "@abseil-cpp//absl/strings:string_view",

ffi_utils/status.cc

@@ -16,39 +16,59 @@
 
 #include <cstdint>
 #include <memory>
-#include <string>
 #include <utility>
 
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
+#include "ffi_utils/cxx_utils.h"
 #include "ffi_utils/status.rs.h"
 #include "include/cxx.h"
 
 namespace secure_aggregation {
 
-FfiStatus MakeFfiStatus() { return FfiStatus{.code = 0, .message = nullptr}; }
+FfiStatus MakeFfiStatus() { return FfiStatus{nullptr}; }
+
 FfiStatus MakeFfiStatus(absl::Status status) {
-  return FfiStatus{static_cast<int>(status.code()),
-                   std::make_unique<std::string>(status.message())};
+  if (status.ok()) {
+    return MakeFfiStatus();
+  }
+  return FfiStatus{std::make_unique<absl::Status>(std::move(status))};
 }
-FfiStatus MakeFfiStatus(int code, std::string message) {
-  return FfiStatus{code, std::make_unique<std::string>(std::move(message))};
+
+FfiStatus MakeFfiStatus(int32_t code, rust::Slice<const uint8_t> message) {
+  return MakeFfiStatus(absl::Status(static_cast<absl::StatusCode>(code),
+                                    ToAbslStringView(message)));
 }
 
-absl::Status UnwrapFfiStatus(const FfiStatus& status) {
-  if (status.code == 0) {
+FfiStatus CloneFfiStatus(const FfiStatus& status) {
+  if (status.ptr == nullptr) {
+    return MakeFfiStatus();
+  }
+  return MakeFfiStatus(*status.ptr);
+}
+
+absl::Status UnwrapFfiStatus(FfiStatus status) {
+  if (status.ptr == nullptr) {
     return absl::OkStatus();
   }
-  absl::string_view message = "";
-  if (status.message != nullptr || status.message->empty()) {
-    message = *status.message;
+  absl::Status out = std::move(*status.ptr);
+  return out;
+}
+
+int32_t FfiStatusCode(const FfiStatus& status) {
+  if (status.ptr == nullptr) {
+    return 0;
   }
-  return absl::Status(static_cast<absl::StatusCode>(status.code), message);
+  return static_cast<int32_t>(status.ptr->code());
 }
 
-FfiStatus MakeFfiStatus(int32_t code, rust::Slice<const uint8_t> message) {
-  return secure_aggregation::MakeFfiStatus(
-      code, std::string(message.begin(), message.end()));
+rust::Slice<const uint8_t> FfiStatusMessage(const FfiStatus& status) {
+  if (status.ptr == nullptr) {
+    return rust::Slice<const uint8_t>();
+  }
+  return rust::Slice<const uint8_t>(
+      reinterpret_cast<const uint8_t*>(status.ptr->message().data()),
+      status.ptr->message().size());
 }
 
 }  // namespace secure_aggregation

ffi_utils/status.h

@@ -27,12 +27,14 @@ struct FfiStatus;
 
 FfiStatus MakeFfiStatus();
 FfiStatus MakeFfiStatus(absl::Status status);
-FfiStatus MakeFfiStatus(int code, std::string message);
-absl::Status UnwrapFfiStatus(const FfiStatus& status);
-
-// To be called from Rust.
 FfiStatus MakeFfiStatus(int32_t code, rust::Slice<const uint8_t> message);
 
+FfiStatus CloneFfiStatus(const FfiStatus& status);
+
+absl::Status UnwrapFfiStatus(FfiStatus status);
+int32_t FfiStatusCode(const FfiStatus& status);
+rust::Slice<const uint8_t> FfiStatusMessage(const FfiStatus& status);
+
 }  // namespace secure_aggregation
 
 #endif  // SECURE_AGGREGATION_FFI_UTILS_STATUS_H_

ffi_utils/status.rs

@@ -24,21 +24,42 @@ use std::borrow::Cow;
 
 #[cxx::bridge(namespace = "secure_aggregation")]
 pub mod ffi {
-    // A simple Status alternative which is cxx-compatible (because it directly uses unique_ptr).
+    unsafe extern "C++" {
+        include!("absl/status/status.h");
+        #[namespace = "absl"]
+        type Status;
+    }
+
+    // A simple Status wrapper which is cxx-compatible (because it directly uses unique_ptr).
     pub struct FfiStatus {
-        pub code: i32,
-        pub message: UniquePtr<CxxString>,
+        // Wrapped absl::Status. A nullptr is interpreted as an OK status.
+        ptr: UniquePtr<Status>,
     }
 
     unsafe extern "C++" {
         include!("ffi_utils/status.h");
+        #[rust_name = "make_ok_ffi_status"]
+        pub fn MakeFfiStatus() -> FfiStatus;
+        #[rust_name = "make_ffi_status"]
         pub fn MakeFfiStatus(code: i32, message: &[u8]) -> FfiStatus;
+        #[rust_name = "ffi_status_code"]
+        pub fn FfiStatusCode(status: &FfiStatus) -> i32;
+        #[rust_name = "ffi_status_message"]
+        pub fn FfiStatusMessage<'a>(status: &'a FfiStatus) -> &'a [u8];
+        #[rust_name = "clone_ffi_status"]
+        pub fn CloneFfiStatus(status: &FfiStatus) -> FfiStatus;
     }
 }
 
 pub type Status = Result<(), StatusError>;
 pub type StatusOr<T> = Result<T, StatusError>;
 
+impl Clone for ffi::FfiStatus {
+    fn clone(&self) -> Self {
+        ffi::clone_ffi_status(self)
+    }
+}
+
 /// All cases of C++ StatusErrorCode except `StatusErrorCode::kOk`.
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
 #[non_exhaustive]
@@ -240,26 +261,27 @@ impl std::error::Error for StatusErrorCodeTryFromError {}
 
 impl From<StatusError> for ffi::FfiStatus {
     fn from(error: StatusError) -> Self {
-        ffi::MakeFfiStatus(error.code as i32, error.message.0.as_slice())
+        ffi::make_ffi_status(error.code as i32, error.message.0.as_slice())
     }
 }
 
 impl From<Status> for ffi::FfiStatus {
     fn from(status: Status) -> Self {
         match status {
-            Ok(()) => ffi::FfiStatus { code: 0, message: cxx::UniquePtr::null() },
+            Ok(()) => ffi::make_ok_ffi_status(),
             Err(error) => error.into(),
         }
     }
 }
 
 pub fn rust_status_from_cpp(status: ffi::FfiStatus) -> Status {
-    if status.code == 0 {
+    let code = ffi::ffi_status_code(&status);
+    if code == 0 {
         Ok(())
     } else {
-        let message = if status.message.is_null() { b"" } else { status.message.as_bytes() };
+        let message = ffi::ffi_status_message(&status);
         Err(StatusError::new(
-            status.code.try_into().unwrap_or(StatusErrorCode::Unknown),
+            code.try_into().unwrap_or(StatusErrorCode::Unknown),
             message,
             core::panic::Location::caller(),
         ))
@@ -369,6 +391,8 @@ impl std::fmt::Debug for MaybeString {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use ffi::ffi_status_code;
+    use ffi::ffi_status_message;
     use googletest::prelude::*;
 
     #[allow(dead_code)]
@@ -396,6 +420,14 @@ mod tests {
         }
     }
 
+    #[gtest]
+    fn test_clone() {
+        let status = ffi::make_ffi_status(1, b"test");
+        let status2 = status.clone();
+        expect_eq!(ffi_status_code(&status2), 1);
+        expect_eq!(ffi_status_message(&status2), b"test");
+    }
+
     #[gtest]
     fn test_try_from() {
         for i in 1..=16 {
@@ -415,13 +447,13 @@ mod tests {
     fn test_ffi_status_from_status_error() {
         let error = StatusError::new_untracked(StatusErrorCode::Cancelled, "test");
         let ffi_status: ffi::FfiStatus = error.into();
-        expect_eq!(ffi_status.code, 1);
-        expect_eq!(ffi_status.message.as_bytes(), b"test");
+        expect_eq!(ffi_status_code(&ffi_status), 1);
+        expect_eq!(ffi_status_message(&ffi_status), b"test");
     }
 
     #[gtest]
     fn test_rust_status_from_cpp() {
-        let ffi_status = ffi::MakeFfiStatus(1, b"test");
+        let ffi_status = ffi::make_ffi_status(1, b"test");
         let rust_status = rust_status_from_cpp(ffi_status);
         assert!(rust_status.is_err());
         expect_eq!(&rust_status.as_ref().err().unwrap().code(), &StatusErrorCode::Cancelled);
@@ -432,15 +464,15 @@ mod tests {
     fn test_ffi_status_from_ok_status() {
         let rust_status = Ok(());
         let ffi_status: ffi::FfiStatus = rust_status.into();
-        expect_eq!(ffi_status.code, 0);
-        expect_eq!(ffi_status.message.is_null(), true);
+        expect_eq!(ffi_status_code(&ffi_status), 0);
+        expect_eq!(ffi_status_message(&ffi_status).is_empty(), true);
     }
 
     #[gtest]
     fn test_ffi_status_from_non_ok_status() {
         let rust_status = Err(StatusError::new_untracked(StatusErrorCode::Cancelled, "test"));
         let ffi_status: ffi::FfiStatus = rust_status.into();
-        expect_eq!(ffi_status.code, StatusErrorCode::Cancelled as i32);
-        expect_eq!(ffi_status.message.as_bytes(), b"test");
+        expect_eq!(ffi_status_code(&ffi_status), StatusErrorCode::Cancelled as i32);
+        expect_eq!(ffi_status_message(&ffi_status), b"test");
     }
 }

shell_wrapper/ahe_test.cc

@@ -54,19 +54,17 @@ constexpr double kSBase = 12.8;
 TEST(AheTest, EncryptDecryptOne) {
   // Create the public parameters.
   std::unique_ptr<std::string> public_seed;
-  auto status = GenerateSingleThreadHkdfSeed(public_seed);
-  SECAGG_ASSERT_OK(UnwrapFfiStatus(status));
+  SECAGG_ASSERT_OK(UnwrapFfiStatus(GenerateSingleThreadHkdfSeed(public_seed)));
   AhePublicParameters public_params;
   SECAGG_ASSERT_OK(UnwrapFfiStatus(CreateAhePublicParameters(
       kLogN, kT, kQs, std::size(kQs), kPublicKeyVariance, kSBase, kSFlood,
       ToRustSlice(*public_seed), &public_params)));
 
   std::unique_ptr<std::string> private_seed;
-  status = GenerateSingleThreadHkdfSeed(private_seed);
-  SECAGG_ASSERT_OK(UnwrapFfiStatus(status));
+  SECAGG_ASSERT_OK(UnwrapFfiStatus(GenerateSingleThreadHkdfSeed(private_seed)));
   SingleThreadHkdfWrapper prng;
-  status = CreateSingleThreadHkdf(ToRustSlice(*private_seed), prng);
-  SECAGG_ASSERT_OK(UnwrapFfiStatus(status));
+  SECAGG_ASSERT_OK(UnwrapFfiStatus(
+      CreateSingleThreadHkdf(ToRustSlice(*private_seed), prng)));
 
   RnsPolynomialWrapper sk_share;
   SECAGG_ASSERT_OK(
@@ -118,19 +116,17 @@ TEST(AheTest, EncryptDecryptOne) {
 TEST(AheTest, ExternCRecoveryTest) {
   // Create the public parameters.
   std::unique_ptr<std::string> public_seed;
-  auto status = GenerateSingleThreadHkdfSeed(public_seed);
-  SECAGG_ASSERT_OK(UnwrapFfiStatus(status));
+  SECAGG_ASSERT_OK(UnwrapFfiStatus(GenerateSingleThreadHkdfSeed(public_seed)));
   AhePublicParameters public_params;
   SECAGG_ASSERT_OK(UnwrapFfiStatus(CreateAhePublicParameters(
       kLogN, kT, kQs, std::size(kQs), kPublicKeyVariance, kSBase, kSFlood,
       ToRustSlice(*public_seed), &public_params)));
 
   std::unique_ptr<std::string> private_seed;
-  status = GenerateSingleThreadHkdfSeed(private_seed);
-  SECAGG_ASSERT_OK(UnwrapFfiStatus(status));
+  SECAGG_ASSERT_OK(UnwrapFfiStatus(GenerateSingleThreadHkdfSeed(private_seed)));
   SingleThreadHkdfWrapper prng;
-  status = CreateSingleThreadHkdf(ToRustSlice(*private_seed), prng);
-  SECAGG_ASSERT_OK(UnwrapFfiStatus(status));
+  SECAGG_ASSERT_OK(UnwrapFfiStatus(
+      CreateSingleThreadHkdf(ToRustSlice(*private_seed), prng)));
 
   std::vector<RnsPolynomialWrapper> secret_key_shares;
   std::vector<RnsPolynomialWrapper> public_key_shares;

shell_wrapper/shell_types.cc

@@ -264,7 +264,7 @@ FfiStatus WriteRnsPolynomialToBuffer128(
     buffer[2 * i + 1] = absl::Uint128High64(coeff_128);
   }
 
-  return {};
+  return MakeFfiStatus();
 }
 
 RnsPolynomialWrapper CloneRnsPolynomialWrapper(const RnsPolynomialWrapper* in) {