fix tests

Author: BewareMyPower

lib/ConsumerImpl.cc

@@ -608,7 +608,8 @@ void ConsumerImpl::messageReceived(const ClientConnectionPtr& cnx, const proto::
         return;
     }
 
-    if (metadata.has_num_messages_in_batch()) {
+    // When the decryption failed, the whole batch message will be treated as a single message.
+    if (metadata.has_num_messages_in_batch() && decryptResult == DECRYPTED) {
         BitSet::Data words(msg.ack_set_size());
         for (int i = 0; i < words.size(); i++) {
             words[i] = msg.ack_set(i);

tests/EncryptionTests.cc

@@ -20,70 +20,105 @@
 #include <pulsar/Client.h>
 #include <pulsar/ConsumerCryptoFailureAction.h>
 
+#include <stdexcept>
+
+#include "PulsarApi.pb.h"
+#include "lib/CompressionCodec.h"
+#include "lib/MessageCrypto.h"
+#include "lib/SharedBuffer.h"
+#include "tests/PulsarFriend.h"
+
 static std::string lookupUrl = "pulsar://localhost:6650";
 
 using namespace pulsar;
 
+static CryptoKeyReaderPtr getDefaultCryptoKeyReader() {
+    return std::make_shared<DefaultCryptoKeyReader>(TEST_CONF_DIR "/public-key.client-rsa.pem",
+                                                    TEST_CONF_DIR "/private-key.client-rsa.pem");
+}
+
+static std::vector<std::string> decryptValue(const Message& message) {
+    if (!message.getEncryptionContext().has_value()) {
+        return {message.getDataAsString()};
+    }
+    auto context = message.getEncryptionContext().value();
+    if (!context->isDecryptionFailed()) {
+        return {message.getDataAsString()};
+    }
+
+    MessageCrypto crypto{"test", false};
+    auto msgImpl = PulsarFriend::getMessageImplPtr(message);
+    SharedBuffer decryptedPayload;
+    // TODO: change the parameters to get context from EncryptionContext directly
+    if (!crypto.decrypt(msgImpl->metadata, msgImpl->payload, getDefaultCryptoKeyReader(), decryptedPayload)) {
+        throw std::runtime_error("Decryption failed");
+    }
+
+    SharedBuffer uncompressedPayload;
+    if (!CompressionCodecProvider::getCodec(context->compressionType())
+             .decode(decryptedPayload, context->uncompressedMessageSize(), uncompressedPayload)) {
+        throw std::runtime_error("Decompression failed");
+    }
+
+    std::vector<std::string> values;
+    if (auto batchSize = message.getEncryptionContext().value()->batchSize(); batchSize > 0) {
+        for (decltype(batchSize) i = 0; i < batchSize; i++) {
+            auto singleMetaSize = uncompressedPayload.readUnsignedInt();
+            proto::SingleMessageMetadata singleMeta;
+            singleMeta.ParseFromArray(uncompressedPayload.data(), singleMetaSize);
+            uncompressedPayload.consume(singleMetaSize);
+
+            auto payload = uncompressedPayload.slice(0, singleMeta.payload_size());
+            uncompressedPayload.consume(payload.readableBytes());
+            values.emplace_back(payload.data(), payload.readableBytes());
+        }
+    } else {
+        // non-batched message
+        values.emplace_back(uncompressedPayload.data(), uncompressedPayload.readableBytes());
+    }
+    return values;
+}
+
 TEST(EncryptionTests, testEncryptionContext) {
     Client client{lookupUrl};
     std::string topic = "test-encryption-context-" + std::to_string(time(nullptr));
 
     ProducerConfiguration producerConf;
     producerConf.setCompressionType(CompressionLZ4);
     producerConf.addEncryptionKey("client-rsa.pem");
-    producerConf.setCryptoKeyReader(std::make_shared<DefaultCryptoKeyReader>(
-        TEST_CONF_DIR "/public-key.client-rsa.pem", TEST_CONF_DIR "/private-key.client-rsa.pem"));
-    // TODO: enable batching
-    producerConf.setBatchingEnabled(false);
+    producerConf.setCryptoKeyReader(getDefaultCryptoKeyReader());
 
     Producer producer;
     ASSERT_EQ(ResultOk, client.createProducer(topic, producerConf, producer));
 
+    std::vector<std::string> sentValues;
+    auto send = [&producer, &sentValues](const std::string& value) {
+        Message msg = MessageBuilder().setContent(value).build();
+        producer.sendAsync(msg, nullptr);
+        sentValues.emplace_back(value);
+    };
+
     for (int i = 0; i < 5; i++) {
-        producer.sendAsync(MessageBuilder().setContent("msg-" + std::to_string(i)).build(), nullptr);
+        send("msg-" + std::to_string(i));
     }
     producer.flush();
-    producer.send(MessageBuilder().setContent("last-msg").build());
+    send("last-msg");
 
     ConsumerConfiguration consumerConf;
     consumerConf.setSubscriptionInitialPosition(InitialPositionEarliest);
     consumerConf.setCryptoFailureAction(ConsumerCryptoFailureAction::CONSUME);
     Consumer consumer;
     ASSERT_EQ(ResultOk, client.subscribe(topic, "sub", consumerConf, consumer));
 
-    std::vector<Message> messages;
-    for (int i = 0; i < 6; i++) {
+    std::vector<std::string> values;
+    for (int i = 0; i < 2; i++) {
         Message msg;
         ASSERT_EQ(ResultOk, consumer.receive(msg, 3000));
-        messages.emplace_back(msg);
-    }
-
-    // TODO: improve it
-    for (int i = 0; i < 6; i++) {
-        auto context = messages[i].getEncryptionContext();
-        if (context.has_value()) {
-            std::cout << i << " keys: ";
-            for (auto&& key : context.value()->keys()) {
-                std::cout << key.first << " => " << key.second.value;
-                if (!key.second.metadata.empty()) {
-                    std::cout << " (metadata: ";
-                    for (auto&& entry : key.second.metadata) {
-                        std::cout << entry.first << "=" << entry.second << " ";
-                    }
-                    std::cout << ")";
-                }
-                std::cout << " ";
-            }
-            std::cout << ", algorithm: " << context.value()->algorithm()
-                      << ", param: " << context.value()->param()
-                      << ", compressionType: " << context.value()->compressionType()
-                      << ", uncompressedMessageSize: " << context.value()->uncompressedMessageSize()
-                      << ", batchSize: " << context.value()->batchSize()
-                      << ", isDecryptionFailed: " << context.value()->isDecryptionFailed() << "\n";
-        } else {
-            std::cerr << i << " no encryption context\n";
+        for (auto&& value : decryptValue(msg)) {
+            values.emplace_back(value);
         }
     }
+    ASSERT_EQ(values, sentValues);
 
     client.close();
 }

tests/PulsarFriend.h

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