Fix encryption context not set for batched messages

Author: BewareMyPower

lib/Commands.cc

@@ -930,6 +930,7 @@ Message Commands::deSerializeSingleMessageInBatch(Message& batchedMessage, int32
                           batchedMessage.impl_->metadata, payload, metadata,
                           batchedMessage.impl_->topicName_);
     singleMessage.impl_->cnx_ = batchedMessage.impl_->cnx_;
+    singleMessage.impl_->encryptionContext_ = batchedMessage.impl_->encryptionContext_;
 
     return singleMessage;
 }

tests/BasicEndToEndTest.cc

@@ -1465,6 +1465,10 @@ TEST(BasicEndToEndTest, testRSAEncryption) {
             expected << msgContent << msgNum;
             ASSERT_EQ(expected.str(), msgReceived.getDataAsString());
             ASSERT_EQ(ResultOk, consumer.acknowledge(msgReceived));
+            auto context = msgReceived.getEncryptionContext();
+            ASSERT_TRUE(context.has_value());
+            ASSERT_EQ(context.value()->keys().size(), 1);
+            ASSERT_EQ(context.value()->keys()[0].key, "client-rsa.pem");
         }
 
         ASSERT_EQ(ResultOk, consumer.unsubscribe());

tests/EncryptionTest.cc

@@ -27,7 +27,6 @@
 #include "lib/CompressionCodec.h"
 #include "lib/MessageCrypto.h"
 #include "lib/SharedBuffer.h"
-#include "tests/PulsarFriend.h"
 
 static std::string lookupUrl = "pulsar://localhost:6650";
 
@@ -122,6 +121,11 @@ static void testDecryption(Client& client, const std::string& topic, bool decryp
     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());
+        } else {
+            ASSERT_FALSE(msg.getEncryptionContext().has_value());
+        }
         for (auto&& value : decryptValue(static_cast<const char*>(msg.getData()), msg.getLength(),
                                          msg.getEncryptionContext())) {
             values.emplace_back(value);

tests/PulsarFriend.h

@@ -217,8 +217,6 @@ class PulsarFriend {
         return waitUntil(std::chrono::seconds(3),
                          [producerImpl] { return !producerImpl->getCnx().expired(); });
     }
-
-    static auto getMessageImplPtr(const Message& message) { return message.impl_; }
 };
 }  // namespace pulsar