From c1c4463c2f1582f17b43a8be6673369f5bcc0e23 Mon Sep 17 00:00:00 2001
From: Sameeksha Vaity <savaity@microsoft.com>
Date: Wed, 6 Oct 2021 14:31:12 -0700
Subject: [PATCH] Fix redirect policy to check header name for valid Redirect
 requests (#24570)

---
 .../http/policy/DefaultRedirectStrategy.java  | 20 +++++++-------
 .../core/http/policy/RedirectPolicyTest.java  | 27 ++++++++++++++++++-
 2 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/DefaultRedirectStrategy.java b/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/DefaultRedirectStrategy.java
index dfdae5eeca7cf..d2e44f197ddb5 100644
--- a/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/DefaultRedirectStrategy.java
+++ b/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/DefaultRedirectStrategy.java
@@ -88,17 +88,19 @@ public DefaultRedirectStrategy(int maxAttempts, String locationHeader, Set<HttpM
     public boolean shouldAttemptRedirect(HttpPipelineCallContext context,
                                          HttpResponse httpResponse, int tryCount,
                                          Set<String> attemptedRedirectUrls) {
-        String redirectUrl = tryGetRedirectHeader(httpResponse.getHeaders(), getLocationHeader());
 
-        if (isValidRedirectCount(tryCount)
-            && redirectUrl != null
-            && !alreadyAttemptedRedirectUrl(redirectUrl, attemptedRedirectUrls)
-            && isValidRedirectStatusCode(httpResponse.getStatusCode())
+        if (isValidRedirectStatusCode(httpResponse.getStatusCode())
+            && isValidRedirectCount(tryCount)
             && isAllowedRedirectMethod(httpResponse.getRequest().getHttpMethod())) {
-            logger.verbose("[Redirecting] Try count: {}, Attempted Redirect URLs: {}", tryCount,
-                attemptedRedirectUrls.toString());
-            attemptedRedirectUrls.add(redirectUrl);
-            return true;
+            String redirectUrl = tryGetRedirectHeader(httpResponse.getHeaders(), getLocationHeader());
+            if (redirectUrl != null && !alreadyAttemptedRedirectUrl(redirectUrl, attemptedRedirectUrls)) {
+                logger.verbose("[Redirecting] Try count: {}, Attempted Redirect URLs: {}", tryCount,
+                    attemptedRedirectUrls.toString());
+                attemptedRedirectUrls.add(redirectUrl);
+                return true;
+            } else {
+                return false;
+            }
         } else {
             return false;
         }
diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/RedirectPolicyTest.java b/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/RedirectPolicyTest.java
index 6934647a1b7be..872ed99804364 100644
--- a/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/RedirectPolicyTest.java
+++ b/sdk/core/azure-core/src/test/java/com/azure/core/http/policy/RedirectPolicyTest.java
@@ -296,11 +296,36 @@ public void redirectForMultipleRequests() throws MalformedURLException {
             new URL("http://localhost/"))).block();
 
         assertEquals(4, httpClient.getCount());
-        // Both requests successfully redirected for same request redirect lcoation
+        // Both requests successfully redirected for same request redirect location
         assertEquals(200, response1.getStatusCode());
         assertEquals(200, response2.getStatusCode());
     }
 
+    @Test
+    public void nonRedirectRequest() throws MalformedURLException {
+        final HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(new NoOpHttpClient() {
+
+                @Override
+                public Mono<HttpResponse> send(HttpRequest request) {
+                    if (request.getUrl().toString().equals("http://localhost/")) {
+                        Map<String, String> headers = new HashMap<>();
+                        HttpHeaders httpHeader = new HttpHeaders(headers);
+                        return Mono.just(new MockHttpResponse(request, 401, httpHeader));
+                    } else {
+                        return Mono.just(new MockHttpResponse(request, 200));
+                    }
+                }
+            })
+            .policies(new RedirectPolicy())
+            .build();
+
+        HttpResponse response = pipeline.send(new HttpRequest(HttpMethod.GET,
+            new URL("http://localhost/"))).block();
+
+        assertEquals(401, response.getStatusCode());
+    }
+
     static class RecordingHttpClient implements HttpClient {
 
         private final AtomicInteger count = new AtomicInteger();