SERVER-28679 Set OP_MSG checksum

2 files changed, 16 insertions, 2 deletions

diff --git a/src/mongo/transport/service_state_machine.cpp b/src/mongo/transport/service_state_machine.cpp
index ebcad102106..8e751d9327c 100644
--- a/src/mongo/transport/service_state_machine.cpp
+++ b/src/mongo/transport/service_state_machine.cpp
@@ -140,14 +140,23 @@ Message makeExhaustMessage(Message requestMsg, DbResponse* dbresponse) {
         return Message();
     }
 
-    // Indicate that the response is part of an exhaust stream.
+    const bool checksumPresent = OpMsg::isFlagSet(requestMsg, OpMsg::kChecksumPresent);
+    OpMsg::removeChecksum(&dbresponse->response);
+    // Indicate that the response is part of an exhaust stream. Re-checksum if needed.
     OpMsg::setFlag(&dbresponse->response, OpMsg::kMoreToCome);
+    if (checksumPresent) {
+        OpMsg::appendChecksum(&dbresponse->response);
+    }
 
     // Return an augmented form of the initial request, which is to be used as the next request to
     // be processed by the database. The id of the response is used as the request id of this
-    // 'synthetic' request.
+    // 'synthetic' request. Re-checksum if needed.
+    OpMsg::removeChecksum(&requestMsg);
     requestMsg.header().setId(dbresponse->response.header().getId());
     requestMsg.header().setResponseToMsgId(dbresponse->response.header().getResponseToMsgId());
+    if (checksumPresent) {
+        OpMsg::appendChecksum(&requestMsg);
+    }
     return requestMsg;
 }
 
@@ -454,10 +463,14 @@ void ServiceStateMachine::_processMessage(ThreadGuard guard) {
     Message& toSink = dbresponse.response;
     if (!toSink.empty()) {
         invariant(!OpMsg::isFlagSet(_inMessage, OpMsg::kMoreToCome));
+        invariant(!OpMsg::isFlagSet(toSink, OpMsg::kChecksumPresent));
 
         // Update the header for the response message.
         toSink.header().setId(nextMessageId());
         toSink.header().setResponseToMsgId(_inMessage.header().getId());
+        if (OpMsg::isFlagSet(_inMessage, OpMsg::kChecksumPresent)) {
+            OpMsg::appendChecksum(&toSink);
+        }
 
         // If the incoming message has the exhaust flag set and is a 'getMore' command, then we
         // bypass the normal RPC behavior. We will sink the response to the network, but we also
diff --git a/src/mongo/transport/transport_layer_asio_test.cpp b/src/mongo/transport/transport_layer_asio_test.cpp
index 7bbaab7cd59..45552516ee0 100644
--- a/src/mongo/transport/transport_layer_asio_test.cpp
+++ b/src/mongo/transport/transport_layer_asio_test.cpp
@@ -259,6 +259,7 @@ public:
         Message msg = builder.finish();
         msg.header().setResponseToMsgId(0);
         msg.header().setId(0);
+        OpMsg::appendChecksum(&msg);
 
         std::error_code ec;
         asio::write(_sock, asio::buffer(msg.buf(), msg.size()), ec);