Fixed an issue in Transfer Manager downloadDirectory that could cause… (#3625)

* Fixed an issue in Transfer Manager downloadDirectory that could cause objects that contain prefix to fail to download

Author: zoewangg

services-custom/s3-transfer-manager/src/it/java/software/amazon/awssdk/transfer/s3/S3TransferManagerDownloadDirectoryIntegrationTest.java

@@ -125,11 +125,24 @@ public void downloadDirectory() throws Exception {
         assertTwoDirectoriesHaveSameStructure(sourceDirectory, directory);
     }
 
+    @ParameterizedTest
+    @ValueSource(strings = {"notes/2021", "notes/2021/"})
+    void downloadDirectory_withPrefix(String prefix) throws Exception {
+        DirectoryDownload downloadDirectory = tm.downloadDirectory(u -> u.destination(directory)
+                                                                         .listObjectsV2RequestTransformer(r -> r.prefix(prefix))
+                                                                         .bucket(TEST_BUCKET));
+        CompletedDirectoryDownload completedDirectoryDownload = downloadDirectory.completionFuture().get(5, TimeUnit.SECONDS);
+        assertThat(completedDirectoryDownload.failedTransfers()).isEmpty();
+
+        assertTwoDirectoriesHaveSameStructure(sourceDirectory.resolve(prefix), directory);
+    }
+
     /**
      * With prefix = "notes", the destination directory structure should be the following:
      * <pre>
      *   {@code
      *      - destination
+     *          - notesMemo.txt
      *          - 2021
      *              - 1.txt
      *              - 2.txt
@@ -139,16 +152,24 @@ public void downloadDirectory() throws Exception {
      *   }
      * </pre>
      */
-    @ParameterizedTest
-    @ValueSource(strings = {"notes/2021", "notes/2021/", "notes"})
-    public void downloadDirectory_withPrefix(String prefix) throws Exception {
+    @Test
+    void downloadDirectory_containsObjectWithPrefixInTheKey_shouldResolveCorrectly() throws Exception {
+        String prefix = "notes";
         DirectoryDownload downloadDirectory = tm.downloadDirectory(u -> u.destination(directory)
                                                                          .listObjectsV2RequestTransformer(r -> r.prefix(prefix))
                                                                          .bucket(TEST_BUCKET));
         CompletedDirectoryDownload completedDirectoryDownload = downloadDirectory.completionFuture().get(5, TimeUnit.SECONDS);
         assertThat(completedDirectoryDownload.failedTransfers()).isEmpty();
 
-        assertTwoDirectoriesHaveSameStructure(sourceDirectory.resolve(prefix), directory);
+        Path expectedDirectory = Files.createTempDirectory("expectedDirectory");
+
+        try {
+            FileUtils.copyDirectory(sourceDirectory.resolve(prefix), expectedDirectory);
+            Files.copy(sourceDirectory.resolve("notesMemo.txt"), expectedDirectory.resolve("notesMemo.txt"));
+            assertTwoDirectoriesHaveSameStructure(expectedDirectory, directory);
+        } finally {
+            FileUtils.cleanUpTestDirectory(expectedDirectory);
+        }
     }
 
     /**
@@ -203,6 +224,7 @@ public void downloadDirectory_withFilter() throws Exception {
             Files.delete(expectedDirectory.resolve("notes/2022"));
             Files.delete(expectedDirectory.resolve("notes/important.txt"));
             Files.delete(expectedDirectory.resolve("notes/2021/1.txt"));
+            Files.delete(expectedDirectory.resolve("notesMemo.txt"));
 
             assertTwoDirectoriesHaveSameStructure(expectedDirectory, directory);
         } finally {
@@ -242,6 +264,7 @@ private static void assertLeftHasRight(Path left, Path right) {
      *      - source
      *          - README.md
      *          - CHANGELOG.md
+     *          - notesMemo.txt
      *          - notes
      *              - 2021
      *                  - 1.txt
@@ -262,6 +285,7 @@ private static Path createLocalTestDirectory() throws IOException {
         Files.createDirectory(Paths.get(directoryName, "notes", "2022"));
         Files.write(Paths.get(directoryName, "README.md"), RandomStringUtils.random(100).getBytes(StandardCharsets.UTF_8));
         Files.write(Paths.get(directoryName, "CHANGELOG.md"), RandomStringUtils.random(100).getBytes(StandardCharsets.UTF_8));
+        Files.write(Paths.get(directoryName, "notesMemo.txt"), RandomStringUtils.random(100).getBytes(StandardCharsets.UTF_8));
         Files.write(Paths.get(directoryName, "notes", "2021", "1.txt"),
                     RandomStringUtils.random(100).getBytes(StandardCharsets.UTF_8));
         Files.write(Paths.get(directoryName, "notes", "2021", "2.txt"),

services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/DownloadDirectoryHelper.java

@@ -147,12 +147,10 @@ private Path determineDestinationPath(DownloadDirectoryRequest downloadDirectory
                                           ListObjectsV2Request listRequest,
                                           S3Object s3Object) {
         FileSystem fileSystem = downloadDirectoryRequest.destination().getFileSystem();
-        // listOBjectv2requests.delimiter
         String delimiter = listRequest.delimiter() == null ? DEFAULT_DELIMITER : listRequest.delimiter();
-        String key = normalizeKey(downloadDirectoryRequest, listRequest, s3Object, delimiter);
+        String key = normalizeKey(listRequest, s3Object.key(), delimiter);
         String relativePath = getRelativePath(fileSystem, delimiter, key);
         Path destinationPath = downloadDirectoryRequest.destination().resolve(relativePath);
-
         validatePath(downloadDirectoryRequest.destination(), destinationPath, s3Object.key());
         return destinationPath;
     }
@@ -200,31 +198,35 @@ private CompletableFuture<CompletedFileDownload> doDownloadSingleFile(DownloadDi
     }
 
     /**
-     * Normalizing the key by stripping the prefix from the s3 object key if the prefix is not empty.
+     * If the prefix is not empty AND the key contains the delimiter, normalize the key by stripping the prefix from the key.
      *
      * If a delimiter is null (not provided by user), use "/" by default.
      *
      * For example: given a request with prefix = "notes/2021"  or "notes/2021/", delimiter = "/" and key = "notes/2021/1.txt",
      * the normalized key should be "1.txt".
      */
-    private static String normalizeKey(DownloadDirectoryRequest downloadDirectoryRequest,
-                                       ListObjectsV2Request listObjectsRequest,
-                                       S3Object s3Object,
+    private static String normalizeKey(ListObjectsV2Request listObjectsRequest,
+                                       String key,
                                        String delimiter) {
-        int delimiterLength = delimiter.length();
+        if (StringUtils.isEmpty(listObjectsRequest.prefix())) {
+            return key;
+        }
 
-        if (!StringUtils.isEmpty(listObjectsRequest.prefix())) {
-            String prefix = listObjectsRequest.prefix();
-            String normalizedKey;
-            if (prefix.endsWith(delimiter)) {
-                normalizedKey = s3Object.key().substring(prefix.length());
-            } else {
-                normalizedKey = s3Object.key().substring(prefix.length() + delimiterLength);
-            }
-            return normalizedKey;
+        String prefix = listObjectsRequest.prefix();
+
+        if (!key.contains(delimiter)) {
+            return key;
+        }
+
+        String normalizedKey;
+
+        if (prefix.endsWith(delimiter)) {
+            normalizedKey = key.substring(prefix.length());
+        } else {
+            normalizedKey = key.substring(prefix.length() + delimiter.length());
         }
+        return normalizedKey;
 
-        return s3Object.key();
     }
 
     private static String getRelativePath(FileSystem fileSystem, String delimiter, String key) {

services-custom/s3-transfer-manager/src/test/java/software/amazon/awssdk/transfer/s3/internal/DownloadDirectoryHelperTest.java

@@ -366,6 +366,42 @@ void downloadDirectory_withPrefix_shouldStripPrefixInDestinationPath(FileSystem
         assertThat(destinations).isEqualTo(expectedPaths);
     }
 
+    @ParameterizedTest
+    @MethodSource("fileSystems")
+    void downloadDirectory_containsObjectWithPrefixInIt_shouldInclude(FileSystem jimfs) {
+        String prefix = "abc";
+        directory = jimfs.getPath("test");
+        String[] keys = {"abc/def/image.jpg", "abc/def/title.jpg", "abcd"};
+        stubSuccessfulListObjects(listObjectsHelper, keys);
+        ArgumentCaptor<DownloadFileRequest> requestArgumentCaptor = ArgumentCaptor.forClass(DownloadFileRequest.class);
+
+        when(singleDownloadFunction.apply(requestArgumentCaptor.capture()))
+            .thenReturn(completedDownload());
+        DirectoryDownload downloadDirectory =
+            downloadDirectoryHelper.downloadDirectory(DownloadDirectoryRequest.builder()
+                                                                              .destination(directory)
+                                                                              .bucket("bucket")
+                                                                              .listObjectsV2RequestTransformer(l -> l.prefix(prefix))
+                                                                              .build());
+        CompletedDirectoryDownload completedDirectoryDownload = downloadDirectory.completionFuture().join();
+        assertThat(completedDirectoryDownload.failedTransfers()).isEmpty();
+
+        List<DownloadFileRequest> actualRequests = requestArgumentCaptor.getAllValues();
+
+        assertThat(actualRequests.size()).isEqualTo(keys.length);
+
+        List<String> destinations =
+            actualRequests.stream().map(u -> u.destination().toString())
+                          .collect(Collectors.toList());
+
+        String jimfsSeparator = jimfs.getSeparator();
+
+        List<String> expectedPaths =
+            Arrays.asList("def/image.jpg", "def/title.jpg", "abcd").stream()
+                  .map(k -> DIRECTORY_NAME + jimfsSeparator + k.replace("/",jimfsSeparator)).collect(Collectors.toList());
+        assertThat(destinations).isEqualTo(expectedPaths);
+    }
+
     @ParameterizedTest
     @MethodSource("fileSystems")
     void downloadDirectory_withDelimiter_shouldHonor(FileSystem jimfs) {