Remove zero padding when unpacking messages in KAHE decryption.

This requires storing the unpacked vector length in KAHE config.

PiperOrigin-RevId: 850474871

Author: tholop

shell_wrapper/BUILD

@@ -242,6 +242,7 @@ cc_library(
         ":status_macros",
         "@abseil-cpp//absl/status",
         "@abseil-cpp//absl/status:statusor",
+        "@abseil-cpp//absl/strings:str_format",
         "@abseil-cpp//absl/strings:string_view",
         "@abseil-cpp//absl/types:span",
         "@shell-encryption//shell_encryption/rns:coefficient_encoder",

shell_wrapper/kahe.cc

@@ -25,6 +25,7 @@
 
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "include/cxx.h"
@@ -347,6 +348,7 @@ FfiStatus PackMessagesRaw(rust::Slice<const uint64_t> messages,
 
 FfiStatus UnpackMessagesRaw(uint64_t packing_base, uint64_t packing_dimension,
                             uint64_t num_packed_values,
+                            uint64_t num_unpacked_values,
                             BigIntVectorWrapper& packed_values,
                             rust::Vec<uint64_t>& out) {
   // Validate the wrappers.
@@ -358,15 +360,25 @@ FfiStatus UnpackMessagesRaw(uint64_t packing_base, uint64_t packing_dimension,
     return MakeFfiStatus(
         absl::InvalidArgumentError("insufficient number of packed values."));
   }
+
+  // `unpacked_messages` is padded with zeros if needed.
   std::vector<uint64_t> unpacked_messages =
       rlwe::UnpackMessagesFlat<secure_aggregation::Integer,
                                secure_aggregation::BigInteger>(
           absl::MakeSpan(*packed_values.ptr).subspan(0, num_packed_values),
           packing_base, packing_dimension);
   packed_values.ptr->erase(packed_values.ptr->begin(),
                            packed_values.ptr->begin() + num_packed_values);
-  for (auto& val : unpacked_messages) {
-    out.push_back(val);
+
+  // Remove padding and copy values to Rust output vector.
+  if (unpacked_messages.size() < num_unpacked_values) {
+    return MakeFfiStatus(absl::InvalidArgumentError(
+        absl::StrFormat("unpacked messages is too short (%d) for the requested "
+                        "number of unpacked values (%d)",
+                        unpacked_messages.size(), num_unpacked_values)));
+  }
+  for (size_t i = 0; i < num_unpacked_values; ++i) {
+    out.push_back(unpacked_messages[i]);
   }
   return MakeFfiStatus();
 }

shell_wrapper/kahe.h

@@ -154,14 +154,16 @@ FfiStatus PackMessagesRaw(rust::Slice<const uint64_t> messages,
                           uint64_t num_packed_values,
                           BigIntVectorWrapper* packed_values);
 
-// Unpacks messages stored at `packed_values[0..num_packed_values]` and appends
-// them to `out`, and removes these packed values from `packed_values`.
+// Unpacks messages stored at `packed_values[0..num_packed_values]`, removes
+// these packed values from `packed_values` and appends the first
+// `num_unpacked_values` messages to `out`.
 // Expects `packed_values.ptr` to be a valid pointer to the vector of packed
 // values, and expects packing_base > 1, packing_dimension > 0,
 // num_packed_values > 0, packing_base^packing_dimension <
 // std::numeric_limits<BigInteger>::max().
 FfiStatus UnpackMessagesRaw(uint64_t packing_base, uint64_t packing_dimension,
                             uint64_t num_packed_values,
+                            uint64_t num_unpacked_values,
                             BigIntVectorWrapper& packed_values,
                             rust::Vec<uint64_t>& out);
 

shell_wrapper/kahe.rs

@@ -22,11 +22,22 @@ use std::collections::HashMap;
 use std::marker::PhantomData;
 use std::mem::MaybeUninit;
 
+/// Configuration for packing and unpacking. Used to convert a long vector of `length` small
+/// integers in [0, `base`) into a short vector of `num_packed_coeffs` large integers in
+/// [0, `base`^`dimension`), and vice versa.
 #[derive(Debug, PartialEq, Clone)]
 pub struct PackedVectorConfig {
+    /// Base for packing.
     pub base: u64,
+
+    /// Number of elements packed into each coefficient.
     pub dimension: u64,
+
+    /// Number of coefficients in the packed vector.
     pub num_packed_coeffs: u64,
+
+    /// Number of elements in the plaintext vector before packing.
+    pub length: u64,
 }
 
 #[cxx::bridge]
@@ -93,6 +104,7 @@ mod ffi {
             packing_base: u64,
             packing_dimension: u64,
             num_packed_values: u64,
+            num_unpacked_values: u64,
             packed_values: &mut BigIntVectorWrapper,
             out: &mut Vec<u64>,
         ) -> FfiStatus;
@@ -260,6 +272,7 @@ pub fn decrypt(
                 packed_vector_config.base,
                 packed_vector_config.dimension,
                 packed_vector_config.num_packed_coeffs,
+                packed_vector_config.length,
                 &mut packed_values,
                 &mut unpacked_values,
             )

shell_wrapper/kahe_test.cc

@@ -337,7 +337,8 @@ TEST(KaheTest, UnpackMessagesRawRemovesConsumedPackedValues) {
   rust::Vec<Integer> unpacked_messages;
   SECAGG_EXPECT_OK(UnwrapFfiStatus(
       UnpackMessagesRaw(packing_base, packing_dimension, num_packed_values,
-                        packed_values, unpacked_messages)));
+                        num_packed_values * packing_dimension, packed_values,
+                        unpacked_messages)));
   EXPECT_EQ(packed_values.ptr->size(), num_packed_values);
   EXPECT_EQ(unpacked_messages.size(), num_packed_values);
   // Unpacked values should match the first half of the original packed values.
@@ -399,8 +400,9 @@ TEST(KaheTest, PackAndEncrypt) {
       .ptr = std::make_unique<std::vector<BigInteger>>(std::move(decrypted))};
   rust::Vec<Integer> unpacked_messages;
   SECAGG_ASSERT_OK(UnwrapFfiStatus(
-      UnpackMessagesRaw(packing_base, num_packing, packed_messages.size(),
-                        decrypted_wrapper, unpacked_messages)));
+      UnpackMessagesRaw(packing_base, num_packing, num_packed_messages,
+                        num_packed_messages * num_packing, decrypted_wrapper,
+                        unpacked_messages)));
   EXPECT_EQ(absl::MakeSpan(unpacked_messages.data(), num_messages),
             absl::MakeSpan(expected_unpacked_messages.data(), num_messages));
   // Check against the original input messages.
@@ -461,7 +463,8 @@ TEST(KaheTest, RawVectorEncryptOnePolynomial) {
   rust::Vec<Integer> unpacked_decrypted_messages;
   SECAGG_ASSERT_OK(UnwrapFfiStatus(UnpackMessagesRaw(
       packing_base, num_packing, decrypted_wrapper.ptr->size(),
-      decrypted_wrapper, unpacked_decrypted_messages)));
+      decrypted_wrapper.ptr->size() * num_packing, decrypted_wrapper,
+      unpacked_decrypted_messages)));
 
   // Filled the whole buffer with right messages.
   EXPECT_EQ(absl::MakeSpan(unpacked_decrypted_messages.data(), num_messages),
@@ -481,6 +484,7 @@ TEST(KaheTest, RawVectorEncryptOnePolynomial) {
   unpacked_decrypted_long_messages.reserve(buffer_length);
   SECAGG_ASSERT_OK(UnwrapFfiStatus(UnpackMessagesRaw(
       packing_base, num_packing, decrypted_long_messages_wrapper.ptr->size(),
+      decrypted_long_messages_wrapper.ptr->size() * num_packing,
       decrypted_long_messages_wrapper, unpacked_decrypted_long_messages)));
 
   // The non-zero messages are identical.
@@ -493,6 +497,38 @@ TEST(KaheTest, RawVectorEncryptOnePolynomial) {
                              unpacked_decrypted_long_messages.size())
                   .subspan(num_messages, padded_length - num_messages),
               ::testing::Each(::testing::Eq(0)));
+
+  // Prepare a fresh packed plaintext.
+  BigIntVectorWrapper packed_messages_2{
+      .ptr = std::make_unique<std::vector<BigInteger>>()};
+  SECAGG_ASSERT_OK(
+      UnwrapFfiStatus(Decrypt(ciphertexts, key, params, &packed_messages_2)));
+  SECAGG_ASSERT_OK(UnwrapFfiStatus(
+      Encrypt(packed_messages_2, key, params, &prng, &ciphertexts)));
+
+  // Now unpack and directly pass the right length to remove padding.
+  rust::Vec<Integer> unpacked_messages_2;
+  SECAGG_ASSERT_OK(UnwrapFfiStatus(UnpackMessagesRaw(
+      packing_base, num_packing, packed_messages_2.ptr->size(), num_messages,
+      packed_messages_2, unpacked_messages_2)));
+  EXPECT_EQ(
+      absl::MakeSpan(unpacked_messages_2.data(), unpacked_messages_2.size()),
+      absl::MakeSpan(input_messages));
+
+  // Finally, check that we fail if we request too many unpacked messages
+  BigIntVectorWrapper packed_messages_3{
+      .ptr = std::make_unique<std::vector<BigInteger>>()};
+  SECAGG_ASSERT_OK(
+      UnwrapFfiStatus(Decrypt(ciphertexts, key, params, &packed_messages_2)));
+  SECAGG_ASSERT_OK(UnwrapFfiStatus(
+      Encrypt(packed_messages_3, key, params, &prng, &ciphertexts)));
+  rust::Vec<Integer> unpacked_messages_3;
+  int num_unpacked_messages_3 = packed_messages_3.ptr->size() * num_packing + 1;
+  EXPECT_THAT(
+      UnwrapFfiStatus(UnpackMessagesRaw(
+          packing_base, num_packing, packed_messages_3.ptr->size(),
+          num_unpacked_messages_3, packed_messages_3, unpacked_messages_3)),
+      StatusIs(absl::StatusCode::kInvalidArgument));
 }
 
 TEST(KaheTest, RawVectorEncryptTwoPolynomials) {
@@ -538,7 +574,7 @@ TEST(KaheTest, RawVectorEncryptTwoPolynomials) {
       UnwrapFfiStatus(Decrypt(ciphertexts, key, params, &decrypted_wrapper)));
   rust::Vec<Integer> unpacked_decrypted_messages;
   SECAGG_ASSERT_OK(UnwrapFfiStatus(UnpackMessagesRaw(
-      packing_base, num_packing, decrypted_wrapper.ptr->size(),
+      packing_base, num_packing, decrypted_wrapper.ptr->size(), num_messages,
       decrypted_wrapper, unpacked_decrypted_messages)));
 
   EXPECT_GE(unpacked_decrypted_messages.size(), num_messages);
@@ -643,7 +679,8 @@ TEST(KaheTest, UnpackMessagesRawFailsIfUnallocatedPackedValues) {
   rust::Vec<Integer> unpacked_messages;
   EXPECT_THAT(UnwrapFfiStatus(UnpackMessagesRaw(
                   packing_base, packing_dimension, num_packed_messages,
-                  bad_packed_values, unpacked_messages)),
+                  num_packed_messages * packing_dimension, bad_packed_values,
+                  unpacked_messages)),
               StatusIs(absl::StatusCode::kInvalidArgument));
 }
 
@@ -658,7 +695,8 @@ TEST(KaheTest, UnpackMessagesRawFailsIfPackedValuesTooShort) {
   rust::Vec<Integer> unpacked_messages;
   EXPECT_THAT(UnwrapFfiStatus(UnpackMessagesRaw(
                   packing_base, packing_dimension, num_packed_messages,
-                  bad_packed_values, unpacked_messages)),
+                  num_packed_messages * packing_dimension, bad_packed_values,
+                  unpacked_messages)),
               StatusIs(absl::StatusCode::kInvalidArgument));
 }
 
@@ -738,7 +776,7 @@ TEST(KaheTest, AddInPlacePolynomial) {
   rust::Vec<Integer> unpacked_decrypted_messages;
   unpacked_decrypted_messages.reserve(num_messages);
   SECAGG_ASSERT_OK(UnwrapFfiStatus(UnpackMessagesRaw(
-      packing_base, num_packing, decrypted_wrapper.ptr->size(),
+      packing_base, num_packing, decrypted_wrapper.ptr->size(), num_messages,
       decrypted_wrapper, unpacked_decrypted_messages)));
   for (int i = 0; i < num_messages; ++i) {
     EXPECT_EQ(input_values1[i] + input_values2[i],

shell_wrapper/kahe_test.rs

@@ -47,13 +47,13 @@ fn encrypt_decrypt() -> Result<()> {
     let plaintext = HashMap::from([(DEFAULT_ID, input_values.as_slice())]);
     let packed_vector_configs = HashMap::from([(
         DEFAULT_ID.to_string(),
-        PackedVectorConfig { base: 10, dimension: 2, num_packed_coeffs: 2 },
+        PackedVectorConfig { base: 10, dimension: 2, num_packed_coeffs: 2, length: 3 },
     )]);
     let ciphertext = encrypt(&plaintext, &packed_vector_configs, &secret_key, &params, &mut prng)?;
 
     let output_values = decrypt(&ciphertext, &secret_key, &params, &packed_vector_configs)?;
     expect_that!(output_values.contains_key(DEFAULT_ID), eq(true));
-    expect_that!(output_values[DEFAULT_ID][..3], container_eq(input_values));
+    expect_that!(output_values[DEFAULT_ID], container_eq(input_values));
     Ok(())
 }
 
@@ -80,14 +80,16 @@ fn encrypt_decrypt_padding() -> Result<()> {
     let input_values: Vec<u64> =
         (0..num_input_values).map(|_| rand::thread_rng().gen_range(0..input_domain)).collect();
 
-    // Encrypt the vector.
+    // Encrypt the vector. Pass a longer length than what we need.
+    let padded_length = (num_packed_coeffs * packing_dimension) as usize;
     let plaintext = HashMap::from([(DEFAULT_ID, input_values.as_slice())]);
     let packed_vector_configs = HashMap::from([(
         DEFAULT_ID.to_string(),
         PackedVectorConfig {
             base: input_domain as u64,
             dimension: packing_dimension as u64,
             num_packed_coeffs: num_packed_coeffs as u64,
+            length: padded_length as u64,
         },
     )]);
     let ciphertext = encrypt(&plaintext, &packed_vector_configs, &secret_key, &params, &mut prng)?;
@@ -97,7 +99,6 @@ fn encrypt_decrypt_padding() -> Result<()> {
     let output_values = &decrypted[DEFAULT_ID];
 
     // Check that message is correctly decrypted with right padding.
-    let padded_length = (num_packed_coeffs * packing_dimension) as usize;
     expect_that!(output_values.len(), eq(padded_length));
     expect_that!(output_values.len(), gt(num_input_values));
     expect_that!(output_values[..num_input_values], container_eq(input_values));
@@ -138,22 +139,17 @@ fn encrypt_decrypt_long() -> Result<()> {
             base: input_domain as u64,
             dimension: packing_dimension as u64,
             num_packed_coeffs: num_packed_coeffs as u64,
+            length: num_input_values as u64,
         },
     )]);
     let ciphertext = encrypt(&plaintext, &packed_vector_configs, &secret_key, &params, &mut prng)?;
 
     let decrypted = decrypt(&ciphertext, &secret_key, &params, &packed_vector_configs)?;
     let output_values = &decrypted[DEFAULT_ID];
 
-    // Check that message is correctly decrypted with right padding.
-    let padded_length = num_packed_coeffs * packing_dimension;
-    expect_that!(output_values.len(), eq(padded_length));
-    expect_that!(output_values.len(), gt(num_input_values));
-    expect_that!(output_values[..num_input_values], container_eq(input_values));
-    expect_that!(
-        output_values[num_input_values..],
-        container_eq(vec![0; padded_length - num_input_values])
-    );
+    // Check that message is correctly decrypted (no padding).
+    expect_that!(output_values.len(), eq(num_input_values));
+    expect_that!(output_values, container_eq(input_values));
 
     // If the input is too long, we should fail.
     let num_values_too_long = num_public_polynomials * poly_capacity + 1;
@@ -195,6 +191,7 @@ fn encrypt_decrypt_two_vectors() -> Result<()> {
                 base: input_domains[0] as u64,
                 dimension: packing_dimensions[0] as u64,
                 num_packed_coeffs: num_packed_coeffs[0] as u64,
+                length: num_input_values[0] as u64,
             },
         ),
         (
@@ -203,6 +200,7 @@ fn encrypt_decrypt_two_vectors() -> Result<()> {
                 base: input_domains[1] as u64,
                 dimension: packing_dimensions[1] as u64,
                 num_packed_coeffs: num_packed_coeffs[1] as u64,
+                length: num_input_values[1] as u64,
             },
         ),
     ]);
@@ -218,26 +216,16 @@ fn encrypt_decrypt_two_vectors() -> Result<()> {
         HashMap::from([(ID0, input_values0.as_slice()), (ID1, input_values1.as_slice())]);
     let ciphertext = encrypt(&plaintext, &packed_vector_configs, &secret_key, &params, &mut prng)?;
 
-    // Decrypt and check the output contains the two vectors that are padded correctly.
+    // Decrypt and check the output contains the two vectors.
     let decrypted = decrypt(&ciphertext, &secret_key, &params, &packed_vector_configs)?;
     verify_that!(decrypted.contains_key(ID0), eq(true))?;
     verify_that!(decrypted.contains_key(ID1), eq(true))?;
 
     let output_values0 = &decrypted[ID0];
     let output_values1 = &decrypted[ID1];
-    expect_that!(output_values0.len(), eq(num_packed_coeffs[0] * packing_dimensions[0]));
-    expect_that!(output_values0.len(), gt(num_input_values[0]));
-    expect_that!(output_values0[..num_input_values[0]], container_eq(input_values0));
-    expect_that!(
-        output_values0[num_input_values[0]..],
-        container_eq(vec![0; num_packed_coeffs[0] * packing_dimensions[0] - num_input_values[0]])
-    );
-    expect_that!(output_values1.len(), eq(num_packed_coeffs[1] * packing_dimensions[1]));
-    expect_that!(output_values1.len(), gt(num_input_values[1]));
-    expect_that!(output_values1[..num_input_values[1]], container_eq(input_values1));
-    expect_that!(
-        output_values1[num_input_values[1]..],
-        container_eq(vec![0; num_packed_coeffs[1] * packing_dimensions[1] - num_input_values[1]])
-    );
+    expect_that!(output_values0.len(), eq(num_input_values[0]));
+    expect_that!(output_values0, container_eq(input_values0));
+    expect_that!(output_values1.len(), eq(num_input_values[1]));
+    expect_that!(output_values1, container_eq(input_values1));
     Ok(())
 }

willow/proto/shell/parameters.proto

@@ -28,6 +28,7 @@ message PackedVectorConfigProto {
   int64 base = 1;
   int64 dimension = 2;
   int64 num_packed_coeffs = 3;
+  int64 length = 4;
 }
 
 // This proto defines the parameters for instantiating the KAHE scheme