PIP-432: Add isEncrypted field to EncryptionContext (#324) (#325)

* PIP-432: Add isEncrypted field to EncryptionContext

* Add doc for EncryptionContext.IsEncrypted

(cherry picked from commit 1eae9d1)

Author: RobertIndie

src/Pulsar.Client/Common/DTO.fs

@@ -242,9 +242,11 @@ type EncryptionContext =
         CompressionType: CompressionType
         UncompressedMessageSize: int
         BatchSize: Nullable<int>
+        // Indicates whether the message payload remains encrypted (true) or has been successfully decrypted (false)
+        IsEncrypted: bool
     }
     with
-        static member internal FromMetadata(metadata: Metadata) =
+        static member internal FromMetadata(metadata: Metadata, isEncrypted: bool) =
             if metadata.EncryptionKeys.Length > 0 then
                 {
                     Keys = metadata.EncryptionKeys
@@ -253,6 +255,7 @@ type EncryptionContext =
                     CompressionType = metadata.CompressionType
                     UncompressedMessageSize = metadata.UncompressedMessageSize
                     BatchSize = if metadata.HasNumMessagesInBatch then Nullable(metadata.NumMessages) else Nullable()
+                    IsEncrypted = isEncrypted
                 } |> Some
             else
                 None

src/Pulsar.Client/Internal/ConsumerImpl.fs

@@ -661,7 +661,7 @@ type internal ConsumerImpl<'T> (consumerConfig: ConsumerConfiguration<'T>, clien
                 trackMessage rawMessage.MessageId
             None
 
-    let handleSingleMessagePayload (rawMessage: RawMessage) msgId payload hasWaitingChannel hasWaitingBatchChannel schemaDecodeFunction =
+    let handleSingleMessagePayload (rawMessage: RawMessage) msgId payload hasWaitingChannel hasWaitingBatchChannel isMessageUndecryptable schemaDecodeFunction =
         if duringSeek.IsSome || (isSameEntry(rawMessage.MessageId) && isPriorEntryIndex(rawMessage.MessageId.EntryId)) then
             // We need to discard entries that were prior to startMessageId
             Log.Logger.LogInformation("{0} Ignoring message from before the startMessageId: {1}", prefix, startMessageId)
@@ -677,7 +677,7 @@ type internal ConsumerImpl<'T> (consumerConfig: ConsumerConfiguration<'T>, clien
                             %msgKey,
                             rawMessage.IsKeyBase64Encoded,
                             rawMessage.Properties,
-                            EncryptionContext.FromMetadata rawMessage.Metadata,
+                            EncryptionContext.FromMetadata(rawMessage.Metadata, isEncrypted = isMessageUndecryptable),
                             getSchemaVersionBytes rawMessage.Metadata.SchemaVersion,
                             rawMessage.Metadata.SequenceId,
                             rawMessage.Metadata.OrderingKey,
@@ -715,15 +715,15 @@ type internal ConsumerImpl<'T> (consumerConfig: ConsumerConfiguration<'T>, clien
                     if isChunkedMessage then
                         match processMessageChunk rawMessage msgId with
                         | Some (chunkedPayload, msgIdWithChunk) ->
-                            handleSingleMessagePayload rawMessage msgIdWithChunk chunkedPayload hasWaitingChannel hasWaitingBatchChannel schemaDecodeFunction
+                            handleSingleMessagePayload rawMessage msgIdWithChunk chunkedPayload hasWaitingChannel hasWaitingBatchChannel isMessageUndecryptable schemaDecodeFunction
                         | None ->
                             ()
                     else
-                        handleSingleMessagePayload rawMessage msgId rawMessage.Payload hasWaitingChannel hasWaitingBatchChannel schemaDecodeFunction
+                        handleSingleMessagePayload rawMessage msgId rawMessage.Payload hasWaitingChannel hasWaitingBatchChannel isMessageUndecryptable schemaDecodeFunction
                 elif rawMessage.Metadata.NumMessages > 0 then
                     // handle batch message enqueuing; uncompressed payload has all messages in batch
                     try
-                        this.ReceiveIndividualMessagesFromBatch rawMessage schemaDecodeFunction
+                        this.ReceiveIndividualMessagesFromBatch rawMessage schemaDecodeFunction isMessageUndecryptable
                     with ex ->
                         Log.Logger.LogError(ex, "{0} Batch reading exception {1}", prefix, msgId)
                         raise <| BatchDeserializationException "Batch reading exception"
@@ -1315,8 +1315,8 @@ type internal ConsumerImpl<'T> (consumerConfig: ConsumerConfiguration<'T>, clien
     do startStatTimer()
     do startChunkTimer()
 
-    abstract member ReceiveIndividualMessagesFromBatch: RawMessage -> (byte [] -> 'T) -> unit
-    default this.ReceiveIndividualMessagesFromBatch (rawMessage: RawMessage) schemaDecodeFunction =
+    abstract member ReceiveIndividualMessagesFromBatch: RawMessage -> (byte [] -> 'T) -> bool -> unit
+    default this.ReceiveIndividualMessagesFromBatch (rawMessage: RawMessage) schemaDecodeFunction isMessageUndecryptable =
         let batchSize = rawMessage.Metadata.NumMessages
         let acker = BatchMessageAcker(batchSize)
         let mutable skippedMessages = 0
@@ -1366,7 +1366,7 @@ type internal ConsumerImpl<'T> (consumerConfig: ConsumerConfiguration<'T>, clien
                                 %msgKey,
                                 singleMessageMetadata.PartitionKeyB64Encoded,
                                 properties,
-                                EncryptionContext.FromMetadata rawMessage.Metadata,
+                                EncryptionContext.FromMetadata(rawMessage.Metadata, isEncrypted = isMessageUndecryptable),
                                 getSchemaVersionBytes rawMessage.Metadata.SchemaVersion,
                                 %(int64 singleMessageMetadata.SequenceId),
                                 singleMessageMetadata.OrderingKey,
@@ -1704,7 +1704,7 @@ and internal ZeroQueueConsumerImpl<'T> (consumerConfig: ConsumerConfiguration<'T
         if this.Waiters.Count > 0 then
             this.SendFlowPermits this.Waiters.Count
 
-    override this.ReceiveIndividualMessagesFromBatch (_: RawMessage) _ =
+    override this.ReceiveIndividualMessagesFromBatch (_: RawMessage) _ _ =
         Log.Logger.LogError("{0} Closing consumer due to unsupported received batch-message with zero receiver queue size", prefix)
         let _ = postAndAsyncReply this.Mb ConsumerMessage.Close
         let exn =

tests/IntegrationTests/MessageCrypto.fs

@@ -95,6 +95,32 @@ type Consumer2KeyReader() =
             { Key = Encoding.UTF8.GetBytes(privateKeysConsumer2.Item keyName); Metadata = null }
 
 
+let consumeMessagesAndCheckEncryption (consumer: IConsumer<byte[]>) number consumerName =
+    task {
+        for i in 1..number do
+            let! message = consumer.ReceiveAsync()
+            let received = Encoding.UTF8.GetString(message.Data)
+            Log.Debug("{0} received {1}", consumerName, received)
+            
+            // Check EncryptionContext
+            match message.EncryptionContext with
+            | Some context ->
+                Log.Debug("{0} message {1} EncryptionContext.IsEncrypted = {2}", consumerName, i, context.IsEncrypted)
+                if context.IsEncrypted then
+                    failwith $"Message {i} should not be encrypted, but IsEncrypted = true"
+            | None ->
+                failwith $"Message {i} should have EncryptionContext but it is None"
+            
+            do! consumer.AcknowledgeAsync(message.MessageId)
+            Log.Debug("{0} acknowledged {1}", consumerName, received)
+            let expected = "Message #" + string i
+            if received.StartsWith(expected) |> not then
+                failwith $"Incorrect message expected {expected} received {received} consumer {consumerName}"
+        
+        Log.Debug("{0} consumed {1} messages, all EncryptionContext.IsEncrypted = false", consumerName, number)
+    }
+
+
 [<Tests>]
 let tests =
     testList "MessageCrypto" [
@@ -129,7 +155,7 @@ let tests =
             let consumerTask =
               Task.Run(fun () ->
                   task {
-                      do! consumeMessages consumer numberOfMessages consumerName
+                      do! consumeMessagesAndCheckEncryption consumer numberOfMessages consumerName
                   } :> Task)
 
             do! Task.WhenAll(producerTask, consumerTask) 
@@ -169,7 +195,7 @@ let tests =
             let consumer1Task =
               Task.Run(fun () ->
                   task {
-                      do! consumeMessages consumer1 numberOfMessages consumerName
+                      do! consumeMessagesAndCheckEncryption consumer1 numberOfMessages consumerName
                   } :> Task)
 
             do! Task.WhenAll(producerTask, consumer1Task) 
@@ -196,7 +222,7 @@ let tests =
             let consumer2Task =
               Task.Run(fun () ->
                   task {
-                      do! consumeMessages consumer2 numberOfMessages consumerName
+                      do! consumeMessagesAndCheckEncryption consumer2 numberOfMessages consumerName
                   } :> Task)
 
             do! Task.WhenAll(producerTask2, consumer2Task) 
@@ -242,7 +268,7 @@ let tests =
             Expect.isNonEmpty context.Param ""
             Expect.isNonEmpty context.Keys ""
             Expect.equal context.CompressionType compressionType ""
-
+            Expect.equal context.IsEncrypted true "Message should remain encrypted when decryption fails"
             do! Task.Delay 100
             Log.Debug("Ended Encryption send message and consume on fail")
         }