Fix encryption context not set when decryption succeeds

BewareMyPower · BewareMyPower · commit e53f8c626b53 · 2025-12-04T20:41:28.000+08:00
diff --git a/lib/Commands.cc b/lib/Commands.cc
@@ -906,7 +906,8 @@ uint64_t Commands::serializeSingleMessagesToBatchPayload(SharedBuffer& batchPayl
 }
 
 Message Commands::deSerializeSingleMessageInBatch(Message& batchedMessage, int32_t batchIndex,
-                                                  int32_t batchSize, const BatchMessageAckerPtr& acker) {
+                                                  int32_t batchSize, const BatchMessageAckerPtr& acker,
+                                                  const optional<EncryptionContext>& encryptionContext) {
     SharedBuffer& uncompressedPayload = batchedMessage.impl_->payload;
 
     // Format of batch message
@@ -927,10 +928,9 @@ Message Commands::deSerializeSingleMessageInBatch(Message& batchedMessage, int32
     auto messageId = MessageIdBuilder::from(m).batchIndex(batchIndex).batchSize(batchSize).build();
     auto batchedMessageId = std::make_shared<BatchedMessageIdImpl>(*(messageId.impl_), acker);
 
-    // TODO: fix the encryption context is not set
     auto msgImpl = std::make_shared<MessageImpl>(messageId, batchedMessage.impl_->brokerEntryMetadata,
                                                  batchedMessage.impl_->metadata, payload, metadata,
-                                                 batchedMessage.impl_->topicName_, std::nullopt);
+                                                 batchedMessage.impl_->topicName_, encryptionContext);
     msgImpl->cnx_ = batchedMessage.impl_->cnx_;
 
     return Message(msgImpl);
diff --git a/lib/Commands.h b/lib/Commands.h
@@ -155,7 +155,8 @@ class Commands {
                                                           const std::vector<Message>& messages);
 
     static Message deSerializeSingleMessageInBatch(Message& batchedMessage, int32_t batchIndex,
-                                                   int32_t batchSize, const BatchMessageAckerPtr& acker);
+                                                   int32_t batchSize, const BatchMessageAckerPtr& acker,
+                                                   const optional<EncryptionContext>& encryptionContext);
 
     static MessageIdImplPtr getMessageIdImpl(const MessageId& messageId);
 
diff --git a/lib/ConsumerImpl.cc b/lib/ConsumerImpl.cc
@@ -620,7 +620,8 @@ void ConsumerImpl::messageReceived(const ClientConnectionPtr& cnx, const proto::
         }
         BitSet ackSet{std::move(words)};
         Lock lock(mutex_);
-        numOfMessageReceived = receiveIndividualMessagesFromBatch(cnx, m, ackSet, msg.redelivery_count());
+        numOfMessageReceived =
+            receiveIndividualMessagesFromBatch(cnx, m, ackSet, msg.redelivery_count(), encryptionContext);
     } else {
         // try convert key value data.
         m.impl_->convertPayloadToKeyValue(config_.getSchema());
@@ -745,9 +746,9 @@ void ConsumerImpl::notifyPendingReceivedCallback(Result result, Message& msg,
 }
 
 // Zero Queue size is not supported with Batch Messages
-uint32_t ConsumerImpl::receiveIndividualMessagesFromBatch(const ClientConnectionPtr& cnx,
-                                                          Message& batchedMessage, const BitSet& ackSet,
-                                                          int redeliveryCount) {
+uint32_t ConsumerImpl::receiveIndividualMessagesFromBatch(
+    const ClientConnectionPtr& cnx, Message& batchedMessage, const BitSet& ackSet, int redeliveryCount,
+    const optional<EncryptionContext>& encryptionContext) {
     auto batchSize = batchedMessage.impl_->metadata.num_messages_in_batch();
     LOG_DEBUG("Received Batch messages of size - " << batchSize
                                                    << " -- msgId: " << batchedMessage.getMessageId());
@@ -759,7 +760,8 @@ uint32_t ConsumerImpl::receiveIndividualMessagesFromBatch(const ClientConnection
     std::vector<Message> possibleToDeadLetter;
     for (int i = 0; i < batchSize; i++) {
         // This is a cheap copy since message contains only one shared pointer (impl_)
-        Message msg = Commands::deSerializeSingleMessageInBatch(batchedMessage, i, batchSize, acker);
+        Message msg =
+            Commands::deSerializeSingleMessageInBatch(batchedMessage, i, batchSize, acker, encryptionContext);
         msg.impl_->setRedeliveryCount(redeliveryCount);
         msg.impl_->setTopicName(batchedMessage.impl_->topicName_);
         msg.impl_->convertPayloadToKeyValue(config_.getSchema());
diff --git a/lib/ConsumerImpl.h b/lib/ConsumerImpl.h
@@ -191,7 +191,8 @@ class ConsumerImpl : public ConsumerImplBase {
     void increaseAvailablePermits(const Message& msg);
     void drainIncomingMessageQueue(size_t count);
     uint32_t receiveIndividualMessagesFromBatch(const ClientConnectionPtr& cnx, Message& batchedMessage,
-                                                const BitSet& ackSet, int redeliveryCount);
+                                                const BitSet& ackSet, int redeliveryCount,
+                                                const optional<EncryptionContext>& encryptionContext);
     bool isPriorBatchIndex(int32_t idx);
     bool isPriorEntryIndex(int64_t idx);
     void brokerConsumerStatsListener(Result, BrokerConsumerStatsImpl, const BrokerConsumerStatsCallback&);
diff --git a/lib/MessageBatch.cc b/lib/MessageBatch.cc
@@ -49,7 +49,8 @@ MessageBatch& MessageBatch::parseFrom(const SharedBuffer& payload, uint32_t batc
 
     auto acker = BatchMessageAckerImpl::create(batchSize);
     for (int i = 0; i < batchSize; ++i) {
-        batch_.push_back(Commands::deSerializeSingleMessageInBatch(batchMessage_, i, batchSize, acker));
+        batch_.push_back(
+            Commands::deSerializeSingleMessageInBatch(batchMessage_, i, batchSize, acker, std::nullopt));
     }
     return *this;
 }
diff --git a/lib/MessageImpl.cc b/lib/MessageImpl.cc
@@ -28,7 +28,7 @@ MessageImpl::MessageImpl(const MessageId& messageId, const proto::BrokerEntryMet
                          const proto::MessageMetadata& metadata, const SharedBuffer& payload,
                          const optional<proto::SingleMessageMetadata>& singleMetadata,
                          const std::shared_ptr<std::string>& topicName,
-                         optional<EncryptionContext>&& encryptionContext)
+                         optional<EncryptionContext> encryptionContext)
     : messageId(messageId),
       brokerEntryMetadata(brokerEntryMetadata),
       metadata(metadata),
diff --git a/lib/MessageImpl.h b/lib/MessageImpl.h
@@ -44,8 +44,7 @@ class MessageImpl {
     MessageImpl(const MessageId& messageId, const proto::BrokerEntryMetadata& brokerEntryMetadata,
                 const proto::MessageMetadata& metadata, const SharedBuffer& payload,
                 const optional<proto::SingleMessageMetadata>& singleMetadata,
-                const std::shared_ptr<std::string>& topicName,
-                optional<EncryptionContext>&& encryptionContext);
+                const std::shared_ptr<std::string>& topicName, optional<EncryptionContext> encryptionContext);
 
     const Message::StringMap& properties();
 
diff --git a/tests/EncryptionTests.cc b/tests/EncryptionTests.cc
@@ -80,10 +80,8 @@ static std::vector<std::string> decryptValue(const Message& message) {
     return values;
 }
 
-TEST(EncryptionTests, testEncryptionContext) {
-    Client client{lookupUrl};
-    std::string topic = "test-encryption-context-" + std::to_string(time(nullptr));
-
+static void testDecryption(Client& client, const std::string& topic, bool decryptionSucceed,
+                           int numMessageReceived) {
     ProducerConfiguration producerConf;
     producerConf.setCompressionType(CompressionLZ4);
     producerConf.addEncryptionKey("client-rsa.pem");
@@ -104,22 +102,49 @@ TEST(EncryptionTests, testEncryptionContext) {
     }
     producer.flush();
     send("last-msg");
+    producer.flush();
+
+    ASSERT_EQ(ResultOk, client.createProducer(topic, producer));
+    send("unencrypted-msg");
+    producer.flush();
+    producer.close();
 
     ConsumerConfiguration consumerConf;
     consumerConf.setSubscriptionInitialPosition(InitialPositionEarliest);
-    consumerConf.setCryptoFailureAction(ConsumerCryptoFailureAction::CONSUME);
+    if (decryptionSucceed) {
+        consumerConf.setCryptoKeyReader(getDefaultCryptoKeyReader());
+    } else {
+        consumerConf.setCryptoFailureAction(ConsumerCryptoFailureAction::CONSUME);
+    }
     Consumer consumer;
     ASSERT_EQ(ResultOk, client.subscribe(topic, "sub", consumerConf, consumer));
 
     std::vector<std::string> values;
-    for (int i = 0; i < 2; i++) {
+    for (int i = 0; i < numMessageReceived; i++) {
         Message msg;
         ASSERT_EQ(ResultOk, consumer.receive(msg, 3000));
+        if (i < numMessageReceived - 1) {
+            ASSERT_TRUE(msg.getEncryptionContext().has_value());
+        }
         for (auto&& value : decryptValue(msg)) {
             values.emplace_back(value);
         }
     }
     ASSERT_EQ(values, sentValues);
+    consumer.close();
+}
 
+TEST(EncryptionTests, testDecryptionSuccess) {
+    Client client{lookupUrl};
+    std::string topic = "test-decryption-success-" + std::to_string(time(nullptr));
+    testDecryption(client, topic, true, 7);
+    client.close();
+}
+
+TEST(EncryptionTests, testDecryptionFailure) {
+    Client client{lookupUrl};
+    std::string topic = "test-decryption-failure-" + std::to_string(time(nullptr));
+    // The 1st batch that has 5 messages cannot be decrypted, so they can be received only once
+    testDecryption(client, topic, false, 3);
     client.close();
 }

Original file line number	Diff line number	Diff line change
`@@ -49,7 +49,8 @@ MessageBatch& MessageBatch::parseFrom(const SharedBuffer& payload, uint32_t batc`
`49`	`49`
`50`	`50`	`auto acker = BatchMessageAckerImpl::create(batchSize);`
`51`	`51`	`for (int i = 0; i < batchSize; ++i) {`
`52`		`- batch_.push_back(Commands::deSerializeSingleMessageInBatch(batchMessage_, i, batchSize, acker));`
	`52`	`+ batch_.push_back(`
	`53`	`+ Commands::deSerializeSingleMessageInBatch(batchMessage_, i, batchSize, acker, std::nullopt));`
`53`	`54`	`}`
`54`	`55`	`return *this;`
`55`	`56`	`}`