Adding sas token compatibility with nio file copy

sdk/storage/azure-storage-blob-nio/CHANGELOG.md

@@ -7,6 +7,7 @@
 ### Breaking Changes
 
 ### Bugs Fixed
+- Fixed an issue where the copy file operation would fail when using SAS token credentials.
 
 ### Other Changes
 

sdk/storage/azure-storage-blob-nio/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "java",
   "TagPrefix": "java/storage/azure-storage-blob-nio",
-  "Tag": "java/storage/azure-storage-blob-nio_263fe946de"
+  "Tag": "java/storage/azure-storage-blob-nio_dd63630d51"
 }

sdk/storage/azure-storage-blob-nio/src/main/java/com/azure/storage/blob/nio/AzureFileSystem.java

@@ -157,8 +157,8 @@ public final class AzureFileSystem extends FileSystem {
     private final Long blockSize;
     private final Long putBlobThreshold;
     private final Integer maxConcurrencyPerRequest;
-    private final Integer downloadResumeRetries;
     private final Map<String, FileStore> fileStores;
+    private final AzureSasCredential sasCredential;
     private FileStore defaultFileStore;
     private boolean closed;
 
@@ -177,7 +177,7 @@ public final class AzureFileSystem extends FileSystem {
             this.blockSize = (Long) config.get(AZURE_STORAGE_UPLOAD_BLOCK_SIZE);
             this.putBlobThreshold = (Long) config.get(AZURE_STORAGE_PUT_BLOB_THRESHOLD);
             this.maxConcurrencyPerRequest = (Integer) config.get(AZURE_STORAGE_MAX_CONCURRENCY_PER_REQUEST);
-            this.downloadResumeRetries = (Integer) config.get(AZURE_STORAGE_DOWNLOAD_RESUME_RETRIES);
+            this.sasCredential = (AzureSasCredential) config.get(AZURE_STORAGE_SAS_TOKEN_CREDENTIAL);
 
             // Initialize and ensure access to FileStores.
             this.fileStores = this.initializeFileStores(config);
@@ -491,4 +491,8 @@ Long getPutBlobThreshold() {
     Integer getMaxConcurrencyPerRequest() {
         return this.maxConcurrencyPerRequest;
     }
+
+    AzureSasCredential getSasCredential() {
+        return this.sasCredential;
+    }
 }

sdk/storage/azure-storage-blob-nio/src/main/java/com/azure/storage/blob/nio/AzureFileSystemProvider.java

@@ -744,6 +744,12 @@ public void copy(Path source, Path destination, CopyOption... copyOptions) throw
                     + "therefore invalid. Destination: " + destinationRes.getPath()));
         }
 
+        // Apply sas token if present
+        AzureFileSystem fileSystem = (AzureFileSystem) source.getFileSystem();
+        String sasToken
+            = fileSystem.getSasCredential() != null ? "?" + fileSystem.getSasCredential().getSignature() : "";
+        String sourceUrl = sourceRes.getBlobClient().getBlobUrl() + sasToken;
+
         /*
         Try to copy the resource at the source path.
 
@@ -754,8 +760,8 @@ public void copy(Path source, Path destination, CopyOption... copyOptions) throw
         first and then do a copy or createDir, which would always be two requests for all resource types.
          */
         try {
-            SyncPoller<BlobCopyInfo, Void> pollResponse = destinationRes.getBlobClient()
-                .beginCopy(sourceRes.getBlobClient().getBlobUrl(), null, null, null, null, requestConditions, null);
+            SyncPoller<BlobCopyInfo, Void> pollResponse
+                = destinationRes.getBlobClient().beginCopy(sourceUrl, null, null, null, null, requestConditions, null);
             pollResponse.waitForCompletion(Duration.ofSeconds(COPY_TIMEOUT_SECONDS));
         } catch (BlobStorageException e) {
             // If the source was not found, it could be because it's a virtual directory. Check the status.

sdk/storage/azure-storage-blob-nio/src/test/java/com/azure/storage/blob/nio/AzureFileSystemProviderTests.java

@@ -3,6 +3,7 @@
 
 package com.azure.storage.blob.nio;
 
+import com.azure.core.credential.AzureSasCredential;
 import com.azure.core.http.HttpHeaders;
 import com.azure.core.http.HttpMethod;
 import com.azure.core.http.HttpPipelineCallContext;
@@ -21,6 +22,10 @@
 import com.azure.storage.blob.specialized.AppendBlobClient;
 import com.azure.storage.blob.specialized.BlockBlobClient;
 import com.azure.storage.common.StorageSharedKeyCredential;
+import com.azure.storage.common.sas.AccountSasPermission;
+import com.azure.storage.common.sas.AccountSasResourceType;
+import com.azure.storage.common.sas.AccountSasService;
+import com.azure.storage.common.sas.AccountSasSignatureValues;
 import com.azure.storage.common.test.shared.extensions.LiveOnly;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -354,6 +359,27 @@ public void copySource(boolean sourceIsDir, boolean sourceIsVirtual, boolean sou
         }
     }
 
+    @Test
+    public void copySourceWithSas() throws IOException {
+        AzureSasCredential cred = new AzureSasCredential(primaryBlobServiceClient
+            .generateAccountSas(new AccountSasSignatureValues(testResourceNamer.now().plusDays(2),
+                AccountSasPermission.parse("rwcdl"), new AccountSasService().setBlobAccess(true),
+                new AccountSasResourceType().setContainer(true).setObject(true))));
+
+        config.put(AzureFileSystem.AZURE_STORAGE_SAS_TOKEN_CREDENTIAL, cred);
+        config.put(AzureFileSystem.AZURE_STORAGE_FILE_STORES, generateContainerName() + "," + generateContainerName());
+
+        AzureFileSystem fs
+            = new AzureFileSystem(new AzureFileSystemProvider(), ENV.getPrimaryAccount().getBlobEndpoint(), config);
+        basicSetupForCopyTest(fs);
+
+        fs.provider().createDirectory(sourcePath);
+        fs.provider().copy(sourcePath, destPath, StandardCopyOption.COPY_ATTRIBUTES);
+
+        assertTrue(destinationClient.exists());
+        checkBlobIsDir(destinationClient);
+    }
+
     @ParameterizedTest
     @CsvSource(value = { "false,false", "true,false", "true,true" })
     public void copyDestination(boolean destinationExists, boolean destinationIsDir) throws IOException {