Fix memory leak and enable OkHttp support for ACR blob updates. (Azure#28024)

* ACR blob updates.

* Disable playback tests on non-windows devices.

* Update the changelog

* Disable client tests for playback mode.

Author: pallavit

sdk/containerregistry/azure-containers-containerregistry/CHANGELOG.md

@@ -1,9 +1,10 @@
 # Release History
 
-## 1.1.0-beta.1 (Unreleased)
+## 1.1.0-beta.1 (2022-04-08)
 
 ### Features Added
 - Added interfaces from `com.azure.core.client.traits` to `ContainerRegistryClientBuilder`.
+- Added support for `ContainerRegistryBlobAsyncClient`.
 
 ### Breaking Changes
 

sdk/containerregistry/azure-containers-containerregistry/src/main/java/com/azure/containers/containerregistry/implementation/ContainerRegistriesImpl.java

@@ -298,7 +298,8 @@ public Mono<Response<Void>> checkDockerV2SupportWithResponseAsync(Context contex
     @ServiceMethod(returns = ReturnType.SINGLE)
     public Mono<Response<ManifestWrapper>> getManifestWithResponseAsync(
             String name, String reference, String accept, Context context) {
-        final String acceptParam = "application/json";
+        // TODO: Bug in autorest that we do not allow picking up acceptParam from the parameters.
+        final String acceptParam = accept;
         return service.getManifest(this.client.getUrl(), name, reference, accept, acceptParam, context);
     }
 

sdk/containerregistry/azure-containers-containerregistry/src/main/java/com/azure/containers/containerregistry/implementation/ContainerRegistryBlobsImpl.java

@@ -86,7 +86,7 @@ Mono<ContainerRegistryBlobsCheckBlobExistsResponse> checkBlobExists(
                 Context context);
 
         @Delete("/v2/{name}/blobs/{digest}")
-        @ExpectedResponses({202})
+        @ExpectedResponses({202, 404})
         @UnexpectedResponseExceptionType(HttpResponseException.class)
         Mono<StreamResponse> deleteBlob(
                 @HostParam("url") String url,

sdk/containerregistry/azure-containers-containerregistry/src/main/java/com/azure/containers/containerregistry/implementation/ContainerRegistryRedirectPolicy.java

@@ -0,0 +1,139 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.containers.containerregistry.implementation;
+
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpMethod;
+import com.azure.core.http.HttpPipelineCallContext;
+import com.azure.core.http.HttpPipelineNextPolicy;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpResponse;
+import com.azure.core.http.policy.HttpPipelinePolicy;
+import com.azure.core.util.CoreUtils;
+import com.azure.core.util.logging.ClientLogger;
+import reactor.core.publisher.Mono;
+
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+
+import static com.azure.containers.containerregistry.implementation.UtilsImpl.DOCKER_DIGEST_HEADER_NAME;
+
+/**
+ * <p> Redirect policy for the container registry.</p>
+ *
+ * <p> This reads some of the headers that are returned from the redirect call that core redirect policy does not handle.</p>
+ */
+public final class ContainerRegistryRedirectPolicy implements HttpPipelinePolicy {
+    private static final ClientLogger LOGGER = new ClientLogger(com.azure.core.http.policy.DefaultRedirectStrategy.class);
+    private static final int MAX_REDIRECT_ATTEMPTS;
+    private static final String REDIRECT_LOCATION_HEADER_NAME;
+    private static final int PERMANENT_REDIRECT_STATUS_CODE;
+    private static final int TEMPORARY_REDIRECT_STATUS_CODE;
+    private static final Set<HttpMethod> REDIRECT_ALLOWED_METHODS;
+    private static final String AUTHORIZATION;
+
+    static {
+        REDIRECT_ALLOWED_METHODS = new HashSet<>(Arrays.asList(HttpMethod.GET, HttpMethod.HEAD));
+        PERMANENT_REDIRECT_STATUS_CODE = 308;
+        TEMPORARY_REDIRECT_STATUS_CODE = 307;
+        REDIRECT_LOCATION_HEADER_NAME = "Location";
+        MAX_REDIRECT_ATTEMPTS = 3;
+        AUTHORIZATION = "Authorization";
+    }
+
+    @Override
+    public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+        return this.attemptRedirect(context, next, context.getHttpRequest(), 1, new HashSet<>());
+    }
+
+    /**
+     * Function to process through the HTTP Response received in the pipeline
+     * and redirect sending the request with new redirect url.
+     */
+    private Mono<HttpResponse> attemptRedirect(HttpPipelineCallContext context, HttpPipelineNextPolicy next, HttpRequest originalHttpRequest, int redirectAttempt, Set<String> attemptedRedirectUrls) {
+        context.setHttpRequest(originalHttpRequest.copy());
+        return next.clone().process().flatMap((httpResponse) -> {
+            if (this.shouldAttemptRedirect(context, httpResponse, redirectAttempt, attemptedRedirectUrls)) {
+                HttpRequest redirectRequestCopy = this.createRedirectRequest(httpResponse);
+                return httpResponse.getBody().ignoreElements()
+                    .then(this.attemptRedirect(context, next, redirectRequestCopy, redirectAttempt + 1, attemptedRedirectUrls))
+                    .flatMap(newResponse -> {
+                        String digest = httpResponse.getHeaders().getValue(DOCKER_DIGEST_HEADER_NAME);
+                        if (digest != null) {
+                            newResponse.getHeaders().add(DOCKER_DIGEST_HEADER_NAME, digest);
+                        }
+                        return Mono.just(newResponse);
+                    });
+            } else {
+                return Mono.just(httpResponse);
+            }
+        });
+    }
+
+    public boolean shouldAttemptRedirect(HttpPipelineCallContext context, HttpResponse httpResponse, int tryCount, Set<String> attemptedRedirectUrls) {
+        if (this.isValidRedirectStatusCode(httpResponse.getStatusCode()) && this.isValidRedirectCount(tryCount) && this.isAllowedRedirectMethod(httpResponse.getRequest().getHttpMethod())) {
+            String redirectUrl = this.tryGetRedirectHeader(httpResponse.getHeaders(), REDIRECT_LOCATION_HEADER_NAME);
+            if (redirectUrl != null && !this.alreadyAttemptedRedirectUrl(redirectUrl, attemptedRedirectUrls)) {
+                LOGGER.verbose("[Redirecting] Try count: {}, Attempted Redirect URLs: {}", tryCount, String.join(",", attemptedRedirectUrls));
+                attemptedRedirectUrls.add(redirectUrl);
+                return true;
+            } else {
+                return false;
+            }
+        } else {
+            return false;
+        }
+    }
+
+    private HttpRequest createRedirectRequest(HttpResponse httpResponse) {
+        String responseLocation = this.tryGetRedirectHeader(httpResponse.getHeaders(), REDIRECT_LOCATION_HEADER_NAME);
+        HttpRequest request = httpResponse.getRequest();
+        request.setUrl(responseLocation);
+        request.getHeaders().remove(AUTHORIZATION);
+        return httpResponse.getRequest().setUrl(responseLocation);
+    }
+
+    private boolean alreadyAttemptedRedirectUrl(String redirectUrl, Set<String> attemptedRedirectUrls) {
+        if (attemptedRedirectUrls.contains(redirectUrl)) {
+            LOGGER.error("Request was redirected more than once to: {}", new Object[]{redirectUrl});
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    private boolean isValidRedirectCount(int tryCount) {
+        if (tryCount >= MAX_REDIRECT_ATTEMPTS) {
+            LOGGER.error("Request has been redirected more than {} times.", new Object[]{MAX_REDIRECT_ATTEMPTS});
+            return false;
+        } else {
+            return true;
+        }
+    }
+
+    private boolean isAllowedRedirectMethod(HttpMethod httpMethod) {
+        if (REDIRECT_ALLOWED_METHODS.contains(httpMethod)) {
+            return true;
+        } else {
+            LOGGER.error("Request was redirected from an invalid redirect allowed method: {}", new Object[]{httpMethod});
+            return false;
+        }
+    }
+
+    private boolean isValidRedirectStatusCode(int statusCode) {
+        return statusCode == PERMANENT_REDIRECT_STATUS_CODE || statusCode == TEMPORARY_REDIRECT_STATUS_CODE;
+    }
+
+    String tryGetRedirectHeader(HttpHeaders headers, String headerName) {
+        String headerValue = headers.getValue(headerName);
+        if (CoreUtils.isNullOrEmpty(headerValue)) {
+            LOGGER.error("Redirect url was null for header name: {}, request redirect was terminated.", headerName);
+            return null;
+        } else {
+            return headerValue;
+        }
+    }
+}
+

sdk/containerregistry/azure-containers-containerregistry/src/main/java/com/azure/containers/containerregistry/implementation/UtilsImpl.java

@@ -62,11 +62,11 @@ public final class UtilsImpl {
     private static final int HTTP_STATUS_CODE_ACCEPTED;
     private static final String CONTINUATION_LINK_HEADER_NAME;
     private static final Pattern CONTINUATION_LINK_PATTERN;
+    private static final ClientLogger LOGGER;
 
-    public static final String OCI_MANIFEST_MEDIA_TYPE;
     public static final String DOCKER_DIGEST_HEADER_NAME;
+    public static final String OCI_MANIFEST_MEDIA_TYPE;
     public static final String CONTAINER_REGISTRY_TRACING_NAMESPACE_VALUE;
-    private static final ClientLogger LOGGER;
 
     static {
         LOGGER = new ClientLogger(UtilsImpl.class);
@@ -76,7 +76,7 @@ public final class UtilsImpl {
         HTTP_STATUS_CODE_NOT_FOUND = 404;
         HTTP_STATUS_CODE_ACCEPTED = 202;
         OCI_MANIFEST_MEDIA_TYPE = "application/vnd.oci.image.manifest.v1+json";
-        DOCKER_DIGEST_HEADER_NAME = "Docker-Content-Digest";
+        DOCKER_DIGEST_HEADER_NAME = "docker-content-digest";
         CONTINUATION_LINK_HEADER_NAME = "Link";
         CONTINUATION_LINK_PATTERN = Pattern.compile("<(.+)>;.*");
         CONTAINER_REGISTRY_TRACING_NAMESPACE_VALUE = "Microsoft.ContainerRegistry";
@@ -126,6 +126,7 @@ public static HttpPipeline buildHttpPipeline(
         policies.add(ClientBuilderUtil.validateAndGetRetryPolicy(retryPolicy, retryOptions));
         policies.add(new CookiePolicy());
         policies.add(new AddDatePolicy());
+        policies.add(new ContainerRegistryRedirectPolicy());
 
         policies.addAll(perRetryPolicies);
         HttpPolicyProviders.addAfterRetryPolicies(policies);
@@ -222,7 +223,7 @@ public static <T> Mono<Response<Void>> deleteResponseToSuccess(Response<T> respo
             return getAcceptedDeleteResponse(responseT, responseT.getStatusCode());
         }
 
-        // In case of 400, we still convert it to success i.e. no-op.
+        // In case of 404, we still convert it to success i.e. no-op.
         return getAcceptedDeleteResponse(responseT, HTTP_STATUS_CODE_ACCEPTED);
     }
 
@@ -335,4 +336,13 @@ public static <T, R> PagedResponse<T> getPagedResponseWithContinuationToken(Page
             null
         );
     }
+
+    /**
+     * Get the digest from the response header if available.
+     * @param response The HttpResponse to parse.
+     * @return The digest value.
+     */
+    public static <T> String getDigestFromHeader(Response<T> response) {
+        return response.getHeaders().getValue(DOCKER_DIGEST_HEADER_NAME);
+    }
 }

sdk/containerregistry/azure-containers-containerregistry/src/main/java/com/azure/containers/containerregistry/implementation/authentication/ContainerRegistryCredentialsPolicy.java

@@ -5,6 +5,7 @@
 
 import com.azure.core.credential.TokenRequestContext;
 import com.azure.core.http.HttpPipelineCallContext;
+import com.azure.core.http.HttpPipelineNextPolicy;
 import com.azure.core.http.HttpResponse;
 import com.azure.core.http.policy.BearerTokenAuthenticationPolicy;
 import com.azure.core.util.logging.ClientLogger;
@@ -84,6 +85,45 @@ public Mono<Void> setAuthorizationHeader(HttpPipelineCallContext context, TokenR
             });
     }
 
+    @Override
+    public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+        if ("http".equals(context.getHttpRequest().getUrl().getProtocol())) {
+            return Mono.error(new RuntimeException("token credentials require a URL using the HTTPS protocol scheme"));
+        }
+
+        // Since we will need to replay this call, adding duplicate to make this replayable.
+        if (context.getHttpRequest().getBody() != null) {
+            context.getHttpRequest().setBody(context.getHttpRequest().getBody().map(buffer -> buffer.duplicate()));
+        }
+
+        HttpPipelineNextPolicy nextPolicy = next.clone();
+        return authorizeRequest(context)
+            .then(Mono.defer(() -> next.process()))
+            .flatMap(httpResponse -> {
+                String authHeader = httpResponse.getHeaderValue(WWW_AUTHENTICATE);
+                if (httpResponse.getStatusCode() == 401 && authHeader != null) {
+                    return authorizeRequestOnChallenge(context, httpResponse).flatMap(retry -> {
+                        if (retry) {
+                            return nextPolicy.process()
+                                .doFinally(ignored -> {
+                                    // Both Netty and OkHttp expect the requestBody to be closed after the connection is closed.
+                                    // Failure to do so results in memory leak.
+                                    // In case of StreamResponse (or other scenarios where we do not eagerly read the response)
+                                    // we let the client close the connection after the stream read.
+                                    // This can cause potential leaks in the scenarios like above, where the policy
+                                    // may intercept the response and prevent it from reaching the client.
+                                    // Hence, the policy needs to ensure that the connection is closed.
+                                    httpResponse.close();
+                                });
+                        } else {
+                            return Mono.just(httpResponse);
+                        }
+                    });
+                }
+                return Mono.just(httpResponse);
+            });
+    }
+
     /**
      * Handles the authentication challenge in the event a 401 response with a WWW-Authenticate authentication
      * challenge header is received after the initial request and returns appropriate {@link TokenRequestContext} to