(WIP) EncryptionContext

TODO: create context, refactor message

Author: BewareMyPower

include/pulsar/EncryptionContext.h

@@ -0,0 +1,91 @@
+/**
+ * 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 "CompressionType.h"
+#include "defines.h"
+
+namespace pulsar {
+
+namespace proto {
+class MessageMetadata;
+}
+
+class Message;
+
+struct PULSAR_PUBLIC EncryptionKey {
+    std::string keyValue;
+    std::unordered_map<std::string, std::string> 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(const proto::MessageMetadata& metadata, bool isDecryptionFailed);
+
+    using KeysType = std::unordered_map<std::string, EncryptionKey>;
+    const KeysType& keys() const noexcept { return keys_; }
+    const std::string& param() const noexcept { return param_; }
+    const std::string& algorithm() const noexcept { return algorithm_; }
+    CompressionType compressionType() const noexcept { return compressionType_; }
+    uint32_t uncompressedMessageSize() const noexcept { return uncompressedMessageSize_; }
+    // Returns 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_; }
+
+    static const EncryptionContext& empty() {
+        static EncryptionContext context;
+        return context;
+    }
+
+   private:
+    const KeysType keys_;
+    const std::string param_;
+    const std::string algorithm_;
+    const CompressionType compressionType_;
+    const uint32_t uncompressedMessageSize_;
+    const int32_t batchSize_;
+    const bool isDecryptionFailed_;
+
+    friend class MessageImpl;
+};
+
+}  // namespace pulsar

include/pulsar/Message.h

@@ -19,10 +19,12 @@
 #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"
@@ -202,6 +204,12 @@ 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:

lib/ConsumerImpl.cc

@@ -548,25 +548,21 @@ 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 decryptResult = decryptMessageIfNeeded(cnx, msg, metadata, payload);
+    if (decryptResult == FAILED) {
+        // Message was discarded due to decryption failure or not consumed due to decryption failure
+        return;
+    }
 
+    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,6 +582,7 @@ void ConsumerImpl::messageReceived(const ClientConnectionPtr& cnx, const proto::
         }
     }
 
+    // TODO: create encryption context
     Message m(messageId, brokerEntryMetadata, metadata, payload);
     m.impl_->cnx_ = cnx.get();
     m.impl_->setTopicName(getTopicPtr());
@@ -812,17 +809,18 @@ 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 proto::CommandMessage& msg,
+                                          const proto::MessageMetadata& metadata, SharedBuffer& payload)
+    -> DecryptResult {
+    if (metadata.encryption_keys_size() == 0) {
+        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");
@@ -833,20 +831,20 @@ bool ConsumerImpl::decryptMessageIfNeeded(const ClientConnectionPtr& cnx, const
             auto messageId = MessageIdBuilder::from(msg.message_id()).build();
             unAckedMessageTrackerPtr_->add(messageId);
         }
-        return false;
+        return FAILED;
     }
 
     SharedBuffer decryptedPayload;
     if (msgCrypto_->decrypt(metadata, 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);
@@ -855,7 +853,7 @@ bool ConsumerImpl::decryptMessageIfNeeded(const ClientConnectionPtr& cnx, const
         auto messageId = MessageIdBuilder::from(msg.message_id()).build();
         unAckedMessageTrackerPtr_->add(messageId);
     }
-    return false;
+    return FAILED;
 }
 
 bool ConsumerImpl::uncompressMessageIfNeeded(const ClientConnectionPtr& cnx,

lib/ConsumerImpl.h

@@ -195,8 +195,14 @@ class ConsumerImpl : public ConsumerImplBase {
     bool isPriorEntryIndex(int64_t idx);
     void brokerConsumerStatsListener(Result, BrokerConsumerStatsImpl, const BrokerConsumerStatsCallback&);
 
-    bool decryptMessageIfNeeded(const ClientConnectionPtr& cnx, const proto::CommandMessage& msg,
-                                const proto::MessageMetadata& metadata, SharedBuffer& payload);
+    enum DecryptResult
+    {
+        DECRYPTED,
+        CONSUME_ENCRYPTED,
+        FAILED
+    };
+    DecryptResult decryptMessageIfNeeded(const ClientConnectionPtr& cnx, const proto::CommandMessage& msg,
+                                         const proto::MessageMetadata& metadata, SharedBuffer& payload);
 
     // TODO - Convert these functions to lambda when we move to C++11
     Result receiveHelper(Message& msg);

lib/EncryptionContext.cc

@@ -0,0 +1,49 @@
+/**
+ * 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.
+ */
+#include <pulsar/EncryptionContext.h>
+
+#include <unordered_map>
+
+#include "PulsarApi.pb.h"
+
+namespace pulsar {
+
+static EncryptionContext::KeysType encryptedKeysFromMetadata(const proto::MessageMetadata& msgMetadata) {
+    EncryptionContext::KeysType keys;
+    for (auto&& key : msgMetadata.encryption_keys()) {
+        decltype(EncryptionKey::metadata) keyMetadata;
+        for (auto&& entry : key.metadata()) {
+            keyMetadata[entry.key()] = entry.value();
+        }
+        keys[key.key()] = EncryptionKey{key.value(), keyMetadata};
+    }
+    return keys;
+}
+
+EncryptionContext::EncryptionContext(const proto::MessageMetadata& msgMetadata, bool isDecryptionFailed)
+
+    : keys_(encryptedKeysFromMetadata(msgMetadata)),
+      param_(msgMetadata.encryption_param()),
+      algorithm_(msgMetadata.encryption_algo()),
+      compressionType_(static_cast<CompressionType>(msgMetadata.compression())),
+      uncompressedMessageSize_(msgMetadata.uncompressed_size()),
+      batchSize_(msgMetadata.has_num_messages_in_batch() ? msgMetadata.num_messages_in_batch() : -1),
+      isDecryptionFailed_(isDecryptionFailed) {}
+
+}  // namespace pulsar

lib/Message.cc

@@ -16,6 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+#include <pulsar/EncryptionContext.h>
 #include <pulsar/Message.h>
 #include <pulsar/MessageBuilder.h>
 #include <pulsar/MessageIdBuilder.h>
@@ -220,6 +221,13 @@ const std::string& Message::getProducerName() const noexcept {
     return impl_->metadata.producer_name();
 }
 
+std::optional<const EncryptionContext*> Message::getEncryptionContext() const {
+    if (!impl_ || !impl_->encryptionContext_.has_value()) {
+        return std::nullopt;
+    }
+    return {&(*impl_->encryptionContext_)};
+}
+
 bool Message::operator==(const Message& msg) const { return getMessageId() == msg.getMessageId(); }
 
 KeyValue Message::getKeyValueData() const { return KeyValue(impl_->keyValuePtr); }

lib/MessageImpl.cc

@@ -17,9 +17,18 @@
  * under the License.
  */
 #include "MessageImpl.h"
+#include <utility>
 
 namespace pulsar {
 
+MessageImpl::MessageImpl(const MessageId& messageId, const proto::MessageMetadata& metadata,
+                         const proto::BrokerEntryMetadata& brokerEntryMetadata, const SharedBuffer& payload,
+                         const optional<proto::SingleMessageMetadata>& singleMetadata,
+                         const std::shared_ptr<std::string>& topicName, bool undecryptedPayload)
+    : encryptionContext_(std::in_place, metadata, undecryptedPayload) {
+    // TODO:
+}
+
 const Message::StringMap& MessageImpl::properties() {
     if (properties_.size() == 0) {
         for (int i = 0; i < metadata.properties_size(); i++) {

lib/MessageImpl.h

@@ -19,14 +19,19 @@
 #ifndef LIB_MESSAGEIMPL_H_
 #define LIB_MESSAGEIMPL_H_
 
+#include <pulsar/EncryptionContext.h>
 #include <pulsar/Message.h>
 #include <pulsar/MessageId.h>
 
+#include <memory>
+#include <optional>
+
 #include "KeyValueImpl.h"
 #include "PulsarApi.pb.h"
 #include "SharedBuffer.h"
 
-using namespace pulsar;
+using std::optional;
+
 namespace pulsar {
 
 class PulsarWrapper;
@@ -35,19 +40,14 @@ class BatchMessageContainer;
 
 class MessageImpl {
    public:
-    const Message::StringMap& properties();
+    // TODO: remove this default constructor
+    explicit MessageImpl() : encryptionContext_(std::nullopt) {}
+    MessageImpl(const MessageId& messageId, const proto::MessageMetadata& metadata,
+                const proto::BrokerEntryMetadata& brokerEntryMetadata, const SharedBuffer& payload,
+                const optional<proto::SingleMessageMetadata>& singleMetadata,
+                const std::shared_ptr<std::string>& topicName, bool undecryptedPayload);
 
-    proto::BrokerEntryMetadata brokerEntryMetadata;
-    proto::MessageMetadata metadata;
-    SharedBuffer payload;
-    std::shared_ptr<KeyValueImpl> keyValuePtr;
-    MessageId messageId;
-    ClientConnection* cnx_;
-    std::shared_ptr<std::string> topicName_;
-    int redeliveryCount_;
-    bool hasSchemaVersion_;
-    const std::string* schemaVersion_;
-    std::weak_ptr<class ConsumerImpl> consumerPtr_;
+    const Message::StringMap& properties();
 
     const std::string& getPartitionKey() const;
     bool hasPartitionKey() const;
@@ -81,6 +81,19 @@ class MessageImpl {
     friend class PulsarWrapper;
     friend class MessageBuilder;
 
+    proto::BrokerEntryMetadata brokerEntryMetadata;
+    proto::MessageMetadata metadata;
+    SharedBuffer payload;
+    std::shared_ptr<KeyValueImpl> keyValuePtr;
+    MessageId messageId;
+    ClientConnection* cnx_;
+    std::shared_ptr<std::string> topicName_;
+    int redeliveryCount_;
+    bool hasSchemaVersion_;
+    const std::string* schemaVersion_;
+    std::weak_ptr<class ConsumerImpl> consumerPtr_;
+    const optional<EncryptionContext> encryptionContext_;
+
    private:
     void setReplicationClusters(const std::vector<std::string>& clusters);
     void setProperty(const std::string& name, const std::string& value);