AzConfig: User provided Customized Headers Policy

sdk/appconfiguration/azure-data-appconfiguration/src/main/java/com/azure/data/appconfiguration/ConfigurationClientBuilder.java

@@ -36,6 +36,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Supplier;
 
 /**
  * This class provides a fluent builder API to help aid the configuration and instantiation of
@@ -97,6 +98,7 @@ public final class ConfigurationClientBuilder {
     private HttpPipelinePolicy retryPolicy;
     private Configuration configuration;
     private ConfigurationServiceVersion version;
+    private Supplier<HttpHeaders> headersSupplier;
 
     /**
      * The constructor with defaults.
@@ -377,6 +379,18 @@ public ConfigurationClientBuilder serviceVersion(ConfigurationServiceVersion ver
         return this;
     }
 
+    /**
+     * Sets the customized headers that is used when each request is sent. It is possible to use new header value for
+     * each request.
+     *
+     * @param headersSupplier user's headers supplier applied to each request.
+     * @return The updated ConfigurationClientBuilder object.
+     */
+    public ConfigurationClientBuilder headersSupplier(Supplier<HttpHeaders> headersSupplier) {
+        this.headersSupplier = headersSupplier;
+        return this;
+    }
+
     private String getBuildEndpoint() {
         if (endpoint != null) {
             return endpoint;

sdk/appconfiguration/azure-data-appconfiguration/src/main/java/com/azure/data/appconfiguration/implementation/CustomHeadersPolicy.java

@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.data.appconfiguration.implementation;
+
+import com.azure.core.http.HttpHeader;
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpPipelineCallContext;
+import com.azure.core.http.HttpPipelineNextPolicy;
+import com.azure.core.http.HttpResponse;
+import com.azure.core.http.policy.HttpPipelinePolicy;
+import reactor.core.publisher.Mono;
+
+import java.util.Objects;
+import java.util.function.Supplier;
+
+public class CustomHeadersPolicy implements HttpPipelinePolicy {
+    private final Supplier<HttpHeaders> customHeaderSupplier;
+
+    public CustomHeadersPolicy(Supplier<HttpHeaders> customHeaderSupplier) {
+        this.customHeaderSupplier = Objects.requireNonNull(customHeaderSupplier, "Custom header cannot be null");
+    }
+
+    /**
+     * Adds the customized headers to authenticate a request to Azure App Configuration service.
+     *
+     * @param context The request context
+     * @param next The next HTTP pipeline policy to process the {@code context's} request after this policy
+     *     completes.
+     * @return A {@link Mono} representing the HTTP response that will arrive asynchronously.
+     */
+    @Override
+    public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+        final HttpHeaders httpHeaders = customHeaderSupplier.get();
+
+        if (Objects.nonNull(httpHeaders) && httpHeaders.getSize() > 0) {
+            for (HttpHeader header : httpHeaders) {
+                context.getHttpRequest().getHeaders().put(header.getName(), header.getValue());
+            }
+        }
+        return next.process();
+    }
+}

sdk/appconfiguration/azure-data-appconfiguration/src/test/java/com/azure/data/appconfiguration/CustomizedHeadersTest.java

@@ -0,0 +1,133 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.data.appconfiguration;
+
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpMethod;
+import com.azure.core.http.HttpPipeline;
+import com.azure.core.http.HttpPipelineBuilder;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpResponse;
+import com.azure.core.http.policy.FixedDelay;
+import com.azure.core.http.policy.RetryPolicy;
+import com.azure.core.test.http.NoOpHttpClient;
+import com.azure.data.appconfiguration.implementation.CustomHeadersPolicy;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+import java.net.URL;
+import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
+import java.time.Duration;
+import java.time.temporal.ChronoUnit;
+import java.util.UUID;
+import java.util.function.Supplier;
+
+public class CustomizedHeadersTest {
+
+    private static final String REQUEST_ID_HEADER = "x-ms-client-request-id";
+    private static final String CORRELATION_ID_HEADER = "x-ms-correlation-request-id";
+    private static final String CORRELATION_CONTEXT_HEADER = "correlation-context";
+
+    private final HttpResponse mockResponse = new HttpResponse(null) {
+        @Override
+        public int getStatusCode() {
+            return 500;
+        }
+        @Override
+        public String getHeaderValue(String name) {
+            return null;
+        }
+        @Override
+        public HttpHeaders getHeaders() {
+            return new HttpHeaders();
+        }
+        @Override
+        public Mono<byte[]> getBodyAsByteArray() {
+            return Mono.empty();
+        }
+        @Override
+        public Flux<ByteBuffer> getBody() {
+            return Flux.empty();
+        }
+        @Override
+        public Mono<String> getBodyAsString() {
+            return Mono.empty();
+        }
+        @Override
+        public Mono<String> getBodyAsString(Charset charset) {
+            return Mono.empty();
+        }
+    };
+
+
+    @Test
+    public void newRequestIdForEachCall() throws Exception {
+
+        String customRequestIdHeaderName = "x-ms-client-custom-request-id";
+        String clientProvidedRequestId = UUID.randomUUID().toString();
+
+        Supplier<HttpHeaders> requestIdSupplier = () ->
+            new HttpHeaders().put(customRequestIdHeaderName, clientProvidedRequestId);
+
+        HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(new NoOpHttpClient() {
+                String firstRequestId = null;
+                @Override
+                public Mono<HttpResponse> send(HttpRequest request) {
+                    if (firstRequestId != null) {
+                        String newRequestId = request.getHeaders().getValue(REQUEST_ID_HEADER);
+                        Assertions.assertNotNull(newRequestId);
+                        Assertions.assertNotEquals(newRequestId, firstRequestId);
+                    }
+                    firstRequestId = request.getHeaders().getValue(REQUEST_ID_HEADER);
+                    if (firstRequestId == null) {
+                        Assertions.fail();
+                    }
+                    return Mono.just(mockResponse);
+                }
+            })
+            .policies(new CustomHeadersPolicy(requestIdSupplier))
+            .build();
+        pipeline.send(new HttpRequest(HttpMethod.GET, new URL("http://localhost/"))).block();
+        pipeline.send(new HttpRequest(HttpMethod.GET, new URL("http://localhost/"))).block();
+    }
+
+
+    @Test
+    public void clientProvidedRequestIdForRetry() throws Exception {
+
+        String customRequestIdHeaderName = "x-ms-client-custom-request-id";
+        String clientProvidedRequestId = UUID.randomUUID().toString();
+
+        Supplier<HttpHeaders> requestIdSupplier = () ->
+            new HttpHeaders().put(customRequestIdHeaderName, clientProvidedRequestId);
+        final HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(new NoOpHttpClient() {
+                String firstRequestId = null;
+
+                @Override
+                public Mono<HttpResponse> send(HttpRequest request) {
+                    if (firstRequestId != null) {
+                        String newRequestId = request.getHeaders().getValue(customRequestIdHeaderName);
+                        Assertions.assertNotNull(newRequestId);
+                        Assertions.assertEquals(newRequestId, firstRequestId);
+                        Assertions.assertEquals(newRequestId, clientProvidedRequestId);
+                    }
+                    firstRequestId = request.getHeaders().getValue(customRequestIdHeaderName);
+                    if (firstRequestId == null) {
+                        Assertions.fail();
+                    }
+                    return Mono.just(mockResponse);
+                }
+            })
+            .policies(new CustomHeadersPolicy(requestIdSupplier),
+                new RetryPolicy(new FixedDelay(1, Duration.of(0, ChronoUnit.SECONDS))))
+            .build();
+
+        pipeline.send(new HttpRequest(HttpMethod.GET, new URL("http://localhost/"))).block();
+    }
+}