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

5 files changed, 39 insertions, 17 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 \
         */