Add method to fail hijacked send messages

3 files changed, 29 insertions, 3 deletions

diff --git a/include/grpcpp/impl/codegen/call_op_set.h b/include/grpcpp/impl/codegen/call_op_set.h
index b4c34a01c9..1f2b88e9e1 100644
--- a/include/grpcpp/impl/codegen/call_op_set.h
+++ b/include/grpcpp/impl/codegen/call_op_set.h
@@ -314,14 +314,19 @@ class CallOpSendMessage {
     // Flags are per-message: clear them after use.
     write_options_.Clear();
   }
-  void FinishOp(bool* status) { send_buf_.Clear(); }
+  void FinishOp(bool* status) {
+    send_buf_.Clear();
+    if (hijacked_ && failed_send_) {
+      *status = false;
+    }
+  }
 
   void SetInterceptionHookPoint(
       InterceptorBatchMethodsImpl* interceptor_methods) {
     if (!send_buf_.Valid()) return;
     interceptor_methods->AddInterceptionHookPoint(
         experimental::InterceptionHookPoints::PRE_SEND_MESSAGE);
-    interceptor_methods->SetSendMessage(&send_buf_);
+    interceptor_methods->SetSendMessage(&send_buf_, &failed_send_);
   }
 
   void SetFinishInterceptionHookPoint(
@@ -333,6 +338,7 @@ class CallOpSendMessage {
 
  private:
   bool hijacked_ = false;
+  bool failed_send_ = false;
   ByteBuffer send_buf_;
   WriteOptions write_options_;
 };
diff --git a/include/grpcpp/impl/codegen/interceptor.h b/include/grpcpp/impl/codegen/interceptor.h
index e449e44a23..47239332c8 100644
--- a/include/grpcpp/impl/codegen/interceptor.h
+++ b/include/grpcpp/impl/codegen/interceptor.h
@@ -118,6 +118,10 @@ class InterceptorBatchMethods {
   // only interceptors after the current interceptor are created from the
   // factory objects registered with the channel.
   virtual std::unique_ptr<ChannelInterface> GetInterceptedChannel() = 0;
+
+  // On a hijacked RPC/ to-be hijacked RPC, this can be called to fail a SEND
+  // MESSAGE op
+  virtual void FailHijackedSendMessage() = 0;
 };
 
 class Interceptor {
diff --git a/include/grpcpp/impl/codegen/interceptor_common.h b/include/grpcpp/impl/codegen/interceptor_common.h
index d0aa23cb0a..601a929afe 100644
--- a/include/grpcpp/impl/codegen/interceptor_common.h
+++ b/include/grpcpp/impl/codegen/interceptor_common.h
@@ -110,12 +110,21 @@ class InterceptorBatchMethodsImpl
 
   Status* GetRecvStatus() override { return recv_status_; }
 
+  void FailHijackedSendMessage() override {
+    GPR_CODEGEN_ASSERT(hooks_[static_cast<size_t>(
+        experimental::InterceptionHookPoints::PRE_SEND_MESSAGE)]);
+    *fail_send_message_ = true;
+  }
+
   std::multimap<grpc::string_ref, grpc::string_ref>* GetRecvTrailingMetadata()
       override {
     return recv_trailing_metadata_->map();
   }
 
-  void SetSendMessage(ByteBuffer* buf) { send_message_ = buf; }
+  void SetSendMessage(ByteBuffer* buf, bool* fail_send_message) {
+    send_message_ = buf;
+    fail_send_message_ = fail_send_message;
+  }
 
   void SetSendInitialMetadata(
       std::multimap<grpc::string, grpc::string>* metadata) {
@@ -334,6 +343,7 @@ class InterceptorBatchMethodsImpl
   std::function<void(void)> callback_;
 
   ByteBuffer* send_message_ = nullptr;
+  bool* fail_send_message_ = nullptr;
 
   std::multimap<grpc::string, grpc::string>* send_initial_metadata_;
 
@@ -451,6 +461,12 @@ class CancelInterceptorBatchMethods
                        "method which has a Cancel notification");
     return std::unique_ptr<ChannelInterface>(nullptr);
   }
+
+  void FailHijackedSendMessage() override {
+    GPR_CODEGEN_ASSERT(false &&
+                       "It is illegal to call FailHijackedSendMessage on a "
+                       "method which has a Cancel notification");
+  }
 };
 }  // namespace internal
 }  // namespace grpc