[SHRINKWRAP-469] Fix for ShrinkWrapDirectoryStream.iterator missing I…

…llegalStateException cases

impl-nio2/src/main/java/org/jboss/shrinkwrap/impl/nio/file/ShrinkWrapDirectoryStream.java

@@ -40,6 +40,10 @@ class ShrinkWrapDirectoryStream implements DirectoryStream<Path> {
 
     private final Path startingPath;
 
+    private boolean closed = false;
+
+    private boolean iteratorReturned = false;
+
     /**
      * Creates a new instance starting from startingPath with is required backing the specified
      * {@link ShrinkWrapFileSystem}, which is required. An optional {@link DirectoryStream.Filter} may be
@@ -72,7 +76,7 @@ class ShrinkWrapDirectoryStream implements DirectoryStream<Path> {
      */
     @Override
     public void close() throws IOException {
-        // NO-OP
+        this.closed = true;
     }
 
     /**
@@ -82,34 +86,49 @@ public void close() throws IOException {
      */
     @Override
     public Iterator<Path> iterator() {
+        if (closed) {
+            throw new IllegalStateException("Directory Stream is closed");
+        } else if (iteratorReturned) {
+            throw new IllegalStateException("Iterator was already returned");
+        }
 
-        // Translate ShrinkWrap API to NIO.2 API Path
-        final Map<ArchivePath, Node> content = this.fs.getArchive().getContent();
-        final Collection<Path> newPaths = new ArrayList<>(content.size());
-        final Collection<ArchivePath> archivePaths = content.keySet();
-        for (final ArchivePath path : archivePaths) {
-            final Path newPath = new ShrinkWrapPath(path, fs);
+        boolean finishedSuccessfully = true;
+        try {
+            // Translate ShrinkWrap API to NIO.2 API Path
+            final Map<ArchivePath, Node> content = this.fs.getArchive().getContent();
+            final Collection<Path> newPaths = new ArrayList<>(content.size());
+            final Collection<ArchivePath> archivePaths = content.keySet();
+            for (final ArchivePath path : archivePaths) {
+                final Path newPath = new ShrinkWrapPath(path, fs);
 
-            if (!newPath.getParent().equals(startingPath)) {
-                continue;
-            }
-
-            // If we have a filter, and it rejects this path
-            try {
-                if (filter != null && !(filter.accept(newPath))) {
-                    // Move along
+                if (!newPath.getParent().equals(startingPath)) {
                     continue;
                 }
-            } catch (IOException ioe) {
-                throw new RuntimeException("Error encountered during filtering", ioe);
+
+                // If we have a filter, and it rejects this path
+                try {
+                    if (filter != null && !(filter.accept(newPath))) {
+                        // Move along
+                        continue;
+                    }
+                } catch (IOException ioe) {
+                    throw new RuntimeException("Error encountered during filtering", ioe);
+                }
+
+                // Add the new Path; the filter either wasn't specified or didn't reject this Path
+                newPaths.add(newPath);
             }
 
-            // Add the new Path; the filter either wasn't specified or didn't reject this Path
-            newPaths.add(newPath);
+            // Return
+            return newPaths.iterator();
+        } catch (Throwable t) {
+            finishedSuccessfully = false;
+            throw t;
+        } finally {
+            if (finishedSuccessfully) {
+                iteratorReturned = true;
+            }
         }
-
-        // Return
-        return newPaths.iterator();
     }
 
 }

impl-nio2/src/test/java/org/jboss/shrinkwrap/impl/nio/file/FilesTestCase.java

@@ -57,7 +57,9 @@
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 
 /**
  * Test cases to assert the ShrinkWrap implementation of the NIO.2 {@link FileSystem} is working as expected via the
@@ -70,6 +72,9 @@ public class FilesTestCase {
     @SuppressWarnings("unused")
     private static final Logger log = Logger.getLogger(FilesTestCase.class.getName());
 
+    @Rule
+    public ExpectedException expectedException = ExpectedException.none();
+
     /**
      * {@link FileSystem} under test
      */
@@ -758,6 +763,39 @@ public void createdDirectoryCanBeStreamed() throws Exception {
         Assert.assertFalse("No further elements expected in stream", it.hasNext());
     }
 
+    @Test
+    public void createdDirectoryStreamThrowsExceptionWhenIsClosed() throws Throwable {
+        // given
+        Path dirPath = fs.getPath("dir");
+        Files.createDirectory(dirPath);
+        Files.createFile(dirPath.resolve("file"));
+        DirectoryStream<Path> stream = Files.newDirectoryStream(dirPath);
+
+        // when
+        stream.close();
+
+        // then
+        expectedException.expect(IllegalStateException.class);
+        Iterator<Path> it = stream.iterator();
+
+    }
+
+    @Test
+    public void createdDirectoryStreamThrowsExceptionWhenIteratorWasReturnedBefore() throws Exception {
+        // given
+        Path dirPath = fs.getPath("dir");
+        Files.createDirectory(dirPath);
+        Files.createFile(dirPath.resolve("file"));
+        DirectoryStream<Path> stream = Files.newDirectoryStream(dirPath);
+
+        // when
+        Iterator<Path> it = stream.iterator();
+
+        // then
+        expectedException.expect(IllegalStateException.class);
+        stream.iterator();
+    }
+
     @Test
     public void directoryStreamDoesNotContainSubfolders() throws Exception {
         Files.createDirectories(fs.getPath("dir/subdir"));