Support checksumming the attachment together with the body

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 {

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;

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;

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 {

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<ChecksumType>(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());

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<const char*>(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<char*>(&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<const char*>(data), size);
-    }
-    auto& val = ControllerPrivateAccessor(const_cast<Controller*>(cntl))
+    uint32_t crc = ComputeCrc32c(in);
+    auto& val = ControllerPrivateAccessor(const_cast<Controller*>(in.cntl))
                     .checksum_value();
     CHECK_EQ(val.size(), sizeof(crc));
     auto expected = *reinterpret_cast<const uint32_t*>(val.data());