Allow extra copy in zero-copy protector integrity-only mode

7 files changed, 82 insertions, 41 deletions

diff --git a/src/core/tsi/alts/handshaker/alts_tsi_handshaker.cc b/src/core/tsi/alts/handshaker/alts_tsi_handshaker.cc
index 9676085380..06b999899a 100644
--- a/src/core/tsi/alts/handshaker/alts_tsi_handshaker.cc
+++ b/src/core/tsi/alts/handshaker/alts_tsi_handshaker.cc
@@ -127,7 +127,8 @@ static tsi_result handshaker_result_create_zero_copy_grpc_protector(
   tsi_result ok = alts_zero_copy_grpc_protector_create(
       reinterpret_cast<const uint8_t*>(result->key_data),
       kAltsAes128GcmRekeyKeyLength, /*is_rekey=*/true, result->is_client,
-      /*is_integrity_only=*/false, max_output_protected_frame_size, protector);
+      /*is_integrity_only=*/false, /*enable_extra_copy=*/false,
+      max_output_protected_frame_size, protector);
   if (ok != TSI_OK) {
     gpr_log(GPR_ERROR, "Failed to create zero-copy grpc protector");
   }
diff --git a/src/core/tsi/alts/zero_copy_frame_protector/alts_grpc_integrity_only_record_protocol.cc b/src/core/tsi/alts/zero_copy_frame_protector/alts_grpc_integrity_only_record_protocol.cc
index 7ba03eb7f0..72adce54bc 100644
--- a/src/core/tsi/alts/zero_copy_frame_protector/alts_grpc_integrity_only_record_protocol.cc
+++ b/src/core/tsi/alts/zero_copy_frame_protector/alts_grpc_integrity_only_record_protocol.cc
@@ -30,6 +30,7 @@
 /* Main struct for alts_grpc_integrity_only_record_protocol.  */
 typedef struct alts_grpc_integrity_only_record_protocol {
   alts_grpc_record_protocol base;
+  bool enable_extra_copy;
   grpc_slice_buffer data_sb;
   unsigned char* tag_buf;
 } alts_grpc_integrity_only_record_protocol;
@@ -46,6 +47,8 @@ static tsi_result alts_grpc_integrity_only_protect(
             "Invalid nullptr arguments to alts_grpc_record_protocol protect.");
     return TSI_INVALID_ARGUMENT;
   }
+  alts_grpc_integrity_only_record_protocol* integrity_only_record_protocol =
+      reinterpret_cast<alts_grpc_integrity_only_record_protocol*>(rp);
   /* Allocates memory for header and tag slices.  */
   grpc_slice header_slice = GRPC_SLICE_MALLOC(rp->header_length);
   grpc_slice tag_slice = GRPC_SLICE_MALLOC(rp->tag_length);
@@ -67,7 +70,16 @@ static tsi_result alts_grpc_integrity_only_protect(
   }
   /* Appends result to protected_slices.  */
   grpc_slice_buffer_add(protected_slices, header_slice);
-  grpc_slice_buffer_move_into(unprotected_slices, protected_slices);
+  if (integrity_only_record_protocol->enable_extra_copy) {
+    /* If extra copy mode is enabled, makes a copy of unprotected_slices.  */
+    for (size_t i = 0; i < unprotected_slices->count; i++) {
+      grpc_slice_buffer_add(protected_slices,
+                            grpc_slice_dup(unprotected_slices->slices[i]));
+    }
+    grpc_slice_buffer_reset_and_unref_internal(unprotected_slices);
+  } else {
+    grpc_slice_buffer_move_into(unprotected_slices, protected_slices);
+  }
   grpc_slice_buffer_add(protected_slices, tag_slice);
   return TSI_OK;
 }
@@ -152,7 +164,7 @@ static const alts_grpc_record_protocol_vtable
 
 tsi_result alts_grpc_integrity_only_record_protocol_create(
     gsec_aead_crypter* crypter, size_t overflow_size, bool is_client,
-    bool is_protect, alts_grpc_record_protocol** rp) {
+    bool is_protect, bool enable_extra_copy, alts_grpc_record_protocol** rp) {
   if (crypter == nullptr || rp == nullptr) {
     gpr_log(GPR_ERROR,
             "Invalid nullptr arguments to alts_grpc_record_protocol create.");
@@ -169,6 +181,7 @@ tsi_result alts_grpc_integrity_only_record_protocol_create(
     gpr_free(impl);
     return result;
   }
+  impl->enable_extra_copy = enable_extra_copy;
   /* Initializes slice buffer for data_sb.  */
   grpc_slice_buffer_init(&impl->data_sb);
   /* Allocates tag buffer.  */
diff --git a/src/core/tsi/alts/zero_copy_frame_protector/alts_grpc_integrity_only_record_protocol.h b/src/core/tsi/alts/zero_copy_frame_protector/alts_grpc_integrity_only_record_protocol.h
index 8d68b27e07..5456d34fad 100644
--- a/src/core/tsi/alts/zero_copy_frame_protector/alts_grpc_integrity_only_record_protocol.h
+++ b/src/core/tsi/alts/zero_copy_frame_protector/alts_grpc_integrity_only_record_protocol.h
@@ -38,6 +38,8 @@
  *   be used at the client or server side.
  * - is_protect: a flag indicating if the alts_grpc_record_protocol instance
  *   will be used for protect or unprotect.
+ *-  enable_extra_copy: a flag indicating if the instance uses one-copy instead
+ *   of zero-copy in the protect operation.
  * - rp: an alts_grpc_record_protocol instance to be returned from
  *   the method.
  *
@@ -46,7 +48,7 @@
  */
 tsi_result alts_grpc_integrity_only_record_protocol_create(
     gsec_aead_crypter* crypter, size_t overflow_size, bool is_client,
-    bool is_protect, alts_grpc_record_protocol** rp);
+    bool is_protect, bool enable_extra_copy, alts_grpc_record_protocol** rp);
 
 #endif /* GRPC_CORE_TSI_ALTS_ZERO_COPY_FRAME_PROTECTOR_ALTS_GRPC_INTEGRITY_ONLY_RECORD_PROTOCOL_H \
         */
diff --git a/src/core/tsi/alts/zero_copy_frame_protector/alts_zero_copy_grpc_protector.cc b/src/core/tsi/alts/zero_copy_frame_protector/alts_zero_copy_grpc_protector.cc
index 608213745e..58aba9b747 100644
--- a/src/core/tsi/alts/zero_copy_frame_protector/alts_zero_copy_grpc_protector.cc
+++ b/src/core/tsi/alts/zero_copy_frame_protector/alts_zero_copy_grpc_protector.cc
@@ -110,7 +110,7 @@ static bool read_frame_size(const grpc_slice_buffer* sb,
  */
 static tsi_result create_alts_grpc_record_protocol(
     const uint8_t* key, size_t key_size, bool is_rekey, bool is_client,
-    bool is_integrity_only, bool is_protect,
+    bool is_integrity_only, bool is_protect, bool enable_extra_copy,
     alts_grpc_record_protocol** record_protocol) {
   if (key == nullptr || record_protocol == nullptr) {
     return TSI_INVALID_ARGUMENT;
@@ -130,13 +130,13 @@ static tsi_result create_alts_grpc_record_protocol(
                                    : kAltsRecordProtocolFrameLimit;
   /* Creates alts_grpc_record_protocol with AEAD crypter ownership transferred.
    */
-  tsi_result result =
-      is_integrity_only
-          ? alts_grpc_integrity_only_record_protocol_create(
-                crypter, overflow_limit, is_client, is_protect, record_protocol)
-          : alts_grpc_privacy_integrity_record_protocol_create(
-                crypter, overflow_limit, is_client, is_protect,
-                record_protocol);
+  tsi_result result = is_integrity_only
+                          ? alts_grpc_integrity_only_record_protocol_create(
+                                crypter, overflow_limit, is_client, is_protect,
+                                enable_extra_copy, record_protocol)
+                          : alts_grpc_privacy_integrity_record_protocol_create(
+                                crypter, overflow_limit, is_client, is_protect,
+                                record_protocol);
   if (result != TSI_OK) {
     gsec_aead_crypter_destroy(crypter);
     return result;
@@ -241,7 +241,8 @@ static const tsi_zero_copy_grpc_protector_vtable
 
 tsi_result alts_zero_copy_grpc_protector_create(
     const uint8_t* key, size_t key_size, bool is_rekey, bool is_client,
-    bool is_integrity_only, size_t* max_protected_frame_size,
+    bool is_integrity_only, bool enable_extra_copy,
+    size_t* max_protected_frame_size,
     tsi_zero_copy_grpc_protector** protector) {
   if (grpc_core::ExecCtx::Get() == nullptr || key == nullptr ||
       protector == nullptr) {
@@ -257,11 +258,11 @@ tsi_result alts_zero_copy_grpc_protector_create(
   /* Creates alts_grpc_record_protocol objects.  */
   tsi_result status = create_alts_grpc_record_protocol(
       key, key_size, is_rekey, is_client, is_integrity_only,
-      /*is_protect=*/true, &impl->record_protocol);
+      /*is_protect=*/true, enable_extra_copy, &impl->record_protocol);
   if (status == TSI_OK) {
     status = create_alts_grpc_record_protocol(
         key, key_size, is_rekey, is_client, is_integrity_only,
-        /*is_protect=*/false, &impl->unrecord_protocol);
+        /*is_protect=*/false, enable_extra_copy, &impl->unrecord_protocol);
     if (status == TSI_OK) {
       /* Sets maximum frame size.  */
       size_t max_protected_frame_size_to_set = kDefaultFrameLength;
diff --git a/src/core/tsi/alts/zero_copy_frame_protector/alts_zero_copy_grpc_protector.h b/src/core/tsi/alts/zero_copy_frame_protector/alts_zero_copy_grpc_protector.h
index 71e953cfc1..515c27ea05 100644
--- a/src/core/tsi/alts/zero_copy_frame_protector/alts_zero_copy_grpc_protector.h
+++ b/src/core/tsi/alts/zero_copy_frame_protector/alts_zero_copy_grpc_protector.h
@@ -35,6 +35,11 @@
  *   server side.
  * - is_integrity_only: a flag indicating if the protector instance will be
  *   used for integrity-only or privacy-integrity mode.
+ * - enable_extra_copy: a flag indicating if the protector instance does one
+ *   extra memory copy during the protect operation for integrity_only mode.
+ *   For the unprotect operation, it is still zero-copy. If application intends
+ *   to modify the data buffer after the protect operation, we can turn on this
+ *   mode to avoid integrity check failure.
  * - max_protected_frame_size: an in/out parameter indicating max frame size
  *   to be used by the protector. If it is nullptr, the default frame size will
  *   be used. Otherwise, the provided frame size will be adjusted (if not
@@ -45,8 +50,8 @@
  */
 tsi_result alts_zero_copy_grpc_protector_create(
     const uint8_t* key, size_t key_size, bool is_rekey, bool is_client,
-    bool is_integrity_only, size_t* max_protected_frame_size,
-    tsi_zero_copy_grpc_protector** protector);
+    bool is_integrity_only, bool enable_extra_copy,
+    size_t* max_protected_frame_size, tsi_zero_copy_grpc_protector** protector);
 
 #endif /* GRPC_CORE_TSI_ALTS_ZERO_COPY_FRAME_PROTECTOR_ALTS_ZERO_COPY_GRPC_PROTECTOR_H \
         */
diff --git a/test/core/tsi/alts/zero_copy_frame_protector/alts_grpc_record_protocol_test.cc b/test/core/tsi/alts/zero_copy_frame_protector/alts_grpc_record_protocol_test.cc
index b763f19d50..3ae64d6f20 100644
--- a/test/core/tsi/alts/zero_copy_frame_protector/alts_grpc_record_protocol_test.cc
+++ b/test/core/tsi/alts/zero_copy_frame_protector/alts_grpc_record_protocol_test.cc
@@ -109,7 +109,7 @@ static void alter_random_byte(grpc_slice_buffer* sb) {
 }
 
 static alts_grpc_record_protocol_test_fixture*
-test_fixture_integrity_only_create(bool rekey) {
+test_fixture_integrity_only_create(bool rekey, bool extra_copy) {
   alts_grpc_record_protocol_test_fixture* fixture =
       static_cast<alts_grpc_record_protocol_test_fixture*>(
           gpr_zalloc(sizeof(alts_grpc_record_protocol_test_fixture)));
@@ -124,41 +124,46 @@ test_fixture_integrity_only_create(bool rekey) {
                  &crypter, nullptr) == GRPC_STATUS_OK);
   GPR_ASSERT(alts_grpc_integrity_only_record_protocol_create(
                  crypter, 8, /*is_client=*/true, /*is_protect=*/true,
-                 &fixture->client_protect) == TSI_OK);
+                 extra_copy, &fixture->client_protect) == TSI_OK);
   /* Create client record protocol for unprotect.  */
   GPR_ASSERT(gsec_aes_gcm_aead_crypter_create(
                  key, key_length, kAesGcmNonceLength, kAesGcmTagLength, rekey,
                  &crypter, nullptr) == GRPC_STATUS_OK);
   GPR_ASSERT(alts_grpc_integrity_only_record_protocol_create(
                  crypter, 8, /*is_client=*/true, /*is_protect=*/false,
-                 &fixture->client_unprotect) == TSI_OK);
+                 extra_copy, &fixture->client_unprotect) == TSI_OK);
   /* Create server record protocol for protect.  */
   GPR_ASSERT(gsec_aes_gcm_aead_crypter_create(
                  key, key_length, kAesGcmNonceLength, kAesGcmTagLength, rekey,
                  &crypter, nullptr) == GRPC_STATUS_OK);
   GPR_ASSERT(alts_grpc_integrity_only_record_protocol_create(
                  crypter, 8, /*is_client=*/false, /*is_protect=*/true,
-                 &fixture->server_protect) == TSI_OK);
+                 extra_copy, &fixture->server_protect) == TSI_OK);
   /* Create server record protocol for unprotect.  */
   GPR_ASSERT(gsec_aes_gcm_aead_crypter_create(
                  key, key_length, kAesGcmNonceLength, kAesGcmTagLength, rekey,
                  &crypter, nullptr) == GRPC_STATUS_OK);
   GPR_ASSERT(alts_grpc_integrity_only_record_protocol_create(
                  crypter, 8, /*is_client=*/false, /*is_protect=*/false,
-                 &fixture->server_unprotect) == TSI_OK);
+                 extra_copy, &fixture->server_unprotect) == TSI_OK);
 
   gpr_free(key);
   return fixture;
 }
 
 static alts_grpc_record_protocol_test_fixture*
-test_fixture_integrity_only_no_rekey_create() {
-  return test_fixture_integrity_only_create(false);
+test_fixture_integrity_only_no_rekey_no_extra_copy_create() {
+  return test_fixture_integrity_only_create(false, false);
 }
 
 static alts_grpc_record_protocol_test_fixture*
 test_fixture_integrity_only_rekey_create() {
-  return test_fixture_integrity_only_create(true);
+  return test_fixture_integrity_only_create(true, false);
+}
+
+static alts_grpc_record_protocol_test_fixture*
+test_fixture_integrity_only_extra_copy_create() {
+  return test_fixture_integrity_only_create(false, true);
 }
 
 static alts_grpc_record_protocol_test_fixture*
@@ -440,9 +445,12 @@ static void alts_grpc_record_protocol_tests(
 }
 
 int main(int argc, char** argv) {
-  alts_grpc_record_protocol_tests(&test_fixture_integrity_only_no_rekey_create);
+  alts_grpc_record_protocol_tests(
+      &test_fixture_integrity_only_no_rekey_no_extra_copy_create);
   alts_grpc_record_protocol_tests(&test_fixture_integrity_only_rekey_create);
   alts_grpc_record_protocol_tests(
+      &test_fixture_integrity_only_extra_copy_create);
+  alts_grpc_record_protocol_tests(
       &test_fixture_privacy_integrity_no_rekey_create);
   alts_grpc_record_protocol_tests(&test_fixture_privacy_integrity_rekey_create);
 
diff --git a/test/core/tsi/alts/zero_copy_frame_protector/alts_zero_copy_grpc_protector_test.cc b/test/core/tsi/alts/zero_copy_frame_protector/alts_zero_copy_grpc_protector_test.cc
index 32159e22f2..3ee8323a31 100644
--- a/test/core/tsi/alts/zero_copy_frame_protector/alts_zero_copy_grpc_protector_test.cc
+++ b/test/core/tsi/alts/zero_copy_frame_protector/alts_zero_copy_grpc_protector_test.cc
@@ -100,7 +100,8 @@ static bool are_slice_buffers_equal(grpc_slice_buffer* first,
 
 static alts_zero_copy_grpc_protector_test_fixture*
 alts_zero_copy_grpc_protector_test_fixture_create(bool rekey,
-                                                  bool integrity_only) {
+                                                  bool integrity_only,
+                                                  bool enable_extra_copy) {
   alts_zero_copy_grpc_protector_test_fixture* fixture =
       static_cast<alts_zero_copy_grpc_protector_test_fixture*>(
           gpr_zalloc(sizeof(alts_zero_copy_grpc_protector_test_fixture)));
@@ -111,10 +112,12 @@ alts_zero_copy_grpc_protector_test_fixture_create(bool rekey,
   gsec_test_random_array(&key, key_length);
   GPR_ASSERT(alts_zero_copy_grpc_protector_create(
                  key, key_length, rekey, /*is_client=*/true, integrity_only,
-                 &max_protected_frame_size, &fixture->client) == TSI_OK);
+                 enable_extra_copy, &max_protected_frame_size,
+                 &fixture->client) == TSI_OK);
   GPR_ASSERT(alts_zero_copy_grpc_protector_create(
                  key, key_length, rekey, /*is_client=*/false, integrity_only,
-                 &max_protected_frame_size, &fixture->server) == TSI_OK);
+                 enable_extra_copy, &max_protected_frame_size,
+                 &fixture->server) == TSI_OK);
   gpr_free(key);
   grpc_core::ExecCtx::Get()->Flush();
   return fixture;
@@ -229,62 +232,70 @@ static void seal_unseal_large_buffer(tsi_zero_copy_grpc_protector* sender,
 
 /* --- Test cases. --- */
 
-static void alts_zero_copy_protector_seal_unseal_small_buffer_tests() {
+static void alts_zero_copy_protector_seal_unseal_small_buffer_tests(
+    bool enable_extra_copy) {
   alts_zero_copy_grpc_protector_test_fixture* fixture =
       alts_zero_copy_grpc_protector_test_fixture_create(
-          /*rekey=*/false, /*integrity_only=*/true);
+          /*rekey=*/false, /*integrity_only=*/true, enable_extra_copy);
   seal_unseal_small_buffer(fixture->client, fixture->server);
   seal_unseal_small_buffer(fixture->server, fixture->client);
   alts_zero_copy_grpc_protector_test_fixture_destroy(fixture);
 
   fixture = alts_zero_copy_grpc_protector_test_fixture_create(
-      /*rekey=*/false, /*integrity_only=*/false);
+      /*rekey=*/false, /*integrity_only=*/false, enable_extra_copy);
   seal_unseal_small_buffer(fixture->client, fixture->server);
   seal_unseal_small_buffer(fixture->server, fixture->client);
   alts_zero_copy_grpc_protector_test_fixture_destroy(fixture);
 
   fixture = alts_zero_copy_grpc_protector_test_fixture_create(
-      /*rekey=*/true, /*integrity_only=*/true);
+      /*rekey=*/true, /*integrity_only=*/true, enable_extra_copy);
   seal_unseal_small_buffer(fixture->client, fixture->server);
   seal_unseal_small_buffer(fixture->server, fixture->client);
   alts_zero_copy_grpc_protector_test_fixture_destroy(fixture);
 
   fixture = alts_zero_copy_grpc_protector_test_fixture_create(
-      /*rekey=*/true, /*integrity_only=*/false);
+      /*rekey=*/true, /*integrity_only=*/false, enable_extra_copy);
   seal_unseal_small_buffer(fixture->client, fixture->server);
   seal_unseal_small_buffer(fixture->server, fixture->client);
   alts_zero_copy_grpc_protector_test_fixture_destroy(fixture);
 }
 
-static void alts_zero_copy_protector_seal_unseal_large_buffer_tests() {
+static void alts_zero_copy_protector_seal_unseal_large_buffer_tests(
+    bool enable_extra_copy) {
   alts_zero_copy_grpc_protector_test_fixture* fixture =
       alts_zero_copy_grpc_protector_test_fixture_create(
-          /*rekey=*/false, /*integrity_only=*/true);
+          /*rekey=*/false, /*integrity_only=*/true, enable_extra_copy);
   seal_unseal_large_buffer(fixture->client, fixture->server);
   seal_unseal_large_buffer(fixture->server, fixture->client);
   alts_zero_copy_grpc_protector_test_fixture_destroy(fixture);
 
   fixture = alts_zero_copy_grpc_protector_test_fixture_create(
-      /*rekey=*/false, /*integrity_only=*/false);
+      /*rekey=*/false, /*integrity_only=*/false, enable_extra_copy);
   seal_unseal_large_buffer(fixture->client, fixture->server);
   seal_unseal_large_buffer(fixture->server, fixture->client);
   alts_zero_copy_grpc_protector_test_fixture_destroy(fixture);
 
   fixture = alts_zero_copy_grpc_protector_test_fixture_create(
-      /*rekey=*/true, /*integrity_only=*/true);
+      /*rekey=*/true, /*integrity_only=*/true, enable_extra_copy);
   seal_unseal_large_buffer(fixture->client, fixture->server);
   seal_unseal_large_buffer(fixture->server, fixture->client);
   alts_zero_copy_grpc_protector_test_fixture_destroy(fixture);
 
   fixture = alts_zero_copy_grpc_protector_test_fixture_create(
-      /*rekey=*/true, /*integrity_only=*/false);
+      /*rekey=*/true, /*integrity_only=*/false, enable_extra_copy);
   seal_unseal_large_buffer(fixture->client, fixture->server);
   seal_unseal_large_buffer(fixture->server, fixture->client);
   alts_zero_copy_grpc_protector_test_fixture_destroy(fixture);
 }
 
 int main(int argc, char** argv) {
-  alts_zero_copy_protector_seal_unseal_small_buffer_tests();
-  alts_zero_copy_protector_seal_unseal_large_buffer_tests();
+  alts_zero_copy_protector_seal_unseal_small_buffer_tests(
+      /*enable_extra_copy=*/false);
+  alts_zero_copy_protector_seal_unseal_small_buffer_tests(
+      /*enable_extra_copy=*/true);
+  alts_zero_copy_protector_seal_unseal_large_buffer_tests(
+      /*enable_extra_copy=*/false);
+  alts_zero_copy_protector_seal_unseal_large_buffer_tests(
+      /*enable_extra_copy=*/true);
   return 0;
 }