[Storage] Fixed WrapKeyInternal in Azure.Storage.Blobs to calls corre…

…ct method (#45352)

* Initial commit

* Changed KeyWapImpl and KeyUnwrappImpl mock

* Test proxy

* Refactored mock impl

* Added testing

* Removed test

sdk/storage/Azure.Storage.Blobs/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "net",
   "TagPrefix": "net/storage/Azure.Storage.Blobs",
-  "Tag": "net/storage/Azure.Storage.Blobs_14eb1d6279"
+  "Tag": "net/storage/Azure.Storage.Blobs_f805dd22f1"
 }

sdk/storage/Azure.Storage.Blobs/src/SpecializedBlobExtensions.ClientSideEncryption.cs

@@ -188,7 +188,7 @@ private static async Task<byte[]> WrapKeyInternal(ReadOnlyMemory<byte> contentEn
                     keyWrapAlgorithm,
                     contentEncryptionKey,
                     cancellationToken).ConfigureAwait(false)
-                : key.UnwrapKey(
+                : key.WrapKey(
                     keyWrapAlgorithm,
                     contentEncryptionKey,
                     cancellationToken);

sdk/storage/Azure.Storage.Blobs/tests/ClientSideEncryptionTests.cs

@@ -1378,6 +1378,32 @@ public async Task UpdateKey(
             await CallCorrectKeyUpdateAsync(blob, useOverrides, mockKey2.Object, mockKeyResolver, version);
 
             // Assert
+            if (IsAsync)
+            {
+                mockKey1.Verify(k => k.WrapKeyAsync(s_algorithmName, IsNotNull<ReadOnlyMemory<byte>>(), s_cancellationToken), Times.Once);
+                mockKey1.Verify(k => k.UnwrapKeyAsync(s_algorithmName, IsNotNull<ReadOnlyMemory<byte>>(), s_cancellationToken), Times.Once);
+                mockKey2.Verify(k => k.WrapKeyAsync(s_algorithmName, IsNotNull<ReadOnlyMemory<byte>>(), s_cancellationToken), Times.Once);
+
+                var mockKey1_firstInvocation = mockKey1.Invocations.First();
+                var mockKey1_lastInvocation = mockKey1.Invocations.Last();
+                var mockKey2_firstInvocation = mockKey2.Invocations.First();
+                Assert.AreEqual(nameof(IKeyEncryptionKey.WrapKeyAsync), mockKey1_firstInvocation.Method.Name);
+                Assert.AreEqual(nameof(IKeyEncryptionKey.UnwrapKeyAsync), mockKey1_lastInvocation.Method.Name);
+                Assert.AreEqual(nameof(IKeyEncryptionKey.WrapKeyAsync), mockKey2_firstInvocation.Method.Name);
+            }
+            else
+            {
+                mockKey1.Verify(k => k.WrapKey(s_algorithmName, IsNotNull<ReadOnlyMemory<byte>>(), s_cancellationToken), Times.Once);
+                mockKey1.Verify(k => k.UnwrapKey(s_algorithmName, IsNotNull<ReadOnlyMemory<byte>>(), s_cancellationToken), Times.Once);
+                mockKey2.Verify(k => k.WrapKey(s_algorithmName, IsNotNull<ReadOnlyMemory<byte>>(), s_cancellationToken), Times.Once);
+
+                var mockKey1_firstInvocation = mockKey1.Invocations.First();
+                var mockKey1_lastInvocation = mockKey1.Invocations.Last();
+                var mockKey2_firstInvocation = mockKey2.Invocations.First();
+                Assert.AreEqual(nameof(IKeyEncryptionKey.WrapKey), mockKey1_firstInvocation.Method.Name);
+                Assert.AreEqual(nameof(IKeyEncryptionKey.UnwrapKey), mockKey1_lastInvocation.Method.Name);
+                Assert.AreEqual(nameof(IKeyEncryptionKey.WrapKey), mockKey2_firstInvocation.Method.Name);
+            }
             await AssertKeyAsync(blob, mockKey2.Object, cek);
         }
 

sdk/storage/Azure.Storage.Common/tests/Shared/ClientSideEncryptionTestExtensions.cs

@@ -57,7 +57,7 @@ public static Mock<IKeyEncryptionKey> GetIKeyEncryptionKey(
                     .Returns<string, ReadOnlyMemory<byte>, CancellationToken>(async (algorithm, wrappedKey, cancellationToken) =>
                     {
                         await Task.Delay(optionalDelay);
-                        return KeyWrapImpl(userKeyBytes, wrappedKey.ToArray());
+                        return KeyUnwrapImpl(userKeyBytes, wrappedKey.ToArray());
                     });
             }
             else
@@ -72,7 +72,7 @@ public static Mock<IKeyEncryptionKey> GetIKeyEncryptionKey(
                     .Returns<string, ReadOnlyMemory<byte>, CancellationToken>((algorithm, wrappedKey, cancellationToken) =>
                     {
                         Thread.Sleep(optionalDelay);
-                        return KeyWrapImpl(userKeyBytes, wrappedKey.ToArray());
+                        return KeyUnwrapImpl(userKeyBytes, wrappedKey.ToArray());
                     });
             }
 
@@ -160,7 +160,7 @@ public static Mock<IKeyEncryptionKeyResolver> GetAlwaysFailsKeyResolver(this Rec
                     Tuple.Create(KeyWrapImpl(userKeyBytes, key), s_algorithmName)));
             keyMock.Setup(k => k.UnwrapKeyAsync(It.IsNotNull<byte[]>(), s_algorithmName, It.IsNotNull<CancellationToken>()))
                 .Returns<byte[], string, CancellationToken>((wrappedKey, algorithm, cancellationToken) => Task.FromResult(
-                    KeyWrapImpl(userKeyBytes, wrappedKey)));
+                    KeyUnwrapImpl(userKeyBytes, wrappedKey)));
 
             return keyMock;
         }
@@ -177,8 +177,29 @@ public static Mock<IKeyEncryptionKeyResolver> GetAlwaysFailsKeyResolver(this Rec
         private static byte[] KeyWrapImpl(byte[] key, byte[] contents)
         {
             var result = new byte[contents.Length];
-            // just bitflip the contents
-            new System.Collections.BitArray(contents).Not().CopyTo(result, 0);
+
+            if (contents.Length == 0) return result;
+
+            // Move each byte one position to the left
+            new Span<byte>(contents, 1, contents.Length - 1).CopyTo(result);
+
+            // Fill the last byte with the first byte (loop-around)
+            result[contents.Length - 1] = contents[0];
+
+            return result;
+        }
+
+        private static byte[] KeyUnwrapImpl(byte[] key, byte[] contents)
+        {
+            var result = new byte[contents.Length];
+
+            if (contents.Length == 0) return result;
+
+            // Move each byte one position to the right
+            new Span<byte>(contents, 0, contents.Length - 1).CopyTo(new Span<byte>(result, 1, result.Length - 1));
+
+            // Fill the first byte with the last byte (loop-around)
+            result[0] = contents[contents.Length - 1];
 
             return result;
         }