Handle AzureSasCredential.Update on retries (#32888)

* Handle AzureSasCredential.Update on retries
Azure · Dec 14, 2022 · e786c1b · e786c1b
1 parent bf0d552
commit e786c1b
Show file tree

Hide file tree

Showing 9 changed files with 112 additions and 3 deletions.
diff --git a/sdk/core/Azure.Core/api/Azure.Core.net461.cs b/sdk/core/Azure.Core/api/Azure.Core.net461.cs
@@ -40,7 +40,11 @@ public partial class AzureSasCredential
     {
         public AzureSasCredential(string signature) { }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
+        public string PreviousSignature { get { throw null; } }
+        [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
         public string Signature { get { throw null; } }
+        [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
+        public long SignatureUpdated { get { throw null; } }
         public void Update(string signature) { }
     }
     [System.FlagsAttribute]

diff --git a/sdk/core/Azure.Core/api/Azure.Core.net5.0.cs b/sdk/core/Azure.Core/api/Azure.Core.net5.0.cs
@@ -40,7 +40,11 @@ public partial class AzureSasCredential
     {
         public AzureSasCredential(string signature) { }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
+        public string PreviousSignature { get { throw null; } }
+        [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
         public string Signature { get { throw null; } }
+        [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
+        public long SignatureUpdated { get { throw null; } }
         public void Update(string signature) { }
     }
     [System.FlagsAttribute]

diff --git a/sdk/core/Azure.Core/api/Azure.Core.net6.0.cs b/sdk/core/Azure.Core/api/Azure.Core.net6.0.cs
@@ -40,7 +40,11 @@ public partial class AzureSasCredential
     {
         public AzureSasCredential(string signature) { }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
+        public string PreviousSignature { get { throw null; } }
+        [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
         public string Signature { get { throw null; } }
+        [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
+        public long SignatureUpdated { get { throw null; } }
         public void Update(string signature) { }
     }
     [System.FlagsAttribute]

diff --git a/sdk/core/Azure.Core/api/Azure.Core.netcoreapp2.1.cs b/sdk/core/Azure.Core/api/Azure.Core.netcoreapp2.1.cs
@@ -40,7 +40,11 @@ public partial class AzureSasCredential
     {
         public AzureSasCredential(string signature) { }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
+        public string PreviousSignature { get { throw null; } }
+        [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
         public string Signature { get { throw null; } }
+        [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
+        public long SignatureUpdated { get { throw null; } }
         public void Update(string signature) { }
     }
     [System.FlagsAttribute]

diff --git a/sdk/core/Azure.Core/api/Azure.Core.netstandard2.0.cs b/sdk/core/Azure.Core/api/Azure.Core.netstandard2.0.cs
@@ -40,7 +40,11 @@ public partial class AzureSasCredential
     {
         public AzureSasCredential(string signature) { }
         [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
+        public string PreviousSignature { get { throw null; } }
+        [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
         public string Signature { get { throw null; } }
+        [System.ComponentModel.EditorBrowsableAttribute(System.ComponentModel.EditorBrowsableState.Never)]
+        public long SignatureUpdated { get { throw null; } }
         public void Update(string signature) { }
     }
     [System.FlagsAttribute]

diff --git a/sdk/core/Azure.Core/src/AzureSasCredential.cs b/sdk/core/Azure.Core/src/AzureSasCredential.cs
@@ -1,6 +1,9 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
 using System.ComponentModel;
 using System.Threading;
 using Azure.Core;
@@ -14,6 +17,8 @@ namespace Azure
     public class AzureSasCredential
     {
         private string _signature;
+        private string _previousSignature;
+        private long _signatureUpdated;
 
         /// <summary>
         /// Shared access signature used to authenticate to an Azure service.
@@ -25,6 +30,27 @@ public string Signature
             private set => Volatile.Write(ref _signature, value);
         }
 
+        /// <summary>
+        /// The ticks value of the DateTimeOffset when the shared access signature was last set.
+        /// </summary>
+        [EditorBrowsable(EditorBrowsableState.Never)]
+        public long SignatureUpdated
+        {
+            get => Volatile.Read(ref _signatureUpdated);
+            private set => Volatile.Write(ref _signatureUpdated, value);
+        }
+
+        /// <summary>
+        /// The previous value af <see cref="Signature"/> after <see cref="Update"/> is called.
+        /// </summary>
+        /// <returns></returns>
+        [EditorBrowsable(EditorBrowsableState.Never)]
+        public string PreviousSignature
+        {
+            get => Volatile.Read(ref _previousSignature);
+            private set => Volatile.Write(ref _previousSignature, value);
+        }
+
         /// <summary>
         /// Initializes a new instance of the <see cref="AzureSasCredential"/> class.
         /// </summary>
@@ -36,7 +62,13 @@ public string Signature
         /// Thrown when the <paramref name="signature"/> is empty.
         /// </exception>
 #pragma warning disable CS8618 // Non-nullable field is uninitialized. Consider declaring as nullable.
-        public AzureSasCredential(string signature) => Update(signature);
+        public AzureSasCredential(string signature)
+        {
+            Argument.AssertNotNullOrWhiteSpace(signature, nameof(signature));
+            SignatureUpdated = DateTimeOffset.UtcNow.Ticks;
+            Signature = signature;
+            PreviousSignature = signature;
+        }
 #pragma warning restore CS8618 // Non-nullable field is uninitialized. Consider declaring as nullable.
 
         /// <summary>
@@ -54,6 +86,9 @@ public string Signature
         public void Update(string signature)
         {
             Argument.AssertNotNullOrWhiteSpace(signature, nameof(signature));
+            SignatureUpdated = DateTimeOffset.UtcNow.Ticks;
+            PreviousSignature = Signature.StartsWith("?", StringComparison.InvariantCulture) ?
+                Signature.Substring(1) : Signature;
             Signature = signature;
         }
     }

diff --git a/sdk/core/Azure.Core/src/Shared/AzureSasCredentialSynchronousPolicy.cs b/sdk/core/Azure.Core/src/Shared/AzureSasCredentialSynchronousPolicy.cs
@@ -31,7 +31,21 @@ public override void OnSendingRequest(HttpMessage message)
             }
             if (!query.Contains(signature))
             {
-                query = string.IsNullOrEmpty(query) ? '?' + signature : query + '&' + signature;
+                bool setSignature = false;
+                // check if the signature has updated since we started processing this message
+                if (message.ProcessingContext.StartTime.Ticks < _credential.SignatureUpdated)
+                {
+                    string previousSig = _credential.PreviousSignature;
+                    if (query.Contains(previousSig))
+                    {
+                        query = query.Replace(previousSig, signature);
+                        setSignature = true;
+                    }
+                }
+                if (!setSignature)
+                {
+                    query = string.IsNullOrEmpty(query) ? '?' + signature : query + '&' + signature;
+                }
                 message.Request.Uri.Query = query;
             }
 

diff --git a/sdk/core/Azure.Core/tests/AzureSasCredentialSynchronousPolicyTests.cs b/sdk/core/Azure.Core/tests/AzureSasCredentialSynchronousPolicyTests.cs
@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using System.Net;
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core.Pipeline;
@@ -92,5 +93,43 @@ public async Task IgnoreDuplicateSas(string signatureValue)
 
             Assert.AreEqual(query, transport.Requests[0].Uri.Query);
         }
+
+        [Test]
+        public async Task VerifyRetryAfterSasCredentialUpdateOperation()
+        {
+            // Arrange
+            const string INITIAL_QUERY_URI = "?foo=bar";
+            const string FIRST_SIGNATURE_VALUE = "?sig=first_signature_value";
+            const string SECOND_SIGNATURE_VALUE = "?sig=second_signature_value";
+            int callCount = 0;
+            var azureSasCredential = new AzureSasCredential(FIRST_SIGNATURE_VALUE);
+            var transport = new MockTransport((req) =>
+            {
+                if (callCount++ == 0)
+                {
+                    Assert.AreEqual("?foo=bar&sig=first_signature_value", req.Uri.Query);
+                    azureSasCredential.Update(SECOND_SIGNATURE_VALUE);
+                    Thread.Sleep(100);
+                    return new MockResponse(429);
+                }
+                else
+                {
+                    Assert.AreEqual("?foo=bar&sig=second_signature_value", req.Uri.Query);
+                    return new MockResponse(200);
+                }
+            });
+            var sasPolicy = new AzureSasCredentialSynchronousPolicy(azureSasCredential);
+
+            // Act + Assert
+            using (Request request = transport.CreateRequest())
+            {
+                request.Method = RequestMethod.Get;
+                request.Uri.Query = INITIAL_QUERY_URI;
+                var pipeline = new HttpPipeline(transport, new HttpPipelinePolicy[] { new DefaultRetryPolicy(new RetryOptions()), sasPolicy });
+
+                Response response = await pipeline.SendRequestAsync(request, CancellationToken.None).ConfigureAwait(false);
+                Assert.AreEqual((int)HttpStatusCode.OK, response.Status);
+            }
+        }
     }
 }
diff --git a/sdk/tables/Azure.Data.Tables/src/Azure.Data.Tables.csproj b/sdk/tables/Azure.Data.Tables/src/Azure.Data.Tables.csproj
@@ -12,7 +12,8 @@
     <IncludeOperationsSharedSource>true</IncludeOperationsSharedSource>
   </PropertyGroup>
   <ItemGroup>
-    <PackageReference Include="Azure.Core" />
+    <!-- <PackageReference Include="Azure.Core" /> -->
+    <ProjectReference Include="..\..\..\core\Azure.Core\src\Azure.Core.csproj" />
     <PackageReference Include="System.Text.Json" />
   </ItemGroup>