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.
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(
InterceptorBatchMethodsImpl* interceptor_methods) {
if (!send_buf_.Valid()) return;
interceptor_methods->AddInterceptionHookPoint(
experimental::InterceptionHookPoints::PRE_SEND_MESSAGE);
interceptor_methods->SetSendMessage(&send_buf_, msg_);
interceptor_methods->SetSendMessage(&send_buf_, msg_, &failed_send_);
}
void SetFinishInterceptionHookPoint(
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
// 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) {
@ -350,6 +366,7 @@ class CallOpSendMessage {
private:
const void* msg_ = nullptr; // The original non-serialized message
bool hijacked_ = false;
bool failed_send_ = false;
ByteBuffer send_buf_;
WriteOptions write_options_;
};

@ -46,9 +46,10 @@ namespace experimental {
/// 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.
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_MESSAGE,
POST_SEND_MESSAGE,
PRE_SEND_STATUS, // server only
PRE_SEND_CLOSE, // client only: WritesDone for stream; after write in unary
/// 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.
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
/// for PRE_SEND_INITIAL_METADATA interceptions. A value of nullptr indicates
/// that this field is not valid.
@ -162,6 +167,10 @@ class InterceptorBatchMethods {
/// started from interceptors without infinite regress through the interceptor
/// list.
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

@ -83,6 +83,8 @@ class InterceptorBatchMethodsImpl
const void* GetSendMessage() override { return orig_send_message_; }
bool GetSendMessageStatus() override { return !*fail_send_message_; }
std::multimap<grpc::string, grpc::string>* GetSendInitialMetadata() override {
return send_initial_metadata_;
}
@ -112,14 +114,22 @@ 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, const void* msg) {
void SetSendMessage(ByteBuffer* buf, const void* msg,
bool* fail_send_message) {
send_message_ = buf;
orig_send_message_ = msg;
fail_send_message_ = fail_send_message;
}
void SetSendInitialMetadata(
@ -339,6 +349,7 @@ class InterceptorBatchMethodsImpl
std::function<void(void)> callback_;
ByteBuffer* send_message_ = nullptr;
bool* fail_send_message_ = nullptr;
const void* orig_send_message_ = nullptr;
std::multimap<grpc::string, grpc::string>* send_initial_metadata_;
@ -392,6 +403,14 @@ class CancelInterceptorBatchMethods
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 {
GPR_CODEGEN_ASSERT(
false &&
@ -465,6 +484,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

@ -339,6 +339,60 @@ class BidiStreamingRpcHijackingInterceptor : public experimental::Interceptor {
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
: public experimental::ClientInterceptorFactoryInterface {
public:
@ -628,6 +682,35 @@ TEST_F(ClientInterceptorsStreamingEnd2endTest, ServerStreamingTest) {
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) {
ChannelArguments args;
DummyInterceptor::Reset();

Loading…
Cancel
Save