[Service Bus] Enable SAS support in connection string for Service Bus (…

…#14939)

* Enable SAS support in connection string for Service Bus

* Fix checkstyle

* Add logs

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusClientBuilder.java

@@ -106,8 +106,7 @@ public ServiceBusClientBuilder connectionString(String connectionString) {
         final ConnectionStringProperties properties = new ConnectionStringProperties(connectionString);
         final TokenCredential tokenCredential;
         try {
-            tokenCredential = new ServiceBusSharedKeyCredential(properties.getSharedAccessKeyName(),
-                properties.getSharedAccessKey(), ServiceBusConstants.TOKEN_VALIDITY);
+            tokenCredential = getTokenCredential(properties);
         } catch (Exception e) {
             throw logger.logExceptionAsError(
                 new AzureException("Could not create the ServiceBusSharedKeyCredential.", e));
@@ -123,6 +122,17 @@ public ServiceBusClientBuilder connectionString(String connectionString) {
         return credential(properties.getEndpoint().getHost(), tokenCredential);
     }
 
+    private TokenCredential getTokenCredential(ConnectionStringProperties properties) {
+        TokenCredential tokenCredential;
+        if (properties.getSharedAccessSignature() == null) {
+            tokenCredential = new ServiceBusSharedKeyCredential(properties.getSharedAccessKeyName(),
+                properties.getSharedAccessKey(), ServiceBusConstants.TOKEN_VALIDITY);
+        } else {
+            tokenCredential = new ServiceBusSharedKeyCredential(properties.getSharedAccessSignature());
+        }
+        return tokenCredential;
+    }
+
     /**
      * Sets the configuration store that is used during construction of the service client.
      *

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/ServiceBusSharedKeyCredential.java

@@ -19,8 +19,10 @@
 import java.security.InvalidKeyException;
 import java.security.NoSuchAlgorithmException;
 import java.time.Duration;
+import java.time.Instant;
 import java.time.OffsetDateTime;
 import java.time.ZoneOffset;
+import java.util.Arrays;
 import java.util.Base64;
 import java.util.Locale;
 import java.util.Objects;
@@ -52,6 +54,7 @@ public class ServiceBusSharedKeyCredential implements TokenCredential {
     private final String policyName;
     private final Mac hmac;
     private final Duration tokenValidity;
+    private final String sharedAccessSignature;
 
     /**
      * Creates an instance that authorizes using the {@code policyName} and {@code sharedAccessKey}.
@@ -112,6 +115,26 @@ public ServiceBusSharedKeyCredential(String policyName, String sharedAccessKey,
             throw logger.logExceptionAsError(new IllegalArgumentException(
                 "'sharedAccessKey' is an invalid value for the hashing algorithm.", e));
         }
+        this.sharedAccessSignature = null;
+    }
+
+    /**
+     * Creates an instance using the provided Shared Access Signature (SAS) string. The credential created using this
+     * constructor will not be refreshed. The expiration time is set to the time defined in "se={
+     * tokenValidationSeconds}`. If the SAS string does not contain this or is in invalid format, then the token
+     * expiration will be set to {@link OffsetDateTime#MAX max duration}.
+     * <p><a href="https://docs.microsoft.com/rest/api/eventhub/generate-sas-token">See how to generate SAS
+     * programmatically.</a></p>
+     *
+     * @param sharedAccessSignature The base64 encoded shared access signature string.
+     * @throws NullPointerException if {@code sharedAccessSignature} is null.
+     */
+    public ServiceBusSharedKeyCredential(String sharedAccessSignature) {
+        this.sharedAccessSignature = Objects.requireNonNull(sharedAccessSignature,
+            "'sharedAccessSignature' cannot be null");
+        this.policyName = null;
+        this.hmac = null;
+        this.tokenValidity = null;
     }
 
     /**
@@ -138,6 +161,10 @@ private AccessToken generateSharedAccessSignature(final String resource) throws
             throw logger.logExceptionAsError(new IllegalArgumentException("resource cannot be empty"));
         }
 
+        if (sharedAccessSignature != null) {
+            return new AccessToken(sharedAccessSignature, getExpirationTime(sharedAccessSignature));
+        }
+
         final String utf8Encoding = UTF_8.name();
         final OffsetDateTime expiresOn = OffsetDateTime.now(ZoneOffset.UTC).plus(tokenValidity);
         final String expiresOnEpochSeconds = Long.toString(expiresOn.toEpochSecond());
@@ -155,4 +182,24 @@ private AccessToken generateSharedAccessSignature(final String resource) throws
 
         return new AccessToken(token, expiresOn);
     }
+
+    private OffsetDateTime getExpirationTime(String sharedAccessSignature) {
+        String[] parts = sharedAccessSignature.split("&");
+        return Arrays.stream(parts)
+            .map(part -> part.split("="))
+            .filter(pair -> pair.length == 2 && pair[0].equalsIgnoreCase("se"))
+            .findFirst()
+            .map(pair -> pair[1])
+            .map(expirationTimeStr -> {
+                try {
+                    long epochSeconds = Long.parseLong(expirationTimeStr);
+                    return Instant.ofEpochSecond(epochSeconds).atOffset(ZoneOffset.UTC);
+                } catch (NumberFormatException exception) {
+                    logger.verbose("Invalid expiration time format in the SAS token: {}. Falling back to max "
+                        + "expiration time.", expirationTimeStr);
+                    return OffsetDateTime.MAX;
+                }
+            })
+            .orElse(OffsetDateTime.MAX);
+    }
 }

sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusClientBuilderTest.java

@@ -289,6 +289,22 @@ public void testProxyOptionsConfiguration(String proxyConfiguration, boolean exp
         Assertions.assertEquals(expectedClientCreation, clientCreated);
     }
 
+    @Test
+    public void testConnectionStringWithSas() {
+        String connectionStringWithEntityPath = "Endpoint=sb://sb-name.servicebus.windows.net/;"
+            + "SharedAccessSignature=SharedAccessSignature test-value;EntityPath=sb-name";
+        assertNotNull(new ServiceBusClientBuilder()
+            .connectionString(connectionStringWithEntityPath));
+
+        assertThrows(IllegalArgumentException.class,
+            () -> new ServiceBusClientBuilder()
+                .connectionString("SharedAccessSignature=SharedAccessSignature test-value;EntityPath=sb-name"));
+
+        assertThrows(IllegalArgumentException.class,
+            () -> new ServiceBusClientBuilder()
+                .connectionString("Endpoint=sb://sb-name.servicebus.windows.net/;EntityPath=sb-name"));
+    }
+
     private static Stream<Arguments> getProxyConfigurations() {
         return Stream.of(
             Arguments.of("http://localhost:8080", true),

sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/implementation/ServiceBusSharedKeyCredentialTest.java

@@ -0,0 +1,58 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.messaging.servicebus.implementation;
+
+import com.azure.core.credential.TokenRequestContext;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import reactor.test.StepVerifier;
+
+import java.time.OffsetDateTime;
+import java.util.stream.Stream;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+/**
+ * Unit tests for {@link ServiceBusSharedKeyCredential}.
+ */
+public class ServiceBusSharedKeyCredentialTest {
+
+    @ParameterizedTest
+    @MethodSource("getSas")
+    public void testSharedAccessSignatureCredential(String sas, OffsetDateTime expectedExpirationTime) {
+        ServiceBusSharedKeyCredential serviceBusSharedKeyCredential = new ServiceBusSharedKeyCredential(sas);
+        StepVerifier.create(serviceBusSharedKeyCredential.getToken(new TokenRequestContext().addScopes("sb://test"
+            + "-entity.servicebus.windows.net/.default")))
+            .assertNext(token -> {
+                assertNotNull(token.getToken());
+                assertEquals(sas, token.getToken());
+                assertEquals(expectedExpirationTime, token.getExpiresAt());
+            })
+            .verifyComplete();
+    }
+
+    private static Stream<Arguments> getSas() {
+        String validSas = "SharedAccessSignature "
+            + "sr=https%3A%2F%2Fentity-name.servicebus.windows.net%2F"
+            + "&sig=encodedsignature%3D"
+            + "&se=1599537084"
+            + "&skn=test-sas-key";
+        String validSasWithNoExpirationTime = "SharedAccessSignature "
+            + "sr=https%3A%2F%2Fentity-name.servicebus.windows.net%2F"
+            + "&sig=encodedsignature%3D"
+            + "&skn=test-sas-key";
+        String validSasInvalidExpirationTimeFormat = "SharedAccessSignature "
+            + "sr=https%3A%2F%2Fentity-name.servicebus.windows.net%2F"
+            + "&sig=encodedsignature%3D"
+            + "&se=se=2020-12-31T13:37:45Z"
+            + "&skn=test-sas-key";
+
+        return Stream.of(
+            Arguments.of(validSas, OffsetDateTime.parse("2020-09-08T03:51:24Z")),
+            Arguments.of(validSasWithNoExpirationTime, OffsetDateTime.MAX),
+            Arguments.of(validSasInvalidExpirationTimeFormat, OffsetDateTime.MAX)
+        );
+    }
+}