From 04fa4328316a70fe776930582f810c5447a95ff2 Mon Sep 17 00:00:00 2001
From: Yash Tibrewal <yashkt@google.com>
Date: Tue, 6 Aug 2019 16:12:51 -0700
Subject: [PATCH] Fail SEND_MESSAGE ops if stream is closed for writes

---
 .../chttp2/transport/chttp2_transport.cc         | 16 +++++-----------
 src/core/lib/surface/call.cc                     |  6 ++++++
 src/core/lib/transport/transport.h               | 12 ++++++++++++
 3 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/src/core/ext/transport/chttp2/transport/chttp2_transport.cc b/src/core/ext/transport/chttp2/transport/chttp2_transport.cc
index 633b35e7d61..1a1c03208eb 100644
--- a/src/core/ext/transport/chttp2/transport/chttp2_transport.cc
+++ b/src/core/ext/transport/chttp2/transport/chttp2_transport.cc
@@ -1532,19 +1532,13 @@ static void perform_stream_op_locked(void* stream_op,
     on_complete->next_data.scratch |= CLOSURE_BARRIER_MAY_COVER_WRITE;
     s->fetching_send_message_finished = add_closure_barrier(op->on_complete);
     if (s->write_closed) {
-      // Return an error unless the client has already received trailing
-      // metadata from the server, since an application using a
-      // streaming call might send another message before getting a
-      // recv_message failure, breaking out of its loop, and then
-      // starting recv_trailing_metadata.
+      op->payload->send_message.stream_write_closed = true;
+      // We should NOT return an error here, so as to avoid a cancel OP being
+      // started. The surface layer will notice that the stream has been closed
+      // for writes and fail the send message op.
       op->payload->send_message.send_message.reset();
       grpc_chttp2_complete_closure_step(
-          t, s, &s->fetching_send_message_finished,
-          t->is_client && s->received_trailing_metadata
-              ? GRPC_ERROR_NONE
-              : GRPC_ERROR_CREATE_REFERENCING_FROM_STATIC_STRING(
-                    "Attempt to send message after stream was closed",
-                    &s->write_closed_error, 1),
+          t, s, &s->fetching_send_message_finished, GRPC_ERROR_NONE,
           "fetching_send_message_finished");
     } else {
       GPR_ASSERT(s->fetching_send_message == nullptr);
diff --git a/src/core/lib/surface/call.cc b/src/core/lib/surface/call.cc
index 13b1d126393..378d148c674 100644
--- a/src/core/lib/surface/call.cc
+++ b/src/core/lib/surface/call.cc
@@ -1175,6 +1175,12 @@ static void post_batch_completion(batch_control* bctl) {
         &call->metadata_batch[0 /* is_receiving */][0 /* is_trailing */]);
   }
   if (bctl->op.send_message) {
+    if (bctl->op.payload->send_message.stream_write_closed &&
+        error == GRPC_ERROR_NONE) {
+      error = grpc_error_add_child(
+          error, GRPC_ERROR_CREATE_FROM_STATIC_STRING(
+                     "Attempt to send message after stream was closed."));
+    }
     call->sending_message = false;
   }
   if (bctl->op.send_trailing_metadata) {
diff --git a/src/core/lib/transport/transport.h b/src/core/lib/transport/transport.h
index 17608f4bc60..c40b290d535 100644
--- a/src/core/lib/transport/transport.h
+++ b/src/core/lib/transport/transport.h
@@ -245,6 +245,18 @@ struct grpc_transport_stream_op_batch_payload {
     // The batch's on_complete will not be called until after the byte
     // stream is orphaned.
     grpc_core::OrphanablePtr<grpc_core::ByteStream> send_message;
+    // Set by the transport if the stream has been closed for writes. If this
+    // is set and send message op is present, we set the operation to be a
+    // failure without sending a cancel OP down the stack. This is so that the
+    // status of the call does not get overwritten by the Cancel OP, which would
+    // be especially problematic if we had received a valid status from the
+    // server.
+    // For send_initial_metadata, it is fine for the status to be overwritten
+    // because at that point, the client will not have received a status.
+    // For send_trailing_metadata, we might overwrite the status if we have
+    // non-zero metadata to send. This is fine because the API does not allow
+    // the client to send trailing metadata.
+    bool stream_write_closed = false;
   } send_message;
 
   struct {