diff --git a/src/brpc/checksum.h b/src/brpc/checksum.h index ff0a98a85b..854b578a25 100644 --- a/src/brpc/checksum.h +++ b/src/brpc/checksum.h @@ -25,8 +25,20 @@ namespace brpc { struct ChecksumIn { + ChecksumIn(const butil::IOBuf* buf, Controller* cntl, + const butil::IOBuf* attachment = NULL) + : buf(buf), cntl(cntl), attachment(attachment) {} + const butil::IOBuf* buf; Controller* cntl; + // Attachment to fold into the checksum in addition to `buf', or NULL if + // the checksum should cover `buf' only (the default, pre-existing + // behavior). When non-NULL, implementations must extend the checksum + // with `buf' first and `attachment' second, so that both sides of the + // RPC (which independently decide whether to pass an attachment here + // based on Controller::request/response_checksum_attachment()) compute + // the same value over the same byte ranges in the same order. + const butil::IOBuf* attachment; }; struct ChecksumHandler { diff --git a/src/brpc/controller.cpp b/src/brpc/controller.cpp index 0bcfb4122d..244d5d237f 100644 --- a/src/brpc/controller.cpp +++ b/src/brpc/controller.cpp @@ -1407,6 +1407,7 @@ void Controller::SaveClientSettings(ClientSettings* s) const { s->connection_type = _connection_type; s->request_compress_type = _request_compress_type; s->request_checksum_type = _request_checksum_type; + s->request_checksum_with_attachment = request_checksum_attachment(); s->log_id = log_id(); s->has_request_code = has_request_code(); s->request_code = _request_code; @@ -1421,6 +1422,7 @@ void Controller::ApplyClientSettings(const ClientSettings& s) { set_connection_type(s.connection_type); set_request_compress_type(s.request_compress_type); set_request_checksum_type(s.request_checksum_type); + set_request_checksum_attachment(s.request_checksum_with_attachment); set_log_id(s.log_id); set_flag(FLAGS_REQUEST_CODE, s.has_request_code); _request_code = s.request_code; diff --git a/src/brpc/controller.h b/src/brpc/controller.h index cb518706ed..5d38de0b0f 100644 --- a/src/brpc/controller.h +++ b/src/brpc/controller.h @@ -145,6 +145,9 @@ friend void policy::ProcessThriftRequest(InputMessageBase*); static const uint32_t FLAGS_READ_PROGRESSIVELY = (1 << 3); static const uint32_t FLAGS_PROGRESSIVE_READER = (1 << 4); static const uint32_t FLAGS_BACKUP_REQUEST = (1 << 5); + // Whether set_request_checksum_type()'s checksum also covers the + // attachment. See set_request_checksum_attachment(). + static const uint32_t FLAGS_REQUEST_CHECKSUM_WITH_ATTACHMENT = (1 << 6); // Let _done delete the correlation_id, used by combo channels to // make lifetime of the correlation_id more flexible. static const uint32_t FLAGS_DESTROY_CID_IN_DONE = (1 << 7); @@ -167,6 +170,9 @@ friend void policy::ProcessThriftRequest(InputMessageBase*); static const uint32_t FLAGS_MANAGE_HTTP_BODY_ON_ERROR = (1 << 21); static const uint32_t FLAGS_WRITE_TO_SOCKET_IN_BACKGROUND = (1 << 22); static const uint32_t FLAGS_ENDING_RPC = (1 << 23); + // Whether set_response_checksum_type()'s checksum also covers the + // attachment. See set_response_checksum_attachment(). + static const uint32_t FLAGS_RESPONSE_CHECKSUM_WITH_ATTACHMENT = (1 << 24); public: struct Inheritable { @@ -260,6 +266,22 @@ friend void policy::ProcessThriftRequest(InputMessageBase*); // Set checksum type for request. void set_request_checksum_type(ChecksumType t) { _request_checksum_type = t; } + // Whether the request's checksum (see set_request_checksum_type) also + // covers the attachment, in addition to the serialized body. Defaults + // to false (checksum covers body only), preserving the pre-existing + // behavior. This setting is sent to the peer along with the request so + // that it recomputes the checksum over the same range; it is meaningless + // (and rejected, see baidu_rpc_protocol.cpp) together with + // request_will_be_read_progressively() since the attachment is not + // fully buffered before the checksum must be verified. + // NOTE: Only enable this if the peer is known to understand the + // checksum_with_attachment field (baidu_std protocol). An older peer + // silently ignores the field and verifies against the body only, which + // will then mismatch the body+attachment checksum computed here. + void set_request_checksum_attachment(bool with_attachment) { + set_flag(FLAGS_REQUEST_CHECKSUM_WITH_ATTACHMENT, with_attachment); + } + // Required by some load balancers. void set_request_code(uint64_t request_code) { add_flag(FLAGS_REQUEST_CODE); @@ -486,6 +508,14 @@ friend void policy::ProcessThriftRequest(InputMessageBase*); // Set checksum type for response. void set_response_checksum_type(ChecksumType t) { _response_checksum_type = t; } + + // Whether the response's checksum (see set_response_checksum_type) also + // covers the attachment, in addition to the serialized body. See + // set_request_checksum_attachment() for details; the same caveat about + // progressive attachment reading applies here. + void set_response_checksum_attachment(bool with_attachment) { + set_flag(FLAGS_RESPONSE_CHECKSUM_WITH_ATTACHMENT, with_attachment); + } // Non-zero when this RPC call is traced (by rpcz or rig). // NOTE: Only valid at server-side, always zero at client-side. @@ -576,6 +606,8 @@ friend void policy::ProcessThriftRequest(InputMessageBase*); CompressType response_compress_type() const { return _response_compress_type; } ChecksumType request_checksum_type() const { return _request_checksum_type; } ChecksumType response_checksum_type() const { return _response_checksum_type; } + bool request_checksum_attachment() const { return has_flag(FLAGS_REQUEST_CHECKSUM_WITH_ATTACHMENT); } + bool response_checksum_attachment() const { return has_flag(FLAGS_RESPONSE_CHECKSUM_WITH_ATTACHMENT); } const HttpHeader& http_request() const { return _http_request != NULL ? *_http_request : DefaultHttpHeader(); } @@ -729,6 +761,7 @@ friend void policy::ProcessThriftRequest(InputMessageBase*); ConnectionType connection_type; CompressType request_compress_type; ChecksumType request_checksum_type; + bool request_checksum_with_attachment; uint64_t log_id; bool has_request_code; int64_t request_code; diff --git a/src/brpc/policy/baidu_rpc_meta.proto b/src/brpc/policy/baidu_rpc_meta.proto index 5591c5dab7..f78066f41d 100644 --- a/src/brpc/policy/baidu_rpc_meta.proto +++ b/src/brpc/policy/baidu_rpc_meta.proto @@ -36,6 +36,7 @@ message RpcMeta { optional ContentType content_type = 10; optional int32 checksum_type = 11; optional bytes checksum_value = 12; + optional bool checksum_with_attachment = 13; } message RpcRequestMeta { diff --git a/src/brpc/policy/baidu_rpc_protocol.cpp b/src/brpc/policy/baidu_rpc_protocol.cpp index 2c5a7e7224..49863b2c06 100644 --- a/src/brpc/policy/baidu_rpc_protocol.cpp +++ b/src/brpc/policy/baidu_rpc_protocol.cpp @@ -148,7 +148,8 @@ ParseResult ParseRpcMessage(butil::IOBuf* source, Socket* socket, bool SerializeRpcMessage(const google::protobuf::Message& message, Controller& cntl, ContentType content_type, CompressType compress_type, ChecksumType checksum_type, - butil::IOBuf* buf) { + butil::IOBuf* buf, + const butil::IOBuf* checksum_attachment) { auto serialize = [&](Serializer& serializer) -> bool { bool ok; if (COMPRESS_TYPE_NONE == compress_type) { @@ -161,7 +162,7 @@ bool SerializeRpcMessage(const google::protobuf::Message& message, } ok = handler->Compress(serializer, buf); } - ChecksumIn checksum_in{buf, &cntl}; + ChecksumIn checksum_in{buf, &cntl, checksum_attachment}; ComputeDataChecksum(checksum_in, checksum_type); return ok; }; @@ -231,8 +232,16 @@ static bool SerializeResponse(const google::protobuf::Message& res, ContentType content_type = cntl.response_content_type(); CompressType compress_type = cntl.response_compress_type(); ChecksumType checksum_type = cntl.response_checksum_type(); + const butil::IOBuf* checksum_attachment = NULL; + if (cntl.response_checksum_attachment()) { + // See the same check in SerializeRpcRequest() for the rationale; + // baidu_std never sets this flag itself but we defend anyway. + if (!cntl.is_response_read_progressively()) { + checksum_attachment = &cntl.response_attachment(); + } + } if (!SerializeRpcMessage(res, cntl, content_type, compress_type, - checksum_type, &buf)) { + checksum_type, &buf, checksum_attachment)) { cntl.SetFailed(ERESPONSE, "Fail to serialize response=%s, " "ContentType=%s, CompressType=%s, ChecksumType=%s", @@ -349,6 +358,9 @@ void SendRpcResponse(int64_t correlation_id, Controller* cntl, meta.set_content_type(cntl->response_content_type()); meta.set_checksum_type(cntl->response_checksum_type()); meta.set_checksum_value(accessor.checksum_value()); + if (cntl->response_checksum_attachment()) { + meta.set_checksum_with_attachment(true); + } if (attached_size > 0) { meta.set_attachment_size(attached_size); } @@ -498,9 +510,10 @@ void EndRunningCallMethodInPool( bool DeserializeRpcMessage(const butil::IOBuf& data, Controller& cntl, ContentType content_type, CompressType compress_type, ChecksumType checksum_type, - google::protobuf::Message* message) { + google::protobuf::Message* message, + const butil::IOBuf* checksum_attachment) { auto deserialize = [&](Deserializer& deserializer) -> bool { - ChecksumIn checksum_in{&data, &cntl}; + ChecksumIn checksum_in{&data, &cntl, checksum_attachment}; bool ok = VerifyDataChecksum(checksum_in, checksum_type); if (!ok) { return ok; @@ -618,6 +631,7 @@ void ProcessRpcRequest(InputMessageBase* msg_base) { cntl->set_request_content_type(meta.content_type()); cntl->set_request_compress_type((CompressType)meta.compress_type()); cntl->set_request_checksum_type((ChecksumType)meta.checksum_type()); + cntl->set_request_checksum_attachment(meta.checksum_with_attachment()); cntl->set_rpc_received_us(msg->received_us()); accessor.set_checksum_value(meta.checksum_value()); accessor.set_server(server) @@ -808,9 +822,16 @@ void ProcessRpcRequest(InputMessageBase* msg_base) { static_cast(meta.checksum_type()); messages = server->options().rpc_pb_message_factory->Get(*svc, *method); + // request_attachment() has already been filled in above (swapped + // out of msg->payload) before we get here, so it's safe to fold + // it into the checksum now when the client asked us to. + const butil::IOBuf* checksum_attachment = + cntl->request_checksum_attachment() ? + &cntl->request_attachment() : NULL; if (!DeserializeRpcMessage(req_buf, *cntl, content_type, compress_type, checksum_type, - messages->Request())) { + messages->Request(), + checksum_attachment)) { cntl->SetFailed( EREQUEST, "Fail to parse request=%s, ContentType=%s, " @@ -987,14 +1008,22 @@ void ProcessRpcResponse(InputMessageBase* msg_base) { cntl->set_response_content_type(content_type); cntl->set_response_compress_type(compress_type); cntl->set_response_checksum_type(checksum_type); + cntl->set_response_checksum_attachment(meta.checksum_with_attachment()); accessor.set_checksum_value(meta.checksum_value()); if (cntl->response()) { + // response_attachment() has already been filled in above (swapped + // out of msg->payload) before we get here, so it's safe to fold + // it into the checksum now when the server told us to. + const butil::IOBuf* checksum_attachment = + cntl->response_checksum_attachment() ? + &cntl->response_attachment() : NULL; if (cntl->response()->GetDescriptor() == SerializedResponse::descriptor()) { ((SerializedResponse*)cntl->response())-> serialized_data().append(*res_buf_ptr); } else if (!DeserializeRpcMessage(*res_buf_ptr, *cntl, content_type, compress_type, checksum_type, - cntl->response())) { + cntl->response(), + checksum_attachment)) { cntl->SetFailed( EREQUEST, "Fail to parse response=%s, ContentType=%s, " @@ -1030,8 +1059,19 @@ void SerializeRpcRequest(butil::IOBuf* request_buf, Controller* cntl, ContentType content_type = cntl->request_content_type(); CompressType compress_type = cntl->request_compress_type(); ChecksumType checksum_type = cntl->request_checksum_type(); + const butil::IOBuf* checksum_attachment = NULL; + if (cntl->request_checksum_attachment()) { + // Progressive reading (HTTP-only feature) hands the attachment to + // the user piece by piece as it arrives, so there's no single, + // complete IOBuf to fold into the checksum here. baidu_std (this + // protocol) never sets FLAGS_READ_PROGRESSIVELY itself, but guard + // against a Controller that's reused/misconfigured across protocols. + if (!cntl->is_response_read_progressively()) { + checksum_attachment = &cntl->request_attachment(); + } + } if (!SerializeRpcMessage(*request, *cntl, content_type, compress_type, - checksum_type, request_buf)) { + checksum_type, request_buf, checksum_attachment)) { return cntl->SetFailed( EREQUEST, "Fail to compress request=%s, " @@ -1065,6 +1105,9 @@ void PackRpcRequest(butil::IOBuf* req_buf, meta.set_compress_type(cntl->request_compress_type()); meta.set_checksum_type(cntl->request_checksum_type()); meta.set_checksum_value(accessor.checksum_value()); + if (cntl->request_checksum_attachment()) { + meta.set_checksum_with_attachment(true); + } } else if (NULL != cntl->sampled_request()) { // Replaying. Keep service-name as the one seen by server. request_meta->set_service_name(cntl->sampled_request()->meta.service_name()); diff --git a/src/brpc/policy/crc32c_checksum.cpp b/src/brpc/policy/crc32c_checksum.cpp index 28b6fab60b..7a3b8ef9d7 100644 --- a/src/brpc/policy/crc32c_checksum.cpp +++ b/src/brpc/policy/crc32c_checksum.cpp @@ -25,33 +25,44 @@ namespace brpc { namespace policy { -void Crc32cCompute(const ChecksumIn& in) { - auto buf = in.buf; - auto cntl = in.cntl; - butil::IOBufAsZeroCopyInputStream wrapper(*buf); +namespace { + +// Extends `crc' over every block of `buf'. Shared by Crc32cCompute and +// Crc32cVerify to keep the byte range/order (`buf' then `attachment', see +// ChecksumIn::attachment) identical on both call sites. +uint32_t ExtendCrc32c(uint32_t crc, const butil::IOBuf& buf) { + butil::IOBufAsZeroCopyInputStream wrapper(buf); const void* data; int size; - uint32_t crc = 0; while (wrapper.Next(&data, &size)) { crc = butil::crc32c::Extend(crc, static_cast(data), size); } + return crc; +} + +// Computes the crc32c over `in.buf', and over `in.attachment' as well when +// the caller opted in (ChecksumIn::attachment != NULL). +uint32_t ComputeCrc32c(const ChecksumIn& in) { + uint32_t crc = ExtendCrc32c(0, *in.buf); + if (in.attachment != NULL) { + crc = ExtendCrc32c(crc, *in.attachment); + } + return crc; +} + +} // namespace + +void Crc32cCompute(const ChecksumIn& in) { + uint32_t crc = ComputeCrc32c(in); RPC_VLOG << "Crc32cCompute crc=" << crc; crc = butil::HostToNet32(butil::crc32c::Mask(crc)); - ControllerPrivateAccessor(cntl).set_checksum_value( + ControllerPrivateAccessor(in.cntl).set_checksum_value( reinterpret_cast(&crc), sizeof(crc)); } bool Crc32cVerify(const ChecksumIn& in) { - auto buf = in.buf; - auto cntl = in.cntl; - butil::IOBufAsZeroCopyInputStream wrapper(*buf); - const void* data; - int size; - uint32_t crc = 0; - while (wrapper.Next(&data, &size)) { - crc = butil::crc32c::Extend(crc, static_cast(data), size); - } - auto& val = ControllerPrivateAccessor(const_cast(cntl)) + uint32_t crc = ComputeCrc32c(in); + auto& val = ControllerPrivateAccessor(const_cast(in.cntl)) .checksum_value(); CHECK_EQ(val.size(), sizeof(crc)); auto expected = *reinterpret_cast(val.data()); diff --git a/test/brpc_checksum_unittest.cpp b/test/brpc_checksum_unittest.cpp new file mode 100644 index 0000000000..f8b86b33ad --- /dev/null +++ b/test/brpc_checksum_unittest.cpp @@ -0,0 +1,244 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Unit tests for Controller::set_request/response_checksum_attachment(), +// i.e. extending checksum coverage (see set_request/response_checksum_type()) +// to also fold in the attachment. + +#include +#include "brpc/channel.h" +#include "brpc/controller.h" +#include "brpc/server.h" +#include "brpc/checksum.h" +#include "brpc/details/controller_private_accessor.h" +#include "brpc/policy/crc32c_checksum.h" +#include "butil/iobuf.h" +#include "echo.pb.h" + +namespace { + +// --------------------------------------------------------------------- +// Focused unit tests on Crc32cCompute/Crc32cVerify: exercise +// ChecksumIn::attachment directly without going through the network, so +// that we can precisely control and corrupt the bytes being checksummed. +// --------------------------------------------------------------------- +class ChecksumAttachmentTest : public ::testing::Test {}; + +TEST_F(ChecksumAttachmentTest, verify_succeeds_when_body_only) { + brpc::Controller cntl; + butil::IOBuf body; + body.append("request body"); + + brpc::ChecksumIn compute_in{&body, &cntl, NULL}; + brpc::policy::Crc32cCompute(compute_in); + + brpc::ChecksumIn verify_in{&body, &cntl, NULL}; + EXPECT_TRUE(brpc::policy::Crc32cVerify(verify_in)); +} + +TEST_F(ChecksumAttachmentTest, verify_succeeds_when_attachment_included) { + brpc::Controller cntl; + butil::IOBuf body; + body.append("request body"); + butil::IOBuf attachment; + attachment.append("some attachment bytes"); + + brpc::ChecksumIn compute_in{&body, &cntl, &attachment}; + brpc::policy::Crc32cCompute(compute_in); + + brpc::ChecksumIn verify_in{&body, &cntl, &attachment}; + EXPECT_TRUE(brpc::policy::Crc32cVerify(verify_in)); +} + +TEST_F(ChecksumAttachmentTest, verify_fails_when_attachment_corrupted) { + brpc::Controller cntl; + butil::IOBuf body; + body.append("request body"); + butil::IOBuf attachment; + attachment.append("some attachment bytes"); + + brpc::ChecksumIn compute_in{&body, &cntl, &attachment}; + brpc::policy::Crc32cCompute(compute_in); + + // Attacker/bit-flip changes one byte of the attachment in flight. + butil::IOBuf corrupted_attachment; + corrupted_attachment.append("some attachment Bytes"); + brpc::ChecksumIn verify_in{&body, &cntl, &corrupted_attachment}; + EXPECT_FALSE(brpc::policy::Crc32cVerify(verify_in)); +} + +TEST_F(ChecksumAttachmentTest, verify_fails_when_attachment_dropped) { + brpc::Controller cntl; + butil::IOBuf body; + body.append("request body"); + butil::IOBuf attachment; + attachment.append("some attachment bytes"); + + // Sender includes the attachment in the checksum... + brpc::ChecksumIn compute_in{&body, &cntl, &attachment}; + brpc::policy::Crc32cCompute(compute_in); + + // ...but receiver (e.g. due to a mismatched + // request/response_checksum_attachment() setting) verifies body only. + brpc::ChecksumIn verify_in{&body, &cntl, NULL}; + EXPECT_FALSE(brpc::policy::Crc32cVerify(verify_in)); +} + +TEST_F(ChecksumAttachmentTest, verify_fails_when_body_and_attachment_swapped) { + // Sanity check that the two parts are not simply concatenated in a + // position-independent way, i.e. order matters and is deterministic. + brpc::Controller cntl; + butil::IOBuf part_a; + part_a.append("AAAA"); + butil::IOBuf part_b; + part_b.append("BBBBBB"); + + brpc::ChecksumIn compute_in{&part_a, &cntl, &part_b}; + brpc::policy::Crc32cCompute(compute_in); + + brpc::ChecksumIn verify_in{&part_b, &cntl, &part_a}; + EXPECT_FALSE(brpc::policy::Crc32cVerify(verify_in)); +} + +// --------------------------------------------------------------------- +// End-to-end test that Controller::set_request/response_checksum_attachment +// is wired all the way through baidu_std serialization/wire format/ +// deserialization without breaking normal RPCs. +// --------------------------------------------------------------------- +class ChecksumEchoServiceImpl : public test::EchoService { +public: + void Echo(google::protobuf::RpcController* cntl_base, + const test::EchoRequest* req, + test::EchoResponse* res, + google::protobuf::Closure* done) override { + brpc::ClosureGuard done_guard(done); + brpc::Controller* cntl = static_cast(cntl_base); + res->set_message("received " + req->message()); + // Echo the attachment back and ask the checksum (set by the test + // below) to also cover it, exercising the response-side path. + cntl->response_attachment().append(cntl->request_attachment()); + if (response_checksum_type_ != brpc::CHECKSUM_TYPE_NONE) { + cntl->set_response_checksum_type(response_checksum_type_); + cntl->set_response_checksum_attachment( + response_checksum_with_attachment_); + } + } + + void set_response_checksum(brpc::ChecksumType type, bool with_attachment) { + response_checksum_type_ = type; + response_checksum_with_attachment_ = with_attachment; + } + +private: + brpc::ChecksumType response_checksum_type_ = brpc::CHECKSUM_TYPE_NONE; + bool response_checksum_with_attachment_ = false; +}; + +class ChecksumAttachmentEndToEndTest : public ::testing::Test { +protected: + void SetUp() override { + ASSERT_EQ(0, server_.AddService(&svc_, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server_.Start(port_, NULL)); + brpc::ChannelOptions options; + ASSERT_EQ(0, channel_.Init(butil::EndPoint(butil::my_ip(), port_), + &options)); + } + + void TearDown() override { + server_.Stop(0); + server_.Join(); + } + + const int port_ = 8934; + brpc::Server server_; + ChecksumEchoServiceImpl svc_; + brpc::Channel channel_; +}; + +TEST_F(ChecksumAttachmentEndToEndTest, request_checksum_with_attachment) { + svc_.set_response_checksum(brpc::CHECKSUM_TYPE_NONE, false); + + brpc::Controller cntl; + cntl.request_attachment().append("request attachment payload"); + cntl.set_request_checksum_type(brpc::CHECKSUM_TYPE_CRC32C); + cntl.set_request_checksum_attachment(true); + + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + test::EchoService::Stub(&channel_).Echo(&cntl, &req, &res, NULL); + + ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText(); + EXPECT_EQ("received " + std::string(__FUNCTION__), res.message()); + EXPECT_EQ("request attachment payload", cntl.response_attachment()); +} + +TEST_F(ChecksumAttachmentEndToEndTest, response_checksum_with_attachment) { + svc_.set_response_checksum(brpc::CHECKSUM_TYPE_CRC32C, true); + + brpc::Controller cntl; + cntl.request_attachment().append("round-trip payload"); + + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + test::EchoService::Stub(&channel_).Echo(&cntl, &req, &res, NULL); + + ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText(); + EXPECT_EQ("received " + std::string(__FUNCTION__), res.message()); + EXPECT_EQ("round-trip payload", cntl.response_attachment()); +} + +TEST_F(ChecksumAttachmentEndToEndTest, both_directions_checksum_with_attachment) { + svc_.set_response_checksum(brpc::CHECKSUM_TYPE_CRC32C, true); + + brpc::Controller cntl; + cntl.request_attachment().append("both directions payload"); + cntl.set_request_checksum_type(brpc::CHECKSUM_TYPE_CRC32C); + cntl.set_request_checksum_attachment(true); + + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + test::EchoService::Stub(&channel_).Echo(&cntl, &req, &res, NULL); + + ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText(); + EXPECT_EQ("received " + std::string(__FUNCTION__), res.message()); + EXPECT_EQ("both directions payload", cntl.response_attachment()); +} + +TEST_F(ChecksumAttachmentEndToEndTest, checksum_without_attachment_still_works) { + // Sanity check that turning the new flag off preserves the pre-existing + // body-only checksum behavior. + svc_.set_response_checksum(brpc::CHECKSUM_TYPE_CRC32C, false); + + brpc::Controller cntl; + cntl.request_attachment().append("legacy behavior payload"); + cntl.set_request_checksum_type(brpc::CHECKSUM_TYPE_CRC32C); + // Deliberately not calling set_request_checksum_attachment(true). + + test::EchoRequest req; + test::EchoResponse res; + req.set_message(__FUNCTION__); + test::EchoService::Stub(&channel_).Echo(&cntl, &req, &res, NULL); + + ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText(); + EXPECT_EQ("received " + std::string(__FUNCTION__), res.message()); + EXPECT_EQ("legacy behavior payload", cntl.response_attachment()); +} + +} // namespace diff --git a/test/brpc_server_unittest.cpp b/test/brpc_server_unittest.cpp index 8508a7986c..1012ce25d7 100644 --- a/test/brpc_server_unittest.cpp +++ b/test/brpc_server_unittest.cpp @@ -74,12 +74,14 @@ DECLARE_bool(use_http_error_code); extern bool SerializeRpcMessage(const google::protobuf::Message& serializer, Controller& cntl, ContentType content_type, CompressType compress_type, - ChecksumType checksum_type, butil::IOBuf* buf); + ChecksumType checksum_type, butil::IOBuf* buf, + const butil::IOBuf* checksum_attachment = NULL); extern bool DeserializeRpcMessage(const butil::IOBuf& deserializer, Controller& cntl, ContentType content_type, CompressType compress_type, ChecksumType checksum_type, - google::protobuf::Message* message); + google::protobuf::Message* message, + const butil::IOBuf* checksum_attachment = NULL); } }