Skip to content

Commit

Permalink
[Client encryption]: Adds JsonNode support on deserialization path (#…
Browse files Browse the repository at this point in the history 
…4787)

# Pull Request Template

## Description

- Added support for System.Text.JsonNode DOM on deserializer and
decryption path
- Drops custom byte[] converter as JsonNode can take Memory<byte>
directly

To be processed after #4780 

## Type of change

Please delete options that are not relevant.

- [] New feature (non-breaking change which adds functionality)
- [] This change requires a documentation update

## Closing issues

Contributes to #4678

---------

Co-authored-by: Juraj Blazek <jublazek@microsoft.com>
Co-authored-by: juraj-blazek <53177060+juraj-blazek@users.noreply.github.com>
Co-authored-by: Santosh Kulkarni <66682828+kr-santosh@users.noreply.github.com>
Co-authored-by: Kiran Kumar Kolli <kirankk@microsoft.com>
  • Loading branch information
@JanHyka @juraj-blazek
@kr-santosh @kirankumarkolli
5 people authored Oct 16, 2024
1 parent 502a5a9 commit f88605c
Show file tree
Hide file tree
Showing 11 changed files with 596 additions and 39 deletions.
111 changes: 110 additions & 1 deletion Microsoft.Azure.Cosmos.Encryption.Custom/src/EncryptionProcessor.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,10 @@ namespace Microsoft.Azure.Cosmos.Encryption.Custom
using System.IO;
using System.Linq;
using System.Text;
#if ENCRYPTION_CUSTOM_PREVIEW && NET8_0_OR_GREATER
using System.Text.Json;
using System.Text.Json.Nodes;
#endif
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Azure.Cosmos.Encryption.Custom.Transformation;
Expand All @@ -28,6 +32,10 @@ internal static class EncryptionProcessor

internal static readonly CosmosJsonDotNetSerializer BaseSerializer = new (JsonSerializerSettings);

#if ENCRYPTION_CUSTOM_PREVIEW && NET8_0_OR_GREATER
private static readonly JsonWriterOptions JsonWriterOptions = new () { SkipValidation = true };
#endif

private static readonly MdeEncryptionProcessor MdeEncryptionProcessor = new ();

/// <remarks>
Expand Down Expand Up @@ -120,6 +128,60 @@ public static async Task<Stream> EncryptAsync(
return (BaseSerializer.ToStream(itemJObj), decryptionContext);
}

public static async Task<(Stream, DecryptionContext)> DecryptAsync(
Stream input,
Encryptor encryptor,
CosmosDiagnosticsContext diagnosticsContext,
JsonProcessor jsonProcessor,
CancellationToken cancellationToken)
{
return jsonProcessor switch
{
JsonProcessor.Newtonsoft => await DecryptAsync(input, encryptor, diagnosticsContext, cancellationToken),
#if ENCRYPTION_CUSTOM_PREVIEW && NET8_0_OR_GREATER
JsonProcessor.SystemTextJson => await DecryptJsonNodeAsync(input, encryptor, diagnosticsContext, cancellationToken),
#endif
_ => throw new InvalidOperationException("Unsupported Json Processor")
};
}

#if ENCRYPTION_CUSTOM_PREVIEW && NET8_0_OR_GREATER
public static async Task<(Stream, DecryptionContext)> DecryptJsonNodeAsync(
Stream input,
Encryptor encryptor,
CosmosDiagnosticsContext diagnosticsContext,
CancellationToken cancellationToken)
{
if (input == null)
{
return (input, null);
}

Debug.Assert(input.CanSeek);
Debug.Assert(encryptor != null);
Debug.Assert(diagnosticsContext != null);

JsonNode document = await JsonNode.ParseAsync(input, cancellationToken: cancellationToken);

(JsonNode decryptedDocument, DecryptionContext context) = await DecryptAsync(document, encryptor, diagnosticsContext, cancellationToken);
if (context == null)
{
input.Position = 0;
return (input, null);
}

await input.DisposeAsync();

MemoryStream ms = new ();
Utf8JsonWriter writer = new (ms, EncryptionProcessor.JsonWriterOptions);

System.Text.Json.JsonSerializer.Serialize(writer, decryptedDocument);

ms.Position = 0;
return (ms, context);
}
#endif

public static async Task<(JObject, DecryptionContext)> DecryptAsync(
JObject document,
Encryptor encryptor,
Expand All @@ -142,6 +204,53 @@ public static async Task<Stream> EncryptAsync(
return (document, decryptionContext);
}

#if ENCRYPTION_CUSTOM_PREVIEW && NET8_0_OR_GREATER
public static async Task<(JsonNode, DecryptionContext)> DecryptAsync(
JsonNode document,
Encryptor encryptor,
CosmosDiagnosticsContext diagnosticsContext,
CancellationToken cancellationToken)
{
Debug.Assert(document != null);

Debug.Assert(encryptor != null);

if (!document.AsObject().TryGetPropertyValue(Constants.EncryptedInfo, out JsonNode encryptionPropertiesNode))
{
return (document, null);
}

EncryptionProperties encryptionProperties;
try
{
encryptionProperties = System.Text.Json.JsonSerializer.Deserialize<EncryptionProperties>(encryptionPropertiesNode);
}
catch (Exception)
{
return (document, null);
}

DecryptionContext decryptionContext = await DecryptInternalAsync(encryptor, diagnosticsContext, document, encryptionProperties, cancellationToken);

return (document, decryptionContext);
}

private static async Task<DecryptionContext> DecryptInternalAsync(Encryptor encryptor, CosmosDiagnosticsContext diagnosticsContext, JsonNode itemNode, EncryptionProperties encryptionProperties, CancellationToken cancellationToken)
{
DecryptionContext decryptionContext = encryptionProperties.EncryptionAlgorithm switch
{
CosmosEncryptionAlgorithm.MdeAeadAes256CbcHmac256Randomized => await MdeEncryptionProcessor.DecryptObjectAsync(
itemNode,
encryptor,
encryptionProperties,
diagnosticsContext,
cancellationToken),
_ => throw new NotSupportedException($"Encryption Algorithm : {encryptionProperties.EncryptionAlgorithm} is not supported."),
};
return decryptionContext;
}
#endif

private static async Task<DecryptionContext> DecryptInternalAsync(Encryptor encryptor, CosmosDiagnosticsContext diagnosticsContext, JObject itemJObj, JObject encryptionPropertiesJObj, CancellationToken cancellationToken)
{
EncryptionProperties encryptionProperties = encryptionPropertiesJObj.ToObject<EncryptionProperties>();
Expand Down Expand Up @@ -225,7 +334,7 @@ private static JObject RetrieveItem(
MaxDepth = 64, // https://github.com/advisories/GHSA-5crp-9r3c-p9vr
};

itemJObj = JsonSerializer.Create(jsonSerializerSettings).Deserialize<JObject>(jsonTextReader);
itemJObj = Newtonsoft.Json.JsonSerializer.Create(jsonSerializerSettings).Deserialize<JObject>(jsonTextReader);
}

return itemJObj;
Expand Down
24 changes: 24 additions & 0 deletions Microsoft.Azure.Cosmos.Encryption.Custom/src/Transformation/JsonNodeSqlSerializer.Preview.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,30 @@ internal virtual (TypeMarker typeMarker, byte[] serializedBytes, int serializedB
return (buffer, length);
}
}

internal virtual JsonNode Deserialize(
TypeMarker typeMarker,
ReadOnlySpan<byte> serializedBytes)
{
switch (typeMarker)
{
case TypeMarker.Boolean:
return JsonValue.Create(SqlBoolSerializer.Deserialize(serializedBytes));
case TypeMarker.Double:
return JsonValue.Create(SqlDoubleSerializer.Deserialize(serializedBytes));
case TypeMarker.Long:
return JsonValue.Create(SqlLongSerializer.Deserialize(serializedBytes));
case TypeMarker.String:
return JsonValue.Create(SqlVarCharSerializer.Deserialize(serializedBytes));
case TypeMarker.Array:
return JsonNode.Parse(serializedBytes);
case TypeMarker.Object:
return JsonNode.Parse(serializedBytes);
default:
Debug.Fail($"Unexpected type marker {typeMarker}");
return null;
}
}
}
}
#endif
15 changes: 15 additions & 0 deletions ...osoft.Azure.Cosmos.Encryption.Custom/src/Transformation/MdeEncryptionProcessor.Preview.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@ namespace Microsoft.Azure.Cosmos.Encryption.Custom.Transformation
{
using System;
using System.IO;
#if NET8_0_OR_GREATER
using System.Text.Json.Nodes;
#endif
using System.Threading;
using System.Threading.Tasks;
using Newtonsoft.Json.Linq;
Expand Down Expand Up @@ -45,6 +48,18 @@ internal async Task<DecryptionContext> DecryptObjectAsync(
{
return await this.JObjectEncryptionProcessor.DecryptObjectAsync(document, encryptor, encryptionProperties, diagnosticsContext, cancellationToken);
}

#if NET8_0_OR_GREATER
internal async Task<DecryptionContext> DecryptObjectAsync(
JsonNode document,
Encryptor encryptor,
EncryptionProperties encryptionProperties,
CosmosDiagnosticsContext diagnosticsContext,
CancellationToken cancellationToken)
{
return await this.JsonNodeEncryptionProcessor.DecryptObjectAsync(document, encryptor, encryptionProperties, diagnosticsContext, cancellationToken);
}
#endif
}
}
#endif
2 changes: 1 addition & 1 deletion Microsoft.Azure.Cosmos.Encryption.Custom/src/Transformation/MdeEncryptionProcessor.Stable.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ internal async Task<DecryptionContext> DecryptObjectAsync(
continue;
}

(byte[] plainText, int decryptedCount) = this.Encryptor.Decrypt(encryptionKey, cipherTextWithTypeMarker, arrayPoolManager);
(byte[] plainText, int decryptedCount) = this.Encryptor.Decrypt(encryptionKey, cipherTextWithTypeMarker, cipherTextWithTypeMarker.Length, arrayPoolManager);

this.Serializer.DeserializeAndAddProperty(
(TypeMarker)cipherTextWithTypeMarker[0],
Expand Down
6 changes: 3 additions & 3 deletions Microsoft.Azure.Cosmos.Encryption.Custom/src/Transformation/MdeEncryptor.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,16 +54,16 @@ internal virtual (byte[], int) Encrypt(DataEncryptionKey encryptionKey, TypeMark
return (encryptedText, encryptedTextLength);
}

internal virtual (byte[] plainText, int plainTextLength) Decrypt(DataEncryptionKey encryptionKey, byte[] cipherText, ArrayPoolManager arrayPoolManager)
internal virtual (byte[] plainText, int plainTextLength) Decrypt(DataEncryptionKey encryptionKey, byte[] cipherText, int cipherTextLength, ArrayPoolManager arrayPoolManager)
{
int plainTextLength = encryptionKey.GetDecryptByteCount(cipherText.Length - 1);
int plainTextLength = encryptionKey.GetDecryptByteCount(cipherTextLength - 1);

byte[] plainText = arrayPoolManager.Rent(plainTextLength);

int decryptedLength = encryptionKey.DecryptData(
cipherText,
cipherTextOffset: 1,
cipherTextLength: cipherText.Length - 1,
cipherTextLength: cipherTextLength - 1,
plainText,
outputOffset: 0);

Expand Down
2 changes: 1 addition & 1 deletion ...zure.Cosmos.Encryption.Custom/src/Transformation/MdeJObjectEncryptionProcessor.Preview.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ internal async Task<DecryptionContext> DecryptObjectAsync(
continue;
}

(byte[] bytes, int processedBytes) = this.Encryptor.Decrypt(encryptionKey, cipherTextWithTypeMarker, arrayPoolManager);
(byte[] bytes, int processedBytes) = this.Encryptor.Decrypt(encryptionKey, cipherTextWithTypeMarker, cipherTextWithTypeMarker.Length, arrayPoolManager);

#if NET8_0_OR_GREATER
if (decompressor != null)
Expand Down
96 changes: 92 additions & 4 deletions ...ure.Cosmos.Encryption.Custom/src/Transformation/MdeJsonNodeEncryptionProcessor.Preview.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,10 @@

namespace Microsoft.Azure.Cosmos.Encryption.Custom.Transformation
{
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text.Json;
using System.Text.Json.Nodes;
using System.Threading;
Expand All @@ -16,14 +18,14 @@ namespace Microsoft.Azure.Cosmos.Encryption.Custom.Transformation

internal class MdeJsonNodeEncryptionProcessor
{
private readonly JsonWriterOptions jsonWriterOptions = new () { SkipValidation = true };

internal JsonNodeSqlSerializer Serializer { get; set; } = new JsonNodeSqlSerializer();

internal MdeEncryptor Encryptor { get; set; } = new MdeEncryptor();

internal JsonSerializerOptions JsonSerializerOptions { get; set; }

private JsonWriterOptions jsonWriterOptions = new () { SkipValidation = true };

public MdeJsonNodeEncryptionProcessor()
{
this.JsonSerializerOptions = new JsonSerializerOptions();
Expand Down Expand Up @@ -102,7 +104,7 @@ public async Task<Stream> EncryptAsync(
#endif
(byte[] encryptedBytes, int encryptedBytesCount) = this.Encryptor.Encrypt(encryptionKey, typeMarker, processedBytes, processedBytesLength, arrayPoolManager);

itemObj[propertyName] = JsonValue.Create(new JsonBytes(encryptedBytes, 0, encryptedBytesCount));
itemObj[propertyName] = JsonValue.Create(new Memory<byte>(encryptedBytes, 0, encryptedBytesCount));
pathsEncrypted.Add(pathToEncrypt);
}

Expand All @@ -125,11 +127,97 @@ public async Task<Stream> EncryptAsync(
MemoryStream ms = new ();
Utf8JsonWriter writer = new (ms, this.jsonWriterOptions);

JsonSerializer.Serialize(writer, document, this.JsonSerializerOptions);
JsonSerializer.Serialize(writer, document);

ms.Position = 0;
return ms;
}

internal async Task<DecryptionContext> DecryptObjectAsync(
JsonNode document,
Encryptor encryptor,
EncryptionProperties encryptionProperties,
CosmosDiagnosticsContext diagnosticsContext,
CancellationToken cancellationToken)
{
_ = diagnosticsContext;

if (encryptionProperties.EncryptionFormatVersion != 3 && encryptionProperties.EncryptionFormatVersion != 4)
{
throw new NotSupportedException($"Unknown encryption format version: {encryptionProperties.EncryptionFormatVersion}. Please upgrade your SDK to the latest version.");
}

using ArrayPoolManager arrayPoolManager = new ();

DataEncryptionKey encryptionKey = await encryptor.GetEncryptionKeyAsync(encryptionProperties.DataEncryptionKeyId, encryptionProperties.EncryptionAlgorithm, cancellationToken);

List<string> pathsDecrypted = new (encryptionProperties.EncryptedPaths.Count());

JsonObject itemObj = document.AsObject();

#if NET8_0_OR_GREATER
BrotliCompressor decompressor = null;
if (encryptionProperties.EncryptionFormatVersion == 4)
{
bool containsCompressed = encryptionProperties.CompressedEncryptedPaths?.Any() == true;
if (encryptionProperties.CompressionAlgorithm != CompressionOptions.CompressionAlgorithm.Brotli && containsCompressed)
{
throw new NotSupportedException($"Unknown compression algorithm {encryptionProperties.CompressionAlgorithm}");
}

if (containsCompressed)
{
decompressor = new ();
}
}
#endif

foreach (string path in encryptionProperties.EncryptedPaths)
{
string propertyName = path[1..];

if (!itemObj.TryGetPropertyValue(propertyName, out JsonNode propertyValue))
{
// malformed document, such record shouldn't be there at all
continue;
}

// can we get to internal JsonNode buffers to avoid string allocation here?
string base64String = propertyValue.GetValue<string>();
byte[] cipherTextWithTypeMarker = arrayPoolManager.Rent((base64String.Length * sizeof(char) * 3 / 4) + 4);
if (!Convert.TryFromBase64Chars(base64String, cipherTextWithTypeMarker, out int cipherTextLength))
{
continue;
}

(byte[] bytes, int processedBytes) = this.Encryptor.Decrypt(encryptionKey, cipherTextWithTypeMarker, cipherTextLength, arrayPoolManager);

#if NET8_0_OR_GREATER
if (decompressor != null)
{
if (encryptionProperties.CompressedEncryptedPaths?.TryGetValue(path, out int decompressedSize) == true)
{
byte[] buffer = arrayPoolManager.Rent(decompressedSize);
processedBytes = decompressor.Decompress(bytes, processedBytes, buffer);

bytes = buffer;
}
}
#endif
document[propertyName] = this.Serializer.Deserialize(
(TypeMarker)cipherTextWithTypeMarker[0],
bytes.AsSpan(0, processedBytes));

pathsDecrypted.Add(path);
}

DecryptionContext decryptionContext = EncryptionProcessor.CreateDecryptionContext(
pathsDecrypted,
encryptionProperties.DataEncryptionKeyId);

itemObj.Remove(Constants.EncryptedInfo);
return decryptionContext;
}
}
}

Expand Down
1 change: 1 addition & 0 deletions ...m/tests/Microsoft.Azure.Cosmos.Encryption.Custom.Performance.Tests/EncryptionBenchmark.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ await EncryptionProcessor.DecryptAsync(
new MemoryStream(this.encryptedData!),
this.encryptor,
new CosmosDiagnosticsContext(),
this.JsonProcessor,
CancellationToken.None);
}

Expand Down
Loading

0 comments on commit f88605c

Please sign in to comment.