Fix secure endpoint

Fixing a few bugs that slipped into the new secure endpoint implementation

PiperOrigin-RevId: 746145618

Author: ctiller

src/core/handshaker/security/secure_endpoint.cc

@@ -79,6 +79,10 @@ class FrameProtector : public RefCounted<FrameProtector> {
                           ->memory_quota()
                           ->CreateMemoryOwner()),
         self_reservation_(memory_owner_.MakeReservation(sizeof(*this))) {
+    GRPC_TRACE_LOG(secure_endpoint, INFO)
+        << "FrameProtector: " << this << " protector: " << protector_
+        << " zero_copy_protector: " << zero_copy_protector_
+        << " leftover_nslices: " << leftover_nslices;
     if (leftover_nslices > 0) {
       leftover_bytes_ = std::make_unique<SliceBuffer>();
       for (size_t i = 0; i < leftover_nslices; i++) {
@@ -106,6 +110,18 @@ class FrameProtector : public RefCounted<FrameProtector> {
   Mutex* read_mu() ABSL_LOCK_RETURNED(read_mu_) { return &read_mu_; }
   Mutex* write_mu() ABSL_LOCK_RETURNED(write_mu_) { return &write_mu_; }
 
+  void TraceOp(absl::string_view op, grpc_slice_buffer* slices) {
+    if (GRPC_TRACE_FLAG_ENABLED(secure_endpoint) && ABSL_VLOG_IS_ON(2)) {
+      size_t i;
+      for (i = 0; i < slices->count; i++) {
+        char* data =
+            grpc_dump_slice(slices->slices[i], GPR_DUMP_HEX | GPR_DUMP_ASCII);
+        VLOG(2) << op << " " << this << ": " << data;
+        gpr_free(data);
+      }
+    }
+  }
+
   void MaybePostReclaimer() {
     if (!has_posted_reclaimer_.exchange(true, std::memory_order_relaxed)) {
       memory_owner_.PostReclaimer(
@@ -146,15 +162,7 @@ class FrameProtector : public RefCounted<FrameProtector> {
   }
 
   void FinishRead(bool ok) {
-    if (GRPC_TRACE_FLAG_ENABLED(secure_endpoint) && ABSL_VLOG_IS_ON(2)) {
-      size_t i;
-      for (i = 0; i < read_buffer_->count; i++) {
-        char* data = grpc_dump_slice(read_buffer_->slices[i],
-                                     GPR_DUMP_HEX | GPR_DUMP_ASCII);
-        VLOG(2) << "READ " << this << ": " << data;
-        gpr_free(data);
-      }
-    }
+    TraceOp("FinishRead", read_buffer_);
     // TODO(yangg) experiment with moving this block after read_cb to see if it
     // helps latency
     source_buffer_.Clear();
@@ -243,6 +251,9 @@ class FrameProtector : public RefCounted<FrameProtector> {
           absl::StrCat("Unwrap failed (", tsi_result_to_string(result), ")"));
     }
 
+    GRPC_TRACE_LOG(secure_endpoint, INFO)
+        << "Unprotect: " << this << " read_status: " << read_status;
+
     return read_status;
   }
 
@@ -252,6 +263,9 @@ class FrameProtector : public RefCounted<FrameProtector> {
   }
 
   bool MaybeCompleteReadImmediately() {
+    GRPC_TRACE_LOG(secure_endpoint, INFO)
+        << "MaybeCompleteReadImmediately: " << this
+        << " leftover_bytes_: " << leftover_bytes_.get();
     if (leftover_bytes_ != nullptr) {
       grpc_slice_buffer_swap(leftover_bytes_->c_slice_buffer(),
                              source_buffer_.c_slice_buffer());
@@ -285,14 +299,7 @@ class FrameProtector : public RefCounted<FrameProtector> {
 
     output_buffer_.Clear();
 
-    if (GRPC_TRACE_FLAG_ENABLED(secure_endpoint) && ABSL_VLOG_IS_ON(2)) {
-      for (size_t i = 0; i < slices->count; i++) {
-        char* data =
-            grpc_dump_slice(slices->slices[i], GPR_DUMP_HEX | GPR_DUMP_ASCII);
-        VLOG(2) << "WRITE " << this << ": " << data;
-        gpr_free(data);
-      }
-    }
+    TraceOp("Protect", slices);
 
     tsi_result result = TSI_OK;
     if (zero_copy_protector_ != nullptr) {
@@ -678,24 +685,28 @@ class SecureEndpoint final : public EventEngine::Endpoint {
       }
       ReadArgs args;
       args.read_hint_bytes = frame_protector_.min_progress_size();
-      return wrapped_ep_->Read(
-                 [impl = Ref()](absl::Status status) mutable {
-                   {
-                     grpc_core::MutexLock lock(
-                         impl->frame_protector_.read_mu());
-                     if (status.ok() && impl->wrapped_ep_ == nullptr) {
-                       status =
-                           absl::CancelledError("secure endpoint shutdown");
-                     }
-                     status =
-                         impl->frame_protector_.Unprotect(std::move(status));
-                   }
-                   auto on_read = std::move(impl->on_read_);
-                   impl.reset();
-                   on_read(status);
-                 },
-                 frame_protector_.source_buffer(), &args) &&
-             MaybeFinishReadImmediately();
+      bool read_completed_immediately = wrapped_ep_->Read(
+          [impl = Ref()](absl::Status status) mutable {
+            {
+              grpc_core::MutexLock lock(impl->frame_protector_.read_mu());
+              if (status.ok() && impl->wrapped_ep_ == nullptr) {
+                status = absl::CancelledError("secure endpoint shutdown");
+              }
+              status = impl->frame_protector_.Unprotect(std::move(status));
+            }
+            if (status.ok()) {
+              impl->frame_protector_.TraceOp(
+                  "Read",
+                  impl->frame_protector_.source_buffer()->c_slice_buffer());
+            }
+            auto on_read = std::move(impl->on_read_);
+            impl->frame_protector_.FinishRead(status.ok());
+            impl.reset();
+            on_read(status);
+          },
+          frame_protector_.source_buffer(), &args);
+      if (read_completed_immediately) return MaybeFinishReadImmediately();
+      return false;
     }
 
     bool Write(absl::AnyInvocable<void(absl::Status)> on_writable,
@@ -717,6 +728,8 @@ class SecureEndpoint final : public EventEngine::Endpoint {
         return false;
       }
       on_write_ = std::move(on_writable);
+      frame_protector_.TraceOp(
+          "Write", frame_protector_.output_buffer()->c_slice_buffer());
       return wrapped_ep_->Write(
           [impl = Ref()](absl::Status status) mutable {
             auto on_write = std::move(impl->on_write_);
@@ -747,7 +760,10 @@ class SecureEndpoint final : public EventEngine::Endpoint {
    private:
     bool MaybeFinishReadImmediately() {
       grpc_core::MutexLock lock(frame_protector_.read_mu());
+      frame_protector_.TraceOp(
+          "Read(Imm)", frame_protector_.source_buffer()->c_slice_buffer());
       auto status = frame_protector_.Unprotect(absl::OkStatus());
+      frame_protector_.FinishRead(status.ok());
       if (status.ok()) return true;
       event_engine_->Run([impl = Ref(), status = std::move(status)]() mutable {
         auto on_read = std::move(impl->on_read_);