Merge branch 'master' into copilot/fix-5154

Author: NaluTripician

Directory.Build.props

@@ -1,9 +1,9 @@
 <Project xmlns="http://schemas.microsoft.com/developer/msbuild/2003">
 	<PropertyGroup>
-		<ClientOfficialVersion>3.54.1</ClientOfficialVersion>
-		<ClientPreviewVersion>3.55.0</ClientPreviewVersion>
-		<ClientPreviewSuffixVersion>preview.1</ClientPreviewSuffixVersion>
-		<DirectVersion>3.41.0</DirectVersion>
+		<ClientOfficialVersion>3.56.0</ClientOfficialVersion>
+		<ClientPreviewVersion>3.57.0</ClientPreviewVersion>
+		<ClientPreviewSuffixVersion>preview.0</ClientPreviewSuffixVersion>
+		<DirectVersion>3.41.2</DirectVersion>
 		<FaultInjectionVersion>1.0.0</FaultInjectionVersion>
 		<FaultInjectionSuffixVersion>beta.0</FaultInjectionSuffixVersion>
 		<EncryptionOfficialVersion>2.0.5</EncryptionOfficialVersion>

Exceptions.md

@@ -33,3 +33,60 @@ Cosmos DB SDK on any IO failure will attempt to retry the failed operation if re
 ## Common error status codes and troubleshooting guide <a id="error-codes"></a>
 
 To see a list of common error code and issues please see [.NET SDK troubleshooting guide](https://docs.microsoft.com/azure/cosmos-db/troubleshoot-dot-net-sdk)
+
+### Disambiguating 404 (Not Found) errors using SubStatusCode <a id="substatus-codes"></a>
+
+When you receive a 404 (Not Found) status code from Cosmos DB, it can indicate two different scenarios:
+1. **Item not found**: The requested item doesn't exist in the container
+2. **Owner resource not found**: The parent resource (container or database) doesn't exist
+
+To distinguish between these cases, check the `SubStatusCode` property:
+
+- **SubStatusCode 0**: Regular item not found (the item doesn't exist in an existing container)
+- **SubStatusCode 1003**: Owner resource not found (the container or database doesn't exist)
+
+#### Example with Typed APIs (throws CosmosException):
+
+```csharp
+try
+{
+    ItemResponse<MyItem> response = await container.ReadItemAsync<MyItem>("itemId", new PartitionKey("partitionKey"));
+}
+catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound)
+{
+    if (ex.SubStatusCode == 1003)
+    {
+        // The container or database doesn't exist
+        Console.WriteLine("Owner resource (container/database) not found");
+    }
+    else
+    {
+        // The item doesn't exist in an existing container
+        Console.WriteLine("Item not found");
+    }
+}
+```
+
+#### Example with Stream APIs (returns ResponseMessage):
+
+```csharp
+ResponseMessage response = await container.ReadItemStreamAsync("itemId", new PartitionKey("partitionKey"));
+
+if (response.StatusCode == HttpStatusCode.NotFound)
+{
+    int subStatusCode = (int)response.Headers.SubStatusCode;
+    
+    if (subStatusCode == 1003)
+    {
+        // The container or database doesn't exist
+        Console.WriteLine("Owner resource (container/database) not found");
+    }
+    else
+    {
+        // The item doesn't exist in an existing container
+        Console.WriteLine("Item not found");
+    }
+}
+```
+
+This distinction is particularly useful when implementing retry logic or error handling strategies, as you may want to handle these scenarios differently (e.g., creating the container if it doesn't exist vs. handling a missing item).

Microsoft.Azure.Cosmos.Encryption.Custom/src/AeadAes/SymmetricKey.cs

@@ -23,10 +23,11 @@ internal class SymmetricKey
         /// <param name="rootKey">root key</param>
         internal SymmetricKey(byte[] rootKey)
         {
-            // Key validation
-            if (rootKey == null || rootKey.Length == 0)
+            ArgumentValidation.ThrowIfNull(rootKey);
+
+            if (rootKey.Length == 0)
             {
-                throw new ArgumentNullException(nameof(rootKey));
+                throw new ArgumentException("The root key cannot be empty.", nameof(rootKey));
             }
 
             this.rootKey = rootKey;

Microsoft.Azure.Cosmos.Encryption.Custom/src/Common/CosmosDiagnosticsContext.cs

@@ -5,32 +5,84 @@
 namespace Microsoft.Azure.Cosmos.Encryption.Custom
 {
     using System;
+    using System.Collections.Generic;
+    using System.Diagnostics;
+    using System.Threading;
 
     /// <summary>
-    /// This is an empty implementation of CosmosDiagnosticsContext which has been plumbed through the DataEncryptionKeyProvider and EncryptionContainer.
-    /// This may help adding diagnostics more easily in future.
+    /// Lightweight diagnostics context for Custom Encryption extension.
+    /// Manages Activity creation for OpenTelemetry integration.
     /// </summary>
     internal class CosmosDiagnosticsContext
     {
-        private static readonly CosmosDiagnosticsContext UnusedSingleton = new ();
-        private static readonly IDisposable UnusedScopeSingleton = new Scope();
+        private static readonly ActivitySource ActivitySource = new ("Microsoft.Azure.Cosmos.Encryption.Custom");
 
+        /// <summary>
+        /// Scope name prefix for MDE (Microsoft.Data.Encryption) encrypt operations.
+        /// </summary>
+        internal const string ScopeEncryptModeSelectionPrefix = "EncryptionProcessor.Encrypt.Mde.";
+
+        /// <summary>
+        /// Scope name prefix for MDE (Microsoft.Data.Encryption) decrypt operations.
+        /// </summary>
+        internal const string ScopeDecryptModeSelectionPrefix = "EncryptionProcessor.Decrypt.Mde.";
+
+        internal CosmosDiagnosticsContext()
+        {
+        }
+
+        /// <summary>
+        /// Creates a new diagnostics context instance.
+        /// </summary>
         public static CosmosDiagnosticsContext Create(RequestOptions options)
         {
             _ = options;
-            return CosmosDiagnosticsContext.UnusedSingleton;
+            return new CosmosDiagnosticsContext();
         }
 
-        public IDisposable CreateScope(string scope)
+        /// <summary>
+        /// Creates a new diagnostic scope for Activity tracking.
+        /// </summary>
+        /// <param name="scope">The name of the scope.</param>
+        /// <returns>A <see cref="Scope"/> that manages an Activity lifecycle.</returns>
+        /// <exception cref="ArgumentNullException">Thrown when <paramref name="scope"/> is null.</exception>
+        /// <exception cref="ArgumentException">Thrown when <paramref name="scope"/> is empty or whitespace.</exception>
+        /// <remarks>
+        /// Use with a <c>using</c> statement to ensure proper disposal.
+        /// </remarks>
+        public Scope CreateScope(string scope)
         {
-            _ = scope;
-            return CosmosDiagnosticsContext.UnusedScopeSingleton;
+            ArgumentValidation.ThrowIfNullOrWhiteSpace(scope, nameof(scope));
+
+            Activity activity = ActivitySource.HasListeners() ? ActivitySource.StartActivity(scope, ActivityKind.Internal) : null;
+
+            return new Scope(activity);
         }
 
-        private class Scope : IDisposable
+        /// <summary>
+        /// Represents a diagnostic scope for Activity tracking.
+        /// Must be used with the 'using' pattern to ensure proper disposal.
+        /// </summary>
+        /// <remarks>
+        /// Dispose() is idempotent - calling it multiple times will only dispose the Activity once.
+        /// </remarks>
+        public sealed class Scope : IDisposable
         {
+            private readonly Activity activity;
+            private bool isDisposed;
+
+            internal Scope(Activity activity)
+            {
+                this.activity = activity;
+            }
+
             public void Dispose()
             {
+                if (!this.isDisposed)
+                {
+                    this.isDisposed = true;
+                    this.activity?.Dispose();
+                }
             }
         }
     }

Microsoft.Azure.Cosmos.Encryption.Custom/src/Common/CosmosJsonDotNetSerializer.cs

@@ -40,14 +40,7 @@ internal CosmosJsonDotNetSerializer(JsonSerializerSettings jsonSerializerSetting
         /// <returns>The object representing the deserialized stream</returns>
         public T FromStream<T>(Stream stream)
         {
-#if NET8_0_OR_GREATER
-            ArgumentNullException.ThrowIfNull(stream);
-#else
-            if (stream == null)
-            {
-                throw new ArgumentNullException(nameof(stream));
-            }
-#endif
+            ArgumentValidation.ThrowIfNull(stream);
 
             if (typeof(Stream).IsAssignableFrom(typeof(T)))
             {
@@ -87,6 +80,41 @@ public MemoryStream ToStream<T>(T input)
             return streamPayload;
         }
 
+        /// <summary>
+        /// Serializes an object directly into the provided output stream (which remains open).
+        /// </summary>
+        /// <typeparam name="T">Type of object being serialized.</typeparam>
+        /// <param name="input">Object to serialize.</param>
+        /// <param name="output">Destination stream. Must be writable. The stream is not disposed by this method.</param>
+        /// <exception cref="ArgumentNullException">Thrown when <paramref name="output"/> is <c>null</c>.</exception>
+        /// <exception cref="ArgumentException">Thrown when <paramref name="output"/> is not writable.</exception>
+        /// <remarks>
+        /// <para>This method serializes the object directly to the provided stream without creating an intermediate MemoryStream,
+        /// reducing memory allocations for large objects.</para>
+        /// <para>After writing, the stream position will be at the end of the written content.
+        /// Callers are responsible for resetting the stream position if needed for subsequent reads.</para>
+        /// </remarks>
+        public void WriteToStream<T>(T input, Stream output)
+        {
+            ArgumentValidation.ThrowIfNull(output);
+
+            if (!output.CanWrite)
+            {
+                throw new ArgumentException("Output stream must be writable", nameof(output));
+            }
+
+            using (StreamWriter streamWriter = new (output, encoding: CosmosJsonDotNetSerializer.DefaultEncoding, bufferSize: 1024, leaveOpen: true))
+            using (JsonTextWriter writer = new (streamWriter))
+            {
+                writer.ArrayPool = JsonArrayPool.Instance;
+                writer.Formatting = Newtonsoft.Json.Formatting.None;
+                JsonSerializer jsonSerializer = this.GetSerializer();
+                jsonSerializer.Serialize(writer, input);
+                writer.Flush();
+                streamWriter.Flush();
+            }
+        }
+
         /// <summary>
         /// JsonSerializer has hit a race conditions with custom settings that cause null reference exception.
         /// To avoid the race condition a new JsonSerializer is created for each call

Microsoft.Azure.Cosmos.Encryption.Custom/src/Common/StreamExtensions.cs

@@ -0,0 +1,31 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Encryption.Custom
+{
+    using System.IO;
+    using System.Threading.Tasks;
+
+    /// <summary>
+    /// Extension methods for Stream to provide compatibility across different .NET versions.
+    /// </summary>
+    internal static class StreamExtensions
+    {
+        /// <summary>
+        /// Asynchronously disposes the stream in a version-compatible way.
+        /// Uses DisposeAsync on .NET 8.0+ and falls back to synchronous Dispose on earlier versions.
+        /// </summary>
+        /// <param name="stream">The stream to dispose.</param>
+        /// <returns>A ValueTask representing the asynchronous dispose operation.</returns>
+        public static ValueTask DisposeCompatAsync(this Stream stream)
+        {
+#if NET8_0_OR_GREATER
+            return stream.DisposeAsync();
+#else
+            stream.Dispose();
+            return default;
+#endif
+        }
+    }
+}

Microsoft.Azure.Cosmos.Encryption.Custom/src/CosmosDataEncryptionKeyProvider.cs

@@ -148,14 +148,7 @@ public async Task InitializeAsync(
                 throw new InvalidOperationException($"{nameof(CosmosDataEncryptionKeyProvider)} has already been initialized.");
             }
 
-#if NET8_0_OR_GREATER
-            ArgumentNullException.ThrowIfNull(database);
-#else
-            if (database == null)
-            {
-                throw new ArgumentNullException(nameof(database));
-            }
-#endif
+            ArgumentValidation.ThrowIfNull(database);
 
             ContainerResponse containerResponse = await database.CreateContainerIfNotExistsAsync(
                 containerId,

Microsoft.Azure.Cosmos.Encryption.Custom/src/DataEncryptionKey.cs

@@ -102,14 +102,7 @@ public static DataEncryptionKey Create(
             byte[] rawKey,
             string encryptionAlgorithm)
         {
-#if NET8_0_OR_GREATER
-            ArgumentNullException.ThrowIfNull(rawKey);
-#else
-            if (rawKey == null)
-            {
-                throw new ArgumentNullException(nameof(rawKey));
-            }
-#endif
+            ArgumentValidation.ThrowIfNull(rawKey);
 
 #pragma warning disable CS0618 // Type or member is obsolete
             if (!string.Equals(encryptionAlgorithm, CosmosEncryptionAlgorithm.AEAes256CbcHmacSha256Randomized))

Microsoft.Azure.Cosmos.Encryption.Custom/src/DataEncryptionKeyContainerCore.cs

@@ -57,24 +57,14 @@ public override async Task<ItemResponse<DataEncryptionKeyProperties>> CreateData
             ItemRequestOptions requestOptions = null,
             CancellationToken cancellationToken = default)
         {
-            if (string.IsNullOrEmpty(id))
-            {
-                throw new ArgumentNullException(nameof(id));
-            }
+            ArgumentValidation.ThrowIfNullOrEmpty(id);
 
             if (!CosmosEncryptionAlgorithm.VerifyIfSupportedAlgorithm(encryptionAlgorithm))
             {
                 throw new ArgumentException(string.Format("Unsupported Encryption Algorithm {0}", encryptionAlgorithm), nameof(encryptionAlgorithm));
             }
 
-#if NET8_0_OR_GREATER
-            ArgumentNullException.ThrowIfNull(encryptionKeyWrapMetadata);
-#else
-            if (encryptionKeyWrapMetadata == null)
-            {
-                throw new ArgumentNullException(nameof(encryptionKeyWrapMetadata));
-            }
-#endif
+            ArgumentValidation.ThrowIfNull(encryptionKeyWrapMetadata);
 
             CosmosDiagnosticsContext diagnosticsContext = CosmosDiagnosticsContext.Create(requestOptions);
 
@@ -159,14 +149,7 @@ public override async Task<ItemResponse<DataEncryptionKeyProperties>> RewrapData
            ItemRequestOptions requestOptions = null,
            CancellationToken cancellationToken = default)
         {
-#if NET8_0_OR_GREATER
-            ArgumentNullException.ThrowIfNull(newWrapMetadata);
-#else
-            if (newWrapMetadata == null)
-            {
-                throw new ArgumentNullException(nameof(newWrapMetadata));
-            }
-#endif
+            ArgumentValidation.ThrowIfNull(newWrapMetadata);
 
             CosmosDiagnosticsContext diagnosticsContext = CosmosDiagnosticsContext.Create(requestOptions);
 

Microsoft.Azure.Cosmos.Encryption.Custom/src/DataEncryptionKeyContainerInlineCore.cs

@@ -56,10 +56,7 @@ public override Task<ItemResponse<DataEncryptionKeyProperties>> ReadDataEncrypti
             ItemRequestOptions requestOptions = null,
             CancellationToken cancellationToken = default)
         {
-            if (string.IsNullOrEmpty(id))
-            {
-                throw new ArgumentNullException(nameof(id));
-            }
+            ArgumentValidation.ThrowIfNullOrEmpty(id);
 
             return TaskHelper.RunInlineIfNeededAsync(() =>
                 this.dataEncryptionKeyContainerCore.ReadDataEncryptionKeyAsync(id, requestOptions, cancellationToken));
@@ -73,19 +70,8 @@ public override Task<ItemResponse<DataEncryptionKeyProperties>> RewrapDataEncryp
            ItemRequestOptions requestOptions = null,
            CancellationToken cancellationToken = default)
         {
-            if (string.IsNullOrEmpty(id))
-            {
-                throw new ArgumentNullException(nameof(id));
-            }
-
-#if NET8_0_OR_GREATER
-            ArgumentNullException.ThrowIfNull(newWrapMetadata);
-#else
-            if (newWrapMetadata == null)
-            {
-                throw new ArgumentNullException(nameof(newWrapMetadata));
-            }
-#endif
+            ArgumentValidation.ThrowIfNullOrEmpty(id);
+            ArgumentValidation.ThrowIfNull(newWrapMetadata);
 
             return TaskHelper.RunInlineIfNeededAsync(() =>
                 this.dataEncryptionKeyContainerCore.RewrapDataEncryptionKeyAsync(id, newWrapMetadata, encryptionAlgorithm, requestOptions, cancellationToken));