Communication: Add TokenCredentialAddHostHeaderPolicy for TokenCreden…

…tial Requests (#24442)

* Communication: Add TokenCredentialAdditionalHeaderPolicy for CallingServerClientBuilder

* Add TokenCredentialAddHostHeaderPolicyTests

* Fixing comment

* Use URL class to get hostname

* Fix style errors

Co-authored-by: Melissa Neubert <[email protected]>

sdk/communication/azure-communication-callingserver/src/main/java/com/azure/communication/callingserver/CallingServerClientBuilder.java

@@ -27,12 +27,13 @@
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.logging.ClientLogger;
 
+import java.net.MalformedURLException;
+import java.net.URL;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
-
 /**
  * Client builder that creates CallingServerAsyncClient and CallingServerClient.
  *
@@ -49,6 +50,7 @@ public final class CallingServerClientBuilder {
     private final ClientLogger logger = new ClientLogger(CallingServerClientBuilder.class);
     private String connectionString;
     private String endpoint;
+    private String hostName;
     private AzureKeyCredential azureKeyCredential;
     private TokenCredential tokenCredential;
     private HttpClient httpClient;
@@ -252,6 +254,14 @@ private AzureCommunicationCallingServerServiceImpl createServiceImpl() {
         }
 
         Objects.requireNonNull(endpoint);
+        if (isTokenCredentialSet) {
+            try {
+                hostName = getHostNameFromEndpoint();
+            } catch (MalformedURLException e) {
+                throw logger.logExceptionAsError(new RuntimeException(e.getMessage()));
+            }
+        }
+
         if (pipeline == null) {
             Objects.requireNonNull(httpClient);
         }
@@ -279,20 +289,25 @@ public CallingServerClientBuilder clientOptions(ClientOptions clientOptions) {
         return this;
     }
 
-    private HttpPipelinePolicy createHttpPipelineAuthPolicy() {
+    private List<HttpPipelinePolicy> createHttpPipelineAuthPolicies() {
         if (tokenCredential != null && azureKeyCredential != null) {
             throw logger.logExceptionAsError(new IllegalArgumentException(
                 "Both 'credential' and 'keyCredential' are set. Just one may be used."));
         }
+
+        List<HttpPipelinePolicy> pipelinePolicies = new ArrayList<>();
         if (tokenCredential != null) {
-            return new BearerTokenAuthenticationPolicy(tokenCredential,
-                "https://communication.azure.com//.default");
+            pipelinePolicies.add(new BearerTokenAuthenticationPolicy(tokenCredential,
+                "https://communication.azure.com//.default"));
+            pipelinePolicies.add(new TokenCredentialAddHostHeaderPolicy(hostName));
         } else if (azureKeyCredential != null) {
-            return new HmacAuthenticationPolicy(azureKeyCredential);
+            pipelinePolicies.add(new HmacAuthenticationPolicy(azureKeyCredential));
         } else {
             throw logger.logExceptionAsError(
                 new IllegalArgumentException("Missing credential information while building a client."));
         }
+
+        return pipelinePolicies;
     }
 
     private HttpPipeline createHttpPipeline(HttpClient httpClient) {
@@ -319,7 +334,7 @@ private HttpPipeline createHttpPipeline(HttpClient httpClient) {
         policyList.add(new RequestIdPolicy());
         policyList.add((retryPolicy == null) ? new RetryPolicy() : retryPolicy);
         policyList.add(new RedirectPolicy());
-        policyList.add(createHttpPipelineAuthPolicy());
+        policyList.addAll(createHttpPipelineAuthPolicies());
         policyList.add(new CookiePolicy());
 
         // Add additional policies
@@ -341,4 +356,8 @@ private HttpLogOptions getHttpLogOptions() {
 
         return httpLogOptions;
     }
+
+    private String getHostNameFromEndpoint() throws MalformedURLException {
+        return new URL(endpoint).getHost();
+    }
 }

sdk/communication/azure-communication-callingserver/src/main/java/com/azure/communication/callingserver/TokenCredentialAddHostHeaderPolicy.java

@@ -0,0 +1,38 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.communication.callingserver;
+
+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;
+
+/**
+ * HttpPipelinePolicy to append request host headers for CallingServerClient TokenCredential requests
+ */
+public final class TokenCredentialAddHostHeaderPolicy implements HttpPipelinePolicy {
+    private static final String X_MS_HOST_HEADER = "x-ms-host";
+    private final String xMsHostValue;
+
+    /**
+     * Created with a non-null resourceHostName value
+     * @param resourceHostName Host name of the ACS resource
+     */
+    public TokenCredentialAddHostHeaderPolicy(String resourceHostName) {
+        Objects.requireNonNull(resourceHostName, "'resourceHostName' cannot be a null value.");
+        xMsHostValue = resourceHostName;
+    }
+
+    @Override
+    public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+        if (!xMsHostValue.isEmpty()) {
+            context.getHttpRequest().setHeader(X_MS_HOST_HEADER, xMsHostValue);
+        }
+        return next.process();
+    }
+}

sdk/communication/azure-communication-callingserver/src/test/java/com/azure/communication/callingserver/TokenCredentialAddHostHeaderPolicyTests.java

@@ -0,0 +1,110 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.communication.callingserver;
+
+import com.azure.core.http.HttpClient;
+import com.azure.core.http.HttpPipeline;
+import com.azure.core.http.HttpPipelineBuilder;
+import com.azure.core.http.HttpMethod;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpResponse;
+import com.azure.core.http.policy.HttpPipelinePolicy;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import org.junit.jupiter.api.Test;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.net.MalformedURLException;
+import java.net.URL;
+
+public class TokenCredentialAddHostHeaderPolicyTests {
+
+    private class NoOpHttpClient implements HttpClient {
+        @Override
+        public Mono<HttpResponse> send(HttpRequest request) {
+            return Mono.empty(); // NOP
+        }
+    }
+
+    private static final String HOST_NAME = "host.communication.azure.com";
+
+    private final HttpPipelinePolicy verifyHeadersPolicy = (context, next) -> {
+        HttpRequest request = context.getHttpRequest();
+        String hostHeaderValue = request.getHeaders().getValue("x-ms-host");
+        assertEquals(HOST_NAME, hostHeaderValue);
+        return next.process();
+    };
+
+    @Test
+    public void getRequestTest() throws MalformedURLException {
+        final TokenCredentialAddHostHeaderPolicy clientPolicy = new TokenCredentialAddHostHeaderPolicy(HOST_NAME);
+
+        final HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(new NoOpHttpClient())
+            .policies(clientPolicy, verifyHeadersPolicy)
+            .build();
+
+        HttpRequest request = new HttpRequest(HttpMethod.GET, new URL("https://localhost?id=b93a5ef4-f622-44d8-a80b-ff983122554e"));
+        StepVerifier.create(pipeline.send(request))
+            .verifyComplete();
+    }
+
+    @Test
+    public void postRequestTest() throws MalformedURLException {
+        final TokenCredentialAddHostHeaderPolicy clientPolicy = new TokenCredentialAddHostHeaderPolicy(HOST_NAME);
+
+        final HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(new NoOpHttpClient())
+            .policies(clientPolicy, verifyHeadersPolicy)
+            .build();
+
+        HttpRequest request = new HttpRequest(HttpMethod.POST, new URL("https://localhost?id=b93a5ef4-f622-44d8-a80b-ff983122554e"));
+        StepVerifier.create(pipeline.send(request))
+            .verifyComplete();
+    }
+
+    @Test
+    public void patchRequestTest() throws MalformedURLException {
+        final TokenCredentialAddHostHeaderPolicy clientPolicy = new TokenCredentialAddHostHeaderPolicy(HOST_NAME);
+
+        final HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(new NoOpHttpClient())
+            .policies(clientPolicy, verifyHeadersPolicy)
+            .build();
+
+        HttpRequest request = new HttpRequest(HttpMethod.PATCH, new URL("https://localhost?id=b93a5ef4-f622-44d8-a80b-ff983122554e"));
+        StepVerifier.create(pipeline.send(request))
+            .verifyComplete();
+    }
+
+
+    @Test
+    public void putRequestTest() throws MalformedURLException {
+        final TokenCredentialAddHostHeaderPolicy clientPolicy = new TokenCredentialAddHostHeaderPolicy(HOST_NAME);
+
+        final HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(new NoOpHttpClient())
+            .policies(clientPolicy, verifyHeadersPolicy)
+            .build();
+
+        HttpRequest request = new HttpRequest(HttpMethod.PUT, new URL("https://localhost?id=b93a5ef4-f622-44d8-a80b-ff983122554e"));
+        StepVerifier.create(pipeline.send(request))
+            .verifyComplete();
+    }
+
+    @Test
+    public void deleteRequestTest() throws MalformedURLException {
+        final TokenCredentialAddHostHeaderPolicy clientPolicy = new TokenCredentialAddHostHeaderPolicy(HOST_NAME);
+
+        final HttpPipeline pipeline = new HttpPipelineBuilder()
+            .httpClient(new NoOpHttpClient())
+            .policies(clientPolicy, verifyHeadersPolicy)
+            .build();
+
+        HttpRequest request = new HttpRequest(HttpMethod.DELETE, new URL("https://localhost?id=b93a5ef4-f622-44d8-a80b-ff983122554e"));
+        StepVerifier.create(pipeline.send(request))
+            .verifyComplete();
+    }
+}