Support Prefixes in AzureKeyCredentialPolicy (Azure#35010)

Support Prefixes in AzureKeyCredentialPolicy

Author: alzimmermsft

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/AzureKeyCredentialPolicy.java

@@ -5,10 +5,12 @@
 
 import com.azure.core.credential.AzureKeyCredential;
 import com.azure.core.http.HttpHeaderName;
+import com.azure.core.http.HttpHeaders;
 import com.azure.core.http.HttpPipelineCallContext;
 import com.azure.core.http.HttpPipelineNextPolicy;
 import com.azure.core.http.HttpPipelineNextSyncPolicy;
 import com.azure.core.http.HttpResponse;
+import com.azure.core.util.FluxUtil;
 import com.azure.core.util.logging.ClientLogger;
 import reactor.core.publisher.Mono;
 
@@ -25,18 +27,7 @@ public final class AzureKeyCredentialPolicy implements HttpPipelinePolicy {
     private static final ClientLogger LOGGER = new ClientLogger(AzureKeyCredentialPolicy.class);
     private final HttpHeaderName name;
     private final AzureKeyCredential credential;
-
-    private final HttpPipelineSyncPolicy inner = new HttpPipelineSyncPolicy() {
-        @Override
-        protected void beforeSendingRequest(HttpPipelineCallContext context) {
-            if ("http".equals(context.getHttpRequest().getUrl().getProtocol())) {
-                throw LOGGER.logExceptionAsError(
-                    new IllegalStateException("Key credentials require HTTPS to prevent leaking the key."));
-            }
-
-            context.getHttpRequest().setHeader(name, credential.getKey());
-        }
-    };
+    private final String prefix;
 
     /**
      * Creates a policy that uses the passed {@link AzureKeyCredential} to set the specified header name.
@@ -47,23 +38,64 @@ protected void beforeSendingRequest(HttpPipelineCallContext context) {
      * @throws IllegalArgumentException If {@code name} is empty.
      */
     public AzureKeyCredentialPolicy(String name, AzureKeyCredential credential) {
-        Objects.requireNonNull(credential, "'credential' cannot be null.");
+        this(name, credential, null);
+    }
+
+    /**
+     * Creates a policy that uses the passed {@link AzureKeyCredential} to set the specified header name.
+     * <p>
+     * The {@code prefix} will be applied before the {@link AzureKeyCredential#getKey()} when setting the header. A
+     * space will be inserted between {@code prefix} and credential.
+     *
+     * @param name The name of the key header that will be set to {@link AzureKeyCredential#getKey()}.
+     * @param credential The {@link AzureKeyCredential} containing the authorization key to use.
+     * @param prefix The prefix to apply before the credential, for example "SharedAccessKey credential".
+     * @throws NullPointerException If {@code name} or {@code credential} is {@code null}.
+     * @throws IllegalArgumentException If {@code name} is empty.
+     */
+    public AzureKeyCredentialPolicy(String name, AzureKeyCredential credential, String prefix) {
+        this(validateName(name), Objects.requireNonNull(credential, "'credential' cannot be null."), prefix);
+    }
+
+    private static HttpHeaderName validateName(String name) {
         Objects.requireNonNull(name, "'name' cannot be null.");
         if (name.isEmpty()) {
             throw LOGGER.logExceptionAsError(new IllegalArgumentException("'name' cannot be empty."));
         }
 
-        this.name = HttpHeaderName.fromString(name);
+        return HttpHeaderName.fromString(name);
+    }
+
+    AzureKeyCredentialPolicy(HttpHeaderName name, AzureKeyCredential credential, String prefix) {
+        this.name = name;
         this.credential = credential;
+        this.prefix = prefix != null ? prefix.trim() : null;
     }
 
     @Override
     public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
-        return inner.process(context, next);
+        if ("http".equals(context.getHttpRequest().getUrl().getProtocol())) {
+            return FluxUtil.monoError(LOGGER,
+                new IllegalStateException("Key credentials require HTTPS to prevent leaking the key."));
+        }
+
+        setCredential(context.getHttpRequest().getHeaders());
+        return next.process();
     }
 
     @Override
     public HttpResponse processSync(HttpPipelineCallContext context, HttpPipelineNextSyncPolicy next) {
-        return inner.processSync(context, next);
+        if ("http".equals(context.getHttpRequest().getUrl().getProtocol())) {
+            throw LOGGER.logExceptionAsError(
+                new IllegalStateException("Key credentials require HTTPS to prevent leaking the key."));
+        }
+
+        setCredential(context.getHttpRequest().getHeaders());
+        return next.processSync();
+    }
+
+    void setCredential(HttpHeaders headers) {
+        String credential = this.credential.getKey();
+        headers.set(name, (prefix == null) ? credential : prefix + " " + credential);
     }
 }

sdk/core/azure-core/src/test/java/com/azure/core/http/policy/AzureKeyCredentialPolicyTests.java

@@ -0,0 +1,36 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.core.http.policy;
+
+import com.azure.core.credential.AzureKeyCredential;
+import com.azure.core.http.HttpHeaderName;
+import com.azure.core.http.HttpHeaders;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import java.util.stream.Stream;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class AzureKeyCredentialPolicyTests {
+    @ParameterizedTest
+    @MethodSource("setCredentialSupplier")
+    public void setCredential(AzureKeyCredentialPolicy policy, String expectedHeader) {
+        HttpHeaders headers = new HttpHeaders();
+        policy.setCredential(headers);
+        assertEquals(expectedHeader, headers.getValue(HttpHeaderName.AUTHORIZATION));
+    }
+
+    private static Stream<Arguments> setCredentialSupplier() {
+        AzureKeyCredential credential = new AzureKeyCredential("asecret");
+        return Stream.of(
+            Arguments.of(new AzureKeyCredentialPolicy(HttpHeaderName.AUTHORIZATION, credential, null), "asecret"),
+            Arguments.of(new AzureKeyCredentialPolicy(HttpHeaderName.AUTHORIZATION, credential, "SharedKeyCredential"),
+                "SharedKeyCredential asecret"),
+            Arguments.of(new AzureKeyCredentialPolicy(HttpHeaderName.AUTHORIZATION, credential, "SharedKeyCredential "),
+                "SharedKeyCredential asecret")
+        );
+    }
+}