From 11b7d80049ed7482b73bb7c80f0a1996f4a7cbfb Mon Sep 17 00:00:00 2001 From: Yang Liming Date: Fri, 24 Jul 2026 18:24:59 +0800 Subject: [PATCH] Support checksumming the attachment together with the body Controller::set_request/response_checksum_type() previously only covered the serialized protobuf body; the attachment (if any) was never protected. Add set_request/response_checksum_attachment(bool) so callers can opt the attachment into the same checksum. - baidu_rpc_meta.proto: add RpcMeta.checksum_with_attachment so the receiver knows whether to fold the attachment into verification. Defaults to false, so old peers that don't understand the field keep verifying against the body only (backward compatible). - Controller: add the two setters/getters, thread the flag through ClientSettings (Save/ApplyClientSettings) so ParallelChannel/ SelectiveChannel sub-controllers inherit it correctly, and reset it in ResetPods(). - ChecksumIn: add an optional `attachment' field consumed by checksum handlers. - crc32c_checksum.cpp: extend the crc32c over body then attachment (in that fixed order) when requested. - baidu_rpc_protocol.cpp: wire checksum_attachment through SerializeRpcMessage/DeserializeRpcMessage and every client/server send/receive path; skip it when progressive attachment reading is enabled since there's no single complete IOBuf to checksum in that case. Add brpc_checksum_unittest.cpp covering Crc32cCompute/Crc32cVerify directly (including corruption/omission/order sensitivity) and an end-to-end Server/Channel test for both request- and response-side attachment checksums. --- src/brpc/checksum.h | 12 ++ src/brpc/controller.cpp | 2 + src/brpc/controller.h | 33 ++++ src/brpc/policy/baidu_rpc_meta.proto | 1 + src/brpc/policy/baidu_rpc_protocol.cpp | 59 +++++- src/brpc/policy/crc32c_checksum.cpp | 43 +++-- test/brpc_checksum_unittest.cpp | 244 +++++++++++++++++++++++++ test/brpc_server_unittest.cpp | 6 +- 8 files changed, 374 insertions(+), 26 deletions(-) create mode 100644 test/brpc_checksum_unittest.cpp 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); } }