From 00249619a0c7e8cb565158a867a8f86352085d92 Mon Sep 17 00:00:00 2001
From: sreeb-msft <111426823+sreeb-msft@users.noreply.github.com>
Date: Mon, 28 Nov 2022 17:08:13 +0530
Subject: [PATCH] HADOOP-18498. ABFS: Remove unwanted ? prefix from SAS Tokens
 (#5136)

This commit parses SAS Tokens and removes the unwanted prefix of '?' from them, if present.

At present, SAS Tokens are provided to the driver through customer implementations of the SASTokenProvider interface. The SAS token providers should not assume that the token will be the first query parameter in the URIs that communicate with the backend. However, it was observed that certain public interfaces provided by Storage to generate SAS can include the '?' as the first character of the SAS Token, which would ideally be the case when it is the first query parameter. Thus, tokens that contain this prefix will lead to an error in the driver due to a clash of query parameters.

To avoid failures for use of such SAS tokens, after receiving the SAS Token from the provider, the code checks for whether any ? prefix is present or not. If yes, it is removed before further usage of the token. This way, users would not have to manually remove the prefix before passing it on as a configuration.

Contributed by Sree Bhattacharya
---
 .../fs/azurebfs/services/AbfsClient.java      |  4 ++
 ...ITestAzureBlobFileSystemDelegationSAS.java | 13 ++++++
 .../MockWithPrefixSASTokenProvider.java       | 40 +++++++++++++++++++
 3 files changed, 57 insertions(+)
 create mode 100644 hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/extensions/MockWithPrefixSASTokenProvider.java

diff --git a/hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/services/AbfsClient.java b/hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/services/AbfsClient.java
index b903dbfd81a..fe2ea35f1df 100644
--- a/hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/services/AbfsClient.java
+++ b/hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/services/AbfsClient.java
@@ -1130,6 +1130,10 @@ public class AbfsClient implements Closeable {
           sasToken = cachedSasToken;
           LOG.trace("Using cached SAS token.");
         }
+        // if SAS Token contains a prefix of ?, it should be removed
+        if (sasToken.charAt(0) == '?') {
+          sasToken = sasToken.substring(1);
+        }
         queryBuilder.setSASToken(sasToken);
         LOG.trace("SAS token fetch complete for {} on {}", operation, path);
       } catch (Exception ex) {
diff --git a/hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/ITestAzureBlobFileSystemDelegationSAS.java b/hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/ITestAzureBlobFileSystemDelegationSAS.java
index edc3930607c..b164689ef80 100644
--- a/hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/ITestAzureBlobFileSystemDelegationSAS.java
+++ b/hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/ITestAzureBlobFileSystemDelegationSAS.java
@@ -479,4 +479,17 @@ public class ITestAzureBlobFileSystemDelegationSAS extends AbstractAbfsIntegrati
         "r--r-----",
         fileStatus.getPermission().toString());
   }
+
+  @Test
+  public void testSASQuesMarkPrefix() throws Exception {
+    AbfsConfiguration testConfig = this.getConfiguration();
+    // the SAS Token Provider is changed
+    testConfig.set(FS_AZURE_SAS_TOKEN_PROVIDER_TYPE, "org.apache.hadoop.fs.azurebfs.extensions.MockWithPrefixSASTokenProvider");
+
+    AzureBlobFileSystem testFs = (AzureBlobFileSystem) FileSystem.newInstance(getRawConfiguration());
+    Path testFile = new Path("/testSASPrefixQuesMark");
+
+    // the creation of this filesystem should work correctly even when a SAS Token is generated with a ? prefix
+    testFs.create(testFile).close();
+  }
 }
diff --git a/hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/extensions/MockWithPrefixSASTokenProvider.java b/hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/extensions/MockWithPrefixSASTokenProvider.java
new file mode 100644
index 00000000000..ed701c4669c
--- /dev/null
+++ b/hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/extensions/MockWithPrefixSASTokenProvider.java
@@ -0,0 +1,40 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.fs.azurebfs.extensions;
+
+import java.io.IOException;
+
+public class MockWithPrefixSASTokenProvider extends MockSASTokenProvider {
+
+    /**
+     * Function to return an already generated SAS Token with a '?' prefix
+     * @param accountName the name of the storage account.
+     * @param fileSystem the name of the fileSystem.
+     * @param path the file or directory path.
+     * @param operation the operation to be performed on the path.
+     * @return
+     * @throws IOException
+     */
+    @Override
+    public String getSASToken(String accountName, String fileSystem, String path,
+                              String operation) throws IOException {
+        String token = super.getSASToken(accountName, fileSystem, path, operation);
+        return "?" + token;
+    }
+}
\ No newline at end of file