From 924efc9542df30c04a7515361277fa5881576837 Mon Sep 17 00:00:00 2001
From: bhou <bhou@netflix.com>
Date: Fri, 3 May 2024 13:52:14 -0700
Subject: [PATCH] Amend the check on IllegalAttachmentFileNameException

---
 .../LocalFileSystemAttachmentServiceImpl.java | 27 ++++++++++--
 ...FileSystemAttachmentServiceImplSpec.groovy | 44 ++++++++++++++++++-
 2 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/genie-web/src/main/java/com/netflix/genie/web/services/impl/LocalFileSystemAttachmentServiceImpl.java b/genie-web/src/main/java/com/netflix/genie/web/services/impl/LocalFileSystemAttachmentServiceImpl.java
index 1d596ad92a..d6cf472412 100644
--- a/genie-web/src/main/java/com/netflix/genie/web/services/impl/LocalFileSystemAttachmentServiceImpl.java
+++ b/genie-web/src/main/java/com/netflix/genie/web/services/impl/LocalFileSystemAttachmentServiceImpl.java
@@ -27,6 +27,7 @@
 import org.springframework.core.io.Resource;
 
 import javax.annotation.Nullable;
+import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URI;
@@ -94,9 +95,24 @@ public Set<URI> saveAttachments(
                 final long attachmentSize = attachment.contentLength();
                 final String filename = attachment.getFilename();
 
-                if (filename != null && (filename.contains("/") || filename.contains("\\"))) {
-                    throw new IllegalAttachmentFileNameException("Attachment filename " + filename + " is illegal. "
-                        + "Filenames should not contain / or \\.");
+                if (filename != null) {
+                    if ((filename.contains("/") || filename.contains("\\")
+                        || filename.equals(".") || filename.equals(".."))) {
+                        throw new IllegalAttachmentFileNameException("Attachment filename " + filename + " is illegal. "
+                            + "Filenames should not be . or .., or contain /, \\.");
+                    }
+
+                    final String attachmentCanonicalPath =
+                        createTempFile(String.valueOf(attachmentsBasePath), filename).getCanonicalPath();
+
+                    final String baseCanonicalPath =
+                        new File(String.valueOf(attachmentsBasePath)).getCanonicalPath();
+
+                    if (!attachmentCanonicalPath.startsWith(baseCanonicalPath)
+                        || attachmentCanonicalPath.equals(baseCanonicalPath)) {
+                        throw new IllegalAttachmentFileNameException("Attachment filename " + filename + " is illegal. "
+                            + "Filenames should not be a relative path.");
+                    }
                 }
 
                 if (attachmentSize > this.attachmentServiceProperties.getMaxSize().toBytes()) {
@@ -124,4 +140,9 @@ public Set<URI> saveAttachments(
 
         return setBuilder.build();
     }
+
+    /* for testing purposes */
+    File createTempFile(final String attachmentsBasePath, final String filename) {
+        return new File(attachmentsBasePath, filename);
+    }
 }
diff --git a/genie-web/src/test/groovy/com/netflix/genie/web/services/impl/LocalFileSystemAttachmentServiceImplSpec.groovy b/genie-web/src/test/groovy/com/netflix/genie/web/services/impl/LocalFileSystemAttachmentServiceImplSpec.groovy
index 75f9cc700d..5cfdc0cd9e 100644
--- a/genie-web/src/test/groovy/com/netflix/genie/web/services/impl/LocalFileSystemAttachmentServiceImplSpec.groovy
+++ b/genie-web/src/test/groovy/com/netflix/genie/web/services/impl/LocalFileSystemAttachmentServiceImplSpec.groovy
@@ -51,7 +51,7 @@ class LocalFileSystemAttachmentServiceImplSpec extends Specification {
         this.serviceProperties.setLocationPrefix(this.temporaryFolder.toUri())
         this.serviceProperties.setMaxSize(DataSize.ofBytes(100))
         this.serviceProperties.setMaxTotalSize(DataSize.ofBytes(150))
-        this.service = new LocalFileSystemAttachmentServiceImpl(serviceProperties)
+        this.service = Mockito.spy(new LocalFileSystemAttachmentServiceImpl(serviceProperties))
     }
 
     def "saveAttachments with no attachments"() {
@@ -179,4 +179,46 @@ class LocalFileSystemAttachmentServiceImplSpec extends Specification {
         then:
         thrown(IllegalAttachmentFileNameException)
     }
+
+    def "reject attachments with illegal filename is ."() {
+        Set<Resource> attachments = new HashSet<Resource>()
+        Resource attachment = Mockito.mock(Resource.class)
+        Mockito.doReturn(".").when(attachment).getFilename()
+        attachments.add(attachment)
+
+        when:
+        service.saveAttachments(null, attachments)
+
+        then:
+        thrown(IllegalAttachmentFileNameException)
+    }
+
+    def "reject attachments with illegal filename is .."() {
+        Set<Resource> attachments = new HashSet<Resource>()
+        Resource attachment = Mockito.mock(Resource.class)
+        Mockito.doReturn("..").when(attachment).getFilename()
+        attachments.add(attachment)
+
+        when:
+        service.saveAttachments(null, attachments)
+
+        then:
+        thrown(IllegalAttachmentFileNameException)
+    }
+
+    def "reject attachments with illegal filename for base path check"() {
+        Set<Resource> attachments = new HashSet<Resource>()
+        Resource attachment = Mockito.mock(Resource.class)
+        Mockito.doReturn("file1.text").when(attachment).getFilename()
+        File file = Mockito.mock(File.class)
+        Mockito.doReturn("/dummy").when(file).getCanonicalPath()
+        Mockito.doReturn(file).when(service).createTempFile(Mockito.anyString(), Mockito.anyString());
+        attachments.add(attachment)
+
+        when:
+        service.saveAttachments(null, attachments)
+
+        then:
+        thrown(IllegalAttachmentFileNameException)
+    }
 }