From b4be22b27c67bc9828cf3826d10b64d3dffd8b54 Mon Sep 17 00:00:00 2001
From: Ivan Krutov <vania-pooh@vania-pooh.com>
Date: Tue, 6 Nov 2018 15:06:39 +0300
Subject: [PATCH] Added ability to limit uploaded files (fixes #599)

---
 docs/s3.adoc |  6 ++++++
 s3_test.go   | 30 ++++++++++++++++++++++++++++++
 upload/s3.go | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+)

diff --git a/docs/s3.adoc b/docs/s3.adoc
index ec7944d6..f7525395 100644
--- a/docs/s3.adoc
+++ b/docs/s3.adoc
@@ -29,6 +29,8 @@ By default uploaded file name is preserved, i.e. S3 path is `/<session-id>.log`
 
 For example, when launching Selenoid with `-s3-key-pattern $browserName/$sessionId/log.txt` files will be accessible as `firefox/0ee0b48b-e29b-6749-b4f1-2277b8f8d6c5/log.txt`. You can also override key pattern for every session with `s3KeyPattern` capability.
 
+Sometimes you may want to upload only video files or files matching some complicated pattern or to not upload some files. To achieve this use `-s3-include-files` and `-s3-exclude-files` flags. These flags accept https://en.wikipedia.org/wiki/Glob_(programming)[globs] such as `*.mp4`.
+
 === S3 CLI Flags
 The following flags are supported by `selenoid` command when compiled with S3 support:
 
@@ -39,6 +41,10 @@ The following flags are supported by `selenoid` command when compiled with S3 su
     S3 bucket name
 -s3-endpoint string
     S3 endpoint URL
+-s3-exclude-files string
+    Pattern used to match and exclude files
+-s3-include-files string
+    Pattern used to match and include files    
 -s3-keep-files
     Do not remove uploaded files
 -s3-key-pattern string
diff --git a/s3_test.go b/s3_test.go
index 2a2bd383..1b9a20e6 100644
--- a/s3_test.go
+++ b/s3_test.go
@@ -90,3 +90,33 @@ func TestGetKey(t *testing.T) {
 	key = upload.GetS3Key(testPattern, input)
 	AssertThat(t, key, EqualTo{"some-user/log.txt"})
 }
+
+func TestFileMatches(t *testing.T) {
+	matches, err := upload.FileMatches("", "", "any-file-name")
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, matches, Is{true})
+
+	matches, err = upload.FileMatches("[", "", "/path/to/file.mp4")
+	AssertThat(t, err, Not{nil})
+	AssertThat(t, matches, Is{false})
+
+	matches, err = upload.FileMatches("", "[", "/path/to/file.mp4")
+	AssertThat(t, err, Not{nil})
+	AssertThat(t, matches, Is{false})
+
+	matches, err = upload.FileMatches("*.mp4", "", "/path/to/file.mp4")
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, matches, Is{true})
+
+	matches, err = upload.FileMatches("*.mp4", "", "/path/to/file.log")
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, matches, Is{false})
+
+	matches, err = upload.FileMatches("*.mp4", "", "/path/to/file.log")
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, matches, Is{false})
+
+	matches, err = upload.FileMatches("", "*.log", "/path/to/file.log")
+	AssertThat(t, err, Is{nil})
+	AssertThat(t, matches, Is{false})
+}
diff --git a/upload/s3.go b/upload/s3.go
index a8f489fb..b8e98426 100644
--- a/upload/s3.go
+++ b/upload/s3.go
@@ -27,6 +27,8 @@ func init() {
 	flag.StringVar(&(s3.KeyPattern), "s3-key-pattern", "$fileName", "S3 bucket name")
 	flag.BoolVar(&(s3.ReducedRedundancy), "s3-reduced-redundancy", false, "Use reduced redundancy storage class")
 	flag.BoolVar(&(s3.KeepFiles), "s3-keep-files", false, "Do not remove uploaded files")
+	flag.StringVar(&(s3.IncludeFiles), "s3-include-files", "", "Pattern used to match and include files")
+	flag.StringVar(&(s3.ExcludeFiles), "s3-exclude-files", "", "Pattern used to match and exclude files")
 	uploader = s3
 }
 
@@ -39,6 +41,8 @@ type S3Uploader struct {
 	KeyPattern        string
 	ReducedRedundancy bool
 	KeepFiles         bool
+	IncludeFiles      string
+	ExcludeFiles      string
 
 	manager *s3manager.Uploader
 }
@@ -64,6 +68,14 @@ func (s3 *S3Uploader) Init() {
 func (s3 *S3Uploader) Upload(input *UploadRequest) error {
 	if s3.manager != nil {
 		filename := input.Filename
+		fileMatches, err := FileMatches(s3.IncludeFiles, s3.ExcludeFiles, filename)
+		if err != nil {
+			return fmt.Errorf("invalid pattern: %v", err)
+		}
+		if !fileMatches {
+			log.Printf("[%d] [SKIPPING_FILE] [%s] [Does not match specified patterns]", input.RequestId, input.Filename)
+			return nil
+		}
 		key := GetS3Key(s3.KeyPattern, input)
 		file, err := os.Open(filename)
 		defer file.Close()
@@ -90,6 +102,26 @@ func (s3 *S3Uploader) Upload(input *UploadRequest) error {
 	return errors.New("S3 uploader is not initialized")
 }
 
+func FileMatches(includedFiles string, excludedFiles string, filename string) (bool, error) {
+	fileIncluded := true
+	if includedFiles != "" {
+		fi, err := filepath.Match(includedFiles, filepath.Base(filename))
+		if err != nil {
+			return false, fmt.Errorf("failed to match included file: %v", err)
+		}
+		fileIncluded = fi
+	}
+	fileExcluded := false
+	if excludedFiles != "" {
+		fe, err := filepath.Match(excludedFiles, filepath.Base(filename))
+		if err != nil {
+			return false, fmt.Errorf("failed to match excluded file: %v", err)
+		}
+		fileExcluded = fe
+	}
+	return fileIncluded && !fileExcluded, nil
+}
+
 func GetS3Key(keyPattern string, input *UploadRequest) string {
 	sess := input.Session
 	pt := keyPattern