Merge pull request #17220 from yashykt/failhijackedsend

Add interceptor method to fail hijacked send messages and get status on POST_SEND_MESSAGE
pull/17179/head^2
Yash Tibrewal 6 years ago committed by GitHub
commit 6de81f54bb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 23
      include/grpcpp/impl/codegen/call_op_set.h
  2. 11
      include/grpcpp/impl/codegen/interceptor.h
  3. 27
      include/grpcpp/impl/codegen/interceptor_common.h
  4. 83
      test/cpp/end2end/client_interceptors_end2end_test.cc

@ -326,21 +326,37 @@ class CallOpSendMessage {
// Flags are per-message: clear them after use. // Flags are per-message: clear them after use.
write_options_.Clear(); write_options_.Clear();
} }
void FinishOp(bool* status) { send_buf_.Clear(); } void FinishOp(bool* status) {
if (!send_buf_.Valid()) {
return;
}
if (hijacked_ && failed_send_) {
// Hijacking interceptor failed this Op
*status = false;
} else if (!*status) {
// This Op was passed down to core and the Op failed
failed_send_ = true;
}
}
void SetInterceptionHookPoint( void SetInterceptionHookPoint(
InterceptorBatchMethodsImpl* interceptor_methods) { InterceptorBatchMethodsImpl* interceptor_methods) {
if (!send_buf_.Valid()) return; if (!send_buf_.Valid()) return;
interceptor_methods->AddInterceptionHookPoint( interceptor_methods->AddInterceptionHookPoint(
experimental::InterceptionHookPoints::PRE_SEND_MESSAGE); experimental::InterceptionHookPoints::PRE_SEND_MESSAGE);
interceptor_methods->SetSendMessage(&send_buf_, msg_); interceptor_methods->SetSendMessage(&send_buf_, msg_, &failed_send_);
} }
void SetFinishInterceptionHookPoint( void SetFinishInterceptionHookPoint(
InterceptorBatchMethodsImpl* interceptor_methods) { InterceptorBatchMethodsImpl* interceptor_methods) {
if (send_buf_.Valid()) {
interceptor_methods->AddInterceptionHookPoint(
experimental::InterceptionHookPoints::POST_SEND_MESSAGE);
}
send_buf_.Clear();
// The contents of the SendMessage value that was previously set // The contents of the SendMessage value that was previously set
// has had its references stolen by core's operations // has had its references stolen by core's operations
interceptor_methods->SetSendMessage(nullptr, nullptr); interceptor_methods->SetSendMessage(nullptr, nullptr, &failed_send_);
} }
void SetHijackingState(InterceptorBatchMethodsImpl* interceptor_methods) { void SetHijackingState(InterceptorBatchMethodsImpl* interceptor_methods) {
@ -350,6 +366,7 @@ class CallOpSendMessage {
private: private:
const void* msg_ = nullptr; // The original non-serialized message const void* msg_ = nullptr; // The original non-serialized message
bool hijacked_ = false; bool hijacked_ = false;
bool failed_send_ = false;
ByteBuffer send_buf_; ByteBuffer send_buf_;
WriteOptions write_options_; WriteOptions write_options_;
}; };

@ -46,9 +46,10 @@ namespace experimental {
/// operation has been requested and it is available. POST_RECV means that a /// operation has been requested and it is available. POST_RECV means that a
/// result is available but has not yet been passed back to the application. /// result is available but has not yet been passed back to the application.
enum class InterceptionHookPoints { enum class InterceptionHookPoints {
/// The first two in this list are for clients and servers /// The first three in this list are for clients and servers
PRE_SEND_INITIAL_METADATA, PRE_SEND_INITIAL_METADATA,
PRE_SEND_MESSAGE, PRE_SEND_MESSAGE,
POST_SEND_MESSAGE,
PRE_SEND_STATUS, // server only PRE_SEND_STATUS, // server only
PRE_SEND_CLOSE, // client only: WritesDone for stream; after write in unary PRE_SEND_CLOSE, // client only: WritesDone for stream; after write in unary
/// The following three are for hijacked clients only and can only be /// The following three are for hijacked clients only and can only be
@ -117,6 +118,10 @@ class InterceptorBatchMethods {
/// only supported for sync and callback APIs at the present moment. /// only supported for sync and callback APIs at the present moment.
virtual const void* GetSendMessage() = 0; virtual const void* GetSendMessage() = 0;
/// Checks whether the SEND MESSAGE op succeeded. Valid for POST_SEND_MESSAGE
/// interceptions.
virtual bool GetSendMessageStatus() = 0;
/// Returns a modifiable multimap of the initial metadata to be sent. Valid /// Returns a modifiable multimap of the initial metadata to be sent. Valid
/// for PRE_SEND_INITIAL_METADATA interceptions. A value of nullptr indicates /// for PRE_SEND_INITIAL_METADATA interceptions. A value of nullptr indicates
/// that this field is not valid. /// that this field is not valid.
@ -162,6 +167,10 @@ class InterceptorBatchMethods {
/// started from interceptors without infinite regress through the interceptor /// started from interceptors without infinite regress through the interceptor
/// list. /// list.
virtual std::unique_ptr<ChannelInterface> GetInterceptedChannel() = 0; 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;
}; };
/// Interface for an interceptor. Interceptor authors must create a class /// Interface for an interceptor. Interceptor authors must create a class

@ -83,6 +83,8 @@ class InterceptorBatchMethodsImpl
const void* GetSendMessage() override { return orig_send_message_; } const void* GetSendMessage() override { return orig_send_message_; }
bool GetSendMessageStatus() override { return !*fail_send_message_; }
std::multimap<grpc::string, grpc::string>* GetSendInitialMetadata() override { std::multimap<grpc::string, grpc::string>* GetSendInitialMetadata() override {
return send_initial_metadata_; return send_initial_metadata_;
} }
@ -112,14 +114,22 @@ class InterceptorBatchMethodsImpl
Status* GetRecvStatus() override { return recv_status_; } 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() std::multimap<grpc::string_ref, grpc::string_ref>* GetRecvTrailingMetadata()
override { override {
return recv_trailing_metadata_->map(); return recv_trailing_metadata_->map();
} }
void SetSendMessage(ByteBuffer* buf, const void* msg) { void SetSendMessage(ByteBuffer* buf, const void* msg,
bool* fail_send_message) {
send_message_ = buf; send_message_ = buf;
orig_send_message_ = msg; orig_send_message_ = msg;
fail_send_message_ = fail_send_message;
} }
void SetSendInitialMetadata( void SetSendInitialMetadata(
@ -339,6 +349,7 @@ class InterceptorBatchMethodsImpl
std::function<void(void)> callback_; std::function<void(void)> callback_;
ByteBuffer* send_message_ = nullptr; ByteBuffer* send_message_ = nullptr;
bool* fail_send_message_ = nullptr;
const void* orig_send_message_ = nullptr; const void* orig_send_message_ = nullptr;
std::multimap<grpc::string, grpc::string>* send_initial_metadata_; std::multimap<grpc::string, grpc::string>* send_initial_metadata_;
@ -392,6 +403,14 @@ class CancelInterceptorBatchMethods
return nullptr; return nullptr;
} }
bool GetSendMessageStatus() override {
GPR_CODEGEN_ASSERT(
false &&
"It is illegal to call GetSendMessageStatus on a method which "
"has a Cancel notification");
return false;
}
const void* GetSendMessage() override { const void* GetSendMessage() override {
GPR_CODEGEN_ASSERT( GPR_CODEGEN_ASSERT(
false && false &&
@ -465,6 +484,12 @@ class CancelInterceptorBatchMethods
"method which has a Cancel notification"); "method which has a Cancel notification");
return std::unique_ptr<ChannelInterface>(nullptr); 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 internal
} // namespace grpc } // namespace grpc

@ -339,6 +339,60 @@ class BidiStreamingRpcHijackingInterceptor : public experimental::Interceptor {
grpc::string msg; grpc::string msg;
}; };
class ClientStreamingRpcHijackingInterceptor
: public experimental::Interceptor {
public:
ClientStreamingRpcHijackingInterceptor(experimental::ClientRpcInfo* info) {
info_ = info;
}
virtual void Intercept(experimental::InterceptorBatchMethods* methods) {
bool hijack = false;
if (methods->QueryInterceptionHookPoint(
experimental::InterceptionHookPoints::PRE_SEND_INITIAL_METADATA)) {
hijack = true;
}
if (methods->QueryInterceptionHookPoint(
experimental::InterceptionHookPoints::PRE_SEND_MESSAGE)) {
if (++count_ > 10) {
methods->FailHijackedSendMessage();
}
}
if (methods->QueryInterceptionHookPoint(
experimental::InterceptionHookPoints::POST_SEND_MESSAGE)) {
EXPECT_FALSE(got_failed_send_);
got_failed_send_ = !methods->GetSendMessageStatus();
}
if (methods->QueryInterceptionHookPoint(
experimental::InterceptionHookPoints::PRE_RECV_STATUS)) {
auto* status = methods->GetRecvStatus();
*status = Status(StatusCode::UNAVAILABLE, "Done sending 10 messages");
}
if (hijack) {
methods->Hijack();
} else {
methods->Proceed();
}
}
static bool GotFailedSend() { return got_failed_send_; }
private:
experimental::ClientRpcInfo* info_;
int count_ = 0;
static bool got_failed_send_;
};
bool ClientStreamingRpcHijackingInterceptor::got_failed_send_ = false;
class ClientStreamingRpcHijackingInterceptorFactory
: public experimental::ClientInterceptorFactoryInterface {
public:
virtual experimental::Interceptor* CreateClientInterceptor(
experimental::ClientRpcInfo* info) override {
return new ClientStreamingRpcHijackingInterceptor(info);
}
};
class BidiStreamingRpcHijackingInterceptorFactory class BidiStreamingRpcHijackingInterceptorFactory
: public experimental::ClientInterceptorFactoryInterface { : public experimental::ClientInterceptorFactoryInterface {
public: public:
@ -628,6 +682,35 @@ TEST_F(ClientInterceptorsStreamingEnd2endTest, ServerStreamingTest) {
EXPECT_EQ(DummyInterceptor::GetNumTimesRun(), 20); EXPECT_EQ(DummyInterceptor::GetNumTimesRun(), 20);
} }
TEST_F(ClientInterceptorsStreamingEnd2endTest, ClientStreamingHijackingTest) {
ChannelArguments args;
std::vector<std::unique_ptr<experimental::ClientInterceptorFactoryInterface>>
creators;
creators.push_back(
std::unique_ptr<ClientStreamingRpcHijackingInterceptorFactory>(
new ClientStreamingRpcHijackingInterceptorFactory()));
auto channel = experimental::CreateCustomChannelWithInterceptors(
server_address_, InsecureChannelCredentials(), args, std::move(creators));
auto stub = grpc::testing::EchoTestService::NewStub(channel);
ClientContext ctx;
EchoRequest req;
EchoResponse resp;
req.mutable_param()->set_echo_metadata(true);
req.set_message("Hello");
string expected_resp = "";
auto writer = stub->RequestStream(&ctx, &resp);
for (int i = 0; i < 10; i++) {
EXPECT_TRUE(writer->Write(req));
expected_resp += "Hello";
}
// The interceptor will reject the 11th message
writer->Write(req);
Status s = writer->Finish();
EXPECT_EQ(s.ok(), false);
EXPECT_TRUE(ClientStreamingRpcHijackingInterceptor::GotFailedSend());
}
TEST_F(ClientInterceptorsStreamingEnd2endTest, BidiStreamingHijackingTest) { TEST_F(ClientInterceptorsStreamingEnd2endTest, BidiStreamingHijackingTest) {
ChannelArguments args; ChannelArguments args;
DummyInterceptor::Reset(); DummyInterceptor::Reset();

Loading…
Cancel
Save