SERVER-28438 Add keyId field to logicalTime object

3 files changed, 85 insertions, 11 deletions

diff --git a/src/mongo/rpc/metadata/logical_time_metadata.cpp b/src/mongo/rpc/metadata/logical_time_metadata.cpp
index 9d79a892e3a..8abdd91edd4 100644
--- a/src/mongo/rpc/metadata/logical_time_metadata.cpp
+++ b/src/mongo/rpc/metadata/logical_time_metadata.cpp
@@ -40,6 +40,8 @@ namespace {
 
 const char kClusterTimeFieldName[] = "clusterTime";
 const char kSignatureFieldName[] = "signature";
+const char kSignatureHashFieldName[] = "hash";
+const char kSignatureKeyIdFieldName[] = "keyId";
 
 }  // unnamed namespace
 
@@ -64,30 +66,49 @@ StatusWith<LogicalTimeMetadata> LogicalTimeMetadata::readFromMetadata(
         return status;
     }
 
-    // Extract BinData type signature and construct a SHA1Block instance from it.
     BSONElement signatureElem;
-    status = bsonExtractTypedField(obj, kSignatureFieldName, BinData, &signatureElem);
+    status = bsonExtractTypedField(obj, kSignatureFieldName, Object, &signatureElem);
+    if (!status.isOK()) {
+        return status;
+    }
+
+    const auto& signatureObj = signatureElem.Obj();
+
+    // Extract BinData type signature hash and construct a SHA1Block instance from it.
+    BSONElement hashElem;
+    status = bsonExtractTypedField(signatureObj, kSignatureHashFieldName, BinData, &hashElem);
     if (!status.isOK()) {
         return status;
     }
 
     int hashLength = 0;
-    auto rawBinSignature = signatureElem.binData(hashLength);
-    BSONBinData proofBinData(rawBinSignature, hashLength, signatureElem.binDataType());
+    auto rawBinSignature = hashElem.binData(hashLength);
+    BSONBinData proofBinData(rawBinSignature, hashLength, hashElem.binDataType());
     auto proofStatus = SHA1Block::fromBinData(proofBinData);
 
     if (!proofStatus.isOK()) {
         return proofStatus.getStatus();
     }
 
+    long long keyId;
+    status = bsonExtractIntegerField(signatureObj, kSignatureKeyIdFieldName, &keyId);
+    if (!status.isOK()) {
+        return status;
+    }
+
     return LogicalTimeMetadata(
-        SignedLogicalTime(LogicalTime(ts), std::move(proofStatus.getValue())));
+        SignedLogicalTime(LogicalTime(ts), std::move(proofStatus.getValue()), keyId));
 }
 
 void LogicalTimeMetadata::writeToMetadata(BSONObjBuilder* metadataBuilder) const {
     BSONObjBuilder subObjBuilder(metadataBuilder->subobjStart(fieldName()));
     _clusterTime.getTime().asTimestamp().append(subObjBuilder.bb(), kClusterTimeFieldName);
-    _clusterTime.getProof().appendAsBinData(subObjBuilder, kSignatureFieldName);
+
+    BSONObjBuilder signatureObjBuilder(subObjBuilder.subobjStart(kSignatureFieldName));
+    _clusterTime.getProof().appendAsBinData(signatureObjBuilder, kSignatureHashFieldName);
+    signatureObjBuilder.append(kSignatureKeyIdFieldName, _clusterTime.getKeyId());
+    signatureObjBuilder.doneFast();
+
     subObjBuilder.doneFast();
 }
 
diff --git a/src/mongo/rpc/metadata/logical_time_metadata.h b/src/mongo/rpc/metadata/logical_time_metadata.h
index 2ff792b58b1..2ec6b7f03b5 100644
--- a/src/mongo/rpc/metadata/logical_time_metadata.h
+++ b/src/mongo/rpc/metadata/logical_time_metadata.h
@@ -42,7 +42,10 @@ namespace rpc {
  * Format:
  * logicalTime: {
  *     clusterTime: <Timestamp>,
- *     signature: <SHA1 hash of clusterTime as BinData>
+ *     signature: {
+ *         hash: <SHA1 hash of clusterTime as BinData>,
+ *         keyId: <long long>
+ *     }
  * }
  */
 class LogicalTimeMetadata {
diff --git a/src/mongo/rpc/metadata/logical_time_metadata_test.cpp b/src/mongo/rpc/metadata/logical_time_metadata_test.cpp
index 61106f3d469..73e9ad64208 100644
--- a/src/mongo/rpc/metadata/logical_time_metadata_test.cpp
+++ b/src/mongo/rpc/metadata/logical_time_metadata_test.cpp
@@ -48,7 +48,9 @@ TEST(LogicalTimeMetadataTest, Roundtrip) {
     proof[19] = 6;
     proof[0] = 12;
 
-    SignedLogicalTime signedTs(LogicalTime(ts), proof);
+    long long keyId = 1;
+
+    SignedLogicalTime signedTs(LogicalTime(ts), proof, keyId);
 
     LogicalTimeMetadata origMetadata(signedTs);
     BSONObjBuilder builder;
@@ -62,15 +64,34 @@ TEST(LogicalTimeMetadataTest, Roundtrip) {
     const auto& parsedTs = parsedMetadata.getSignedTime();
     ASSERT_EQ(ts.asTimestamp(), parsedTs.getTime().asTimestamp());
     ASSERT_TRUE(SHA1Block(proof) == parsedTs.getProof());
+    ASSERT_TRUE(keyId == parsedTs.getKeyId());
 }
 
 TEST(LogicalTimeMetadataTest, MissingClusterTimeShouldFailToParse) {
     std::array<uint8_t, 20> proof;
     proof.fill(0);
 
+    long long keyId = 1;
+
+    BSONObjBuilder builder;
+    BSONObjBuilder subObjBuilder(builder.subobjStart("logicalTime"));
+    BSONObjBuilder signatureObjBuilder(subObjBuilder.subobjStart("signature"));
+    signatureObjBuilder.append("hash", BSONBinData(proof.data(), proof.size(), BinDataGeneral));
+    signatureObjBuilder.append("keyId", keyId);
+    signatureObjBuilder.doneFast();
+    subObjBuilder.doneFast();
+
+    auto serializedObj = builder.done();
+    auto status = LogicalTimeMetadata::readFromMetadata(serializedObj).getStatus();
+    ASSERT_EQ(ErrorCodes::NoSuchKey, status);
+}
+
+TEST(LogicalTimeMetadataTest, MissingSignatureShouldFailToParse) {
+    const auto ts = Timestamp(100, 200);
+
     BSONObjBuilder builder;
     BSONObjBuilder subObjBuilder(builder.subobjStart("logicalTime"));
-    subObjBuilder.append("signature", BSONBinData(proof.data(), proof.size(), BinDataGeneral));
+    ts.append(subObjBuilder.bb(), "clusterTime");
     subObjBuilder.doneFast();
 
     auto serializedObj = builder.done();
@@ -78,12 +99,36 @@ TEST(LogicalTimeMetadataTest, MissingClusterTimeShouldFailToParse) {
     ASSERT_EQ(ErrorCodes::NoSuchKey, status);
 }
 
-TEST(LogicalTimeMetadataTest, MissingProofShouldFailToParse) {
+TEST(LogicalTimeMetadataTest, MissingHashShouldFailToParse) {
     const auto ts = Timestamp(100, 200);
 
+    long long keyId = 1;
+
     BSONObjBuilder builder;
     BSONObjBuilder subObjBuilder(builder.subobjStart("logicalTime"));
     ts.append(subObjBuilder.bb(), "clusterTime");
+    BSONObjBuilder signatureObjBuilder(subObjBuilder.subobjStart("signature"));
+    signatureObjBuilder.append("keyId", keyId);
+    signatureObjBuilder.doneFast();
+    subObjBuilder.doneFast();
+
+    auto serializedObj = builder.done();
+    auto status = LogicalTimeMetadata::readFromMetadata(serializedObj).getStatus();
+    ASSERT_EQ(ErrorCodes::NoSuchKey, status);
+}
+
+TEST(LogicalTimeMetadataTest, MissingKeyIdShouldFailToParse) {
+    const auto ts = Timestamp(100, 200);
+
+    std::array<uint8_t, 20> proof;
+    proof.fill(0);
+
+    BSONObjBuilder builder;
+    BSONObjBuilder subObjBuilder(builder.subobjStart("logicalTime"));
+    ts.append(subObjBuilder.bb(), "clusterTime");
+    BSONObjBuilder signatureObjBuilder(subObjBuilder.subobjStart("signature"));
+    signatureObjBuilder.append("hash", BSONBinData(proof.data(), proof.size(), BinDataGeneral));
+    signatureObjBuilder.doneFast();
     subObjBuilder.doneFast();
 
     auto serializedObj = builder.done();
@@ -97,10 +142,15 @@ TEST(LogicalTimeMetadataTest, ProofWithWrongLengthShouldFailToParse) {
     std::array<uint8_t, 10> proof;
     proof.fill(0);
 
+    long long keyId = 1;
+
     BSONObjBuilder builder;
     BSONObjBuilder subObjBuilder(builder.subobjStart("logicalTime"));
     ts.append(subObjBuilder.bb(), "clusterTime");
-    subObjBuilder.append("signature", BSONBinData(proof.data(), proof.size(), BinDataGeneral));
+    BSONObjBuilder signatureObjBuilder(subObjBuilder.subobjStart("signature"));
+    signatureObjBuilder.append("hash", BSONBinData(proof.data(), proof.size(), BinDataGeneral));
+    signatureObjBuilder.append("keyId", keyId);
+    signatureObjBuilder.doneFast();
     subObjBuilder.doneFast();
 
     auto serializedObj = builder.done();