Initial commit

Author: BewareMyPower

include/pulsar/EncryptionContext.h

@@ -0,0 +1,119 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#pragma once
+
+#include <cstdint>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "CompressionType.h"
+#include "defines.h"
+
+namespace pulsar {
+
+namespace proto {
+class MessageMetadata;
+}
+
+class Message;
+
+struct PULSAR_PUBLIC EncryptionKey {
+    std::string key;
+    std::string value;
+    std::unordered_map<std::string, std::string> metadata;
+
+    explicit EncryptionKey() = default;
+
+    // Support in-place construction
+    EncryptionKey(const std::string& key, const std::string& value,
+                  const decltype(EncryptionKey::metadata)& metadata)
+        : key(key), value(value), metadata(metadata) {}
+};
+
+/**
+ * It contains encryption and compression information in it using which application can decrypt consumed
+ * message with encrypted-payload.
+ */
+class PULSAR_PUBLIC EncryptionContext {
+   public:
+    explicit EncryptionContext()
+        : compressionType_(CompressionNone),
+          uncompressedMessageSize_(0),
+          batchSize_(-1),
+          isDecryptionFailed_(false) {}
+    EncryptionContext(const EncryptionContext&) = default;
+    EncryptionContext(EncryptionContext&&) noexcept = default;
+    EncryptionContext(const proto::MessageMetadata& metadata, bool isDecryptionFailed);
+
+    using KeysType = std::vector<EncryptionKey>;
+
+    /**
+     * @return the map of encryption keys used for the message
+     */
+    const KeysType& keys() const noexcept { return keys_; }
+
+    /**
+     * @return the encryption parameter used for the message
+     */
+    const std::string& param() const noexcept { return param_; }
+
+    /**
+     * @return the encryption algorithm used for the message
+     */
+    const std::string& algorithm() const noexcept { return algorithm_; }
+
+    /**
+     * @return the compression type used for the message
+     */
+    CompressionType compressionType() const noexcept { return compressionType_; }
+
+    /**
+     * @return the uncompressed message size if the message is compressed, 0 otherwise
+     */
+    uint32_t uncompressedMessageSize() const noexcept { return uncompressedMessageSize_; }
+
+    /**
+     * @return the batch size if the message is part of a batch, -1 otherwise
+     */
+    int32_t batchSize() const noexcept { return batchSize_; }
+
+    /**
+     * When the `ConsumerConfiguration#getCryptoFailureAction` is set to `CONSUME`, the message will still be
+     * returned even if the decryption failed, in this case, the message payload is still not decrypted but
+     * users have no way to know that. This method is provided to let users know whether the decryption
+     * failed.
+     *
+     * @return whether the decryption failed
+     */
+    bool isDecryptionFailed() const noexcept { return isDecryptionFailed_; }
+
+   private:
+    KeysType keys_;
+    std::string param_;
+    std::string algorithm_;
+    CompressionType compressionType_;
+    uint32_t uncompressedMessageSize_;
+    int32_t batchSize_;
+    bool isDecryptionFailed_;
+
+    friend class ConsumerImpl;
+};
+
+}  // namespace pulsar

include/pulsar/Message.h

@@ -19,22 +19,18 @@
 #ifndef MESSAGE_HPP_
 #define MESSAGE_HPP_
 
+#include <pulsar/EncryptionContext.h>
 #include <pulsar/defines.h>
 
 #include <map>
 #include <memory>
+#include <optional>
 #include <string>
 
 #include "KeyValue.h"
 #include "MessageId.h"
 
 namespace pulsar {
-namespace proto {
-class CommandMessage;
-class BrokerEntryMetadata;
-class MessageMetadata;
-class SingleMessageMetadata;
-}  // namespace proto
 
 class SharedBuffer;
 class MessageBuilder;
@@ -202,19 +198,19 @@ class PULSAR_PUBLIC Message {
      */
     const std::string& getProducerName() const noexcept;
 
+    /**
+     * @return the optional encryption context that is present when the message is encrypted, the pointer is
+     * valid as the Message instance is alive
+     */
+    std::optional<const EncryptionContext*> getEncryptionContext() const;
+
     bool operator==(const Message& msg) const;
 
    protected:
     typedef std::shared_ptr<MessageImpl> MessageImplPtr;
     MessageImplPtr impl_;
 
-    Message(MessageImplPtr& impl);
-    Message(const MessageId& messageId, proto::BrokerEntryMetadata& brokerEntryMetadata,
-            proto::MessageMetadata& metadata, SharedBuffer& payload);
-    /// Used for Batch Messages
-    Message(const MessageId& messageId, proto::BrokerEntryMetadata& brokerEntryMetadata,
-            proto::MessageMetadata& metadata, SharedBuffer& payload,
-            proto::SingleMessageMetadata& singleMetadata, const std::shared_ptr<std::string>& topicName);
+    Message(const MessageImplPtr& impl);
     friend class PartitionedProducerImpl;
     friend class MultiTopicsConsumerImpl;
     friend class MessageBuilder;

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
@@ -926,12 +927,13 @@ Message Commands::deSerializeSingleMessageInBatch(Message& batchedMessage, int32
     const MessageId& m = batchedMessage.impl_->messageId;
     auto messageId = MessageIdBuilder::from(m).batchIndex(batchIndex).batchSize(batchSize).build();
     auto batchedMessageId = std::make_shared<BatchedMessageIdImpl>(*(messageId.impl_), acker);
-    Message singleMessage(MessageId{batchedMessageId}, batchedMessage.impl_->brokerEntryMetadata,
-                          batchedMessage.impl_->metadata, payload, metadata,
-                          batchedMessage.impl_->topicName_);
-    singleMessage.impl_->cnx_ = batchedMessage.impl_->cnx_;
 
-    return singleMessage;
+    auto msgImpl = std::make_shared<MessageImpl>(messageId, batchedMessage.impl_->brokerEntryMetadata,
+                                                 batchedMessage.impl_->metadata, payload, metadata,
+                                                 batchedMessage.impl_->topicName_, encryptionContext);
+    msgImpl->cnx_ = batchedMessage.impl_->cnx_;
+
+    return Message(msgImpl);
 }
 
 MessageIdImplPtr Commands::getMessageIdImpl(const MessageId& messageId) { return messageId.impl_; }

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);
 

lib/ConsumerImpl.cc

@@ -548,25 +548,27 @@ void ConsumerImpl::messageReceived(const ClientConnectionPtr& cnx, const proto::
                                    bool& isChecksumValid, proto::BrokerEntryMetadata& brokerEntryMetadata,
                                    proto::MessageMetadata& metadata, SharedBuffer& payload) {
     LOG_DEBUG(getName() << "Received Message -- Size: " << payload.readableBytes());
-
-    if (!decryptMessageIfNeeded(cnx, msg, metadata, payload)) {
-        // Message was discarded or not consumed due to decryption failure
-        return;
-    }
-
     if (!isChecksumValid) {
         // Message discarded for checksum error
         discardCorruptedMessage(cnx, msg.message_id(), CommandAck_ValidationError_ChecksumMismatch);
         return;
     }
 
-    auto redeliveryCount = msg.redelivery_count();
-    const bool isMessageUndecryptable =
-        metadata.encryption_keys_size() > 0 && !config_.getCryptoKeyReader().get() &&
-        config_.getCryptoFailureAction() == ConsumerCryptoFailureAction::CONSUME;
+    auto encryptionContext = metadata.encryption_keys_size() > 0
+                                 ? optional<EncryptionContext>{std::in_place, metadata, false}
+                                 : std::nullopt;
 
+    auto decryptResult = decryptMessageIfNeeded(cnx, encryptionContext, payload, msg.message_id());
+    if (decryptResult == FAILED) {
+        // Message was discarded due to decryption failure or not consumed due to decryption failure
+        return;
+    } else if (decryptResult == CONSUME_ENCRYPTED) {
+        encryptionContext->isDecryptionFailed_ = true;
+    }
+
+    auto redeliveryCount = msg.redelivery_count();
     const bool isChunkedMessage = metadata.num_chunks_from_msg() > 1;
-    if (!isMessageUndecryptable && !isChunkedMessage) {
+    if (decryptResult == DECRYPTED && !isChunkedMessage) {
         if (!uncompressMessageIfNeeded(cnx, msg.message_id(), metadata, payload, true)) {
             // Message was discarded on decompression error
             return;
@@ -586,9 +588,9 @@ void ConsumerImpl::messageReceived(const ClientConnectionPtr& cnx, const proto::
         }
     }
 
-    Message m(messageId, brokerEntryMetadata, metadata, payload);
+    Message m{std::make_shared<MessageImpl>(messageId, brokerEntryMetadata, metadata, payload, std::nullopt,
+                                            getTopicPtr(), std::move(encryptionContext))};
     m.impl_->cnx_ = cnx.get();
-    m.impl_->setTopicName(getTopicPtr());
     m.impl_->setRedeliveryCount(msg.redelivery_count());
 
     if (metadata.has_schema_version()) {
@@ -610,14 +612,16 @@ 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);
         }
         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());
@@ -742,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());
@@ -756,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());
@@ -812,50 +817,51 @@ uint32_t ConsumerImpl::receiveIndividualMessagesFromBatch(const ClientConnection
     return batchSize - skippedMessages;
 }
 
-bool ConsumerImpl::decryptMessageIfNeeded(const ClientConnectionPtr& cnx, const proto::CommandMessage& msg,
-                                          const proto::MessageMetadata& metadata, SharedBuffer& payload) {
-    if (!metadata.encryption_keys_size()) {
-        return true;
+auto ConsumerImpl::decryptMessageIfNeeded(const ClientConnectionPtr& cnx,
+                                          const optional<EncryptionContext>& context, SharedBuffer& payload,
+                                          const proto::MessageIdData& msgId) -> DecryptResult {
+    if (!context.has_value()) {
+        return DECRYPTED;
     }
 
     // If KeyReader is not configured throw exception based on config param
     if (!config_.isEncryptionEnabled()) {
         if (config_.getCryptoFailureAction() == ConsumerCryptoFailureAction::CONSUME) {
             LOG_WARN(getName() << "CryptoKeyReader is not implemented. Consuming encrypted message.");
-            return true;
+            return CONSUME_ENCRYPTED;
         } else if (config_.getCryptoFailureAction() == ConsumerCryptoFailureAction::DISCARD) {
             LOG_WARN(getName() << "Skipping decryption since CryptoKeyReader is not implemented and config "
                                   "is set to discard");
-            discardCorruptedMessage(cnx, msg.message_id(), CommandAck_ValidationError_DecryptionError);
+            discardCorruptedMessage(cnx, msgId, CommandAck_ValidationError_DecryptionError);
         } else {
             LOG_ERROR(getName() << "Message delivery failed since CryptoKeyReader is not implemented to "
                                    "consume encrypted message");
-            auto messageId = MessageIdBuilder::from(msg.message_id()).build();
+            auto messageId = MessageIdBuilder::from(msgId).build();
             unAckedMessageTrackerPtr_->add(messageId);
         }
-        return false;
+        return FAILED;
     }
 
     SharedBuffer decryptedPayload;
-    if (msgCrypto_->decrypt(metadata, payload, config_.getCryptoKeyReader(), decryptedPayload)) {
+    if (msgCrypto_->decrypt(*context, payload, config_.getCryptoKeyReader(), decryptedPayload)) {
         payload = decryptedPayload;
-        return true;
+        return DECRYPTED;
     }
 
     if (config_.getCryptoFailureAction() == ConsumerCryptoFailureAction::CONSUME) {
         // Note, batch message will fail to consume even if config is set to consume
         LOG_WARN(
             getName() << "Decryption failed. Consuming encrypted message since config is set to consume.");
-        return true;
+        return CONSUME_ENCRYPTED;
     } else if (config_.getCryptoFailureAction() == ConsumerCryptoFailureAction::DISCARD) {
         LOG_WARN(getName() << "Discarding message since decryption failed and config is set to discard");
-        discardCorruptedMessage(cnx, msg.message_id(), CommandAck_ValidationError_DecryptionError);
+        discardCorruptedMessage(cnx, msgId, CommandAck_ValidationError_DecryptionError);
     } else {
         LOG_ERROR(getName() << "Message delivery failed since unable to decrypt incoming message");
-        auto messageId = MessageIdBuilder::from(msg.message_id()).build();
+        auto messageId = MessageIdBuilder::from(msgId).build();
         unAckedMessageTrackerPtr_->add(messageId);
     }
-    return false;
+    return FAILED;
 }
 
 bool ConsumerImpl::uncompressMessageIfNeeded(const ClientConnectionPtr& cnx,