From 5a57899bb23302c0fc764f75306e09370803aee6 Mon Sep 17 00:00:00 2001
From: clsr <clsr@clsr.net>
Date: Fri, 17 Jun 2016 18:18:12 +0200
Subject: Add --whitelist

---
 USAGE      | 15 ++++++++++-----
 main.go    |  2 ++
 storage.go | 51 +++++++++++++++++++++++++++++++++++----------------
 3 files changed, 47 insertions(+), 21 deletions(-)

diff --git a/USAGE b/USAGE
index b19f508..69db247 100644
--- a/USAGE
+++ b/USAGE
@@ -33,15 +33,13 @@ Running
 		--https HOST:PORT
 			serves HTTPS on HOST:PORT
 			needs --cert and --key
-			example: --https example.com:443
+			example: --https example.com:443 --cert ssl/cert.pem --key ssl/cert.key
 
 		--cert PATH
 			uses PATH as the TLS certificate for HTTPS
-			example: --cert ssl/cert.pem
 
 		--key PATH
 			uses PATH as the TLS certificate key for HTTPS
-			example: --key ssl/cert.key
 
 		--redirect-https
 			redirect HTTP request to HTTPS
@@ -69,13 +67,20 @@ Running
 			equivalent bash example: --bytes $((1024 * 1024 * 10))
 
 		--filter-ext EXTS
-			forbids file extensions contained in the comma-separated list EXTS
+			filter file extensions contained in the comma-separated list EXTS
+			forbids extensions by default, unless --whitelist is in effect
 			example: --filter-ext exe,dll,scr
 
 		--filter-mime TYPES
-			forbids MIME types contained in the comma-separated list TYPES
+			filter MIME types contained in the comma-separated list TYPES
+			forbids types by default, unless --whitelist is in effect
 			example: --filters-mime application/x-dosexec
 
+		--whitelist
+			treat file extension and MIME type filters as whitelists instead of blacklists
+			forbids any upload whose type or extension is not on at least one of the filters
+			example: --whitelist --filter-ext png,jpg,gif --filter-mime=
+
 		--contact EMAIL
 			sets the contact email address to EMAIL
 			example: --contact contact@example.com
diff --git a/main.go b/main.go
index ae2c9f5..2dcac62 100644
--- a/main.go
+++ b/main.go
@@ -70,6 +70,7 @@ func main() {
 	maxSize := flag.Int64("max-size", 50*1024*1024, "max filesize in bytes")
 	filterMime := flag.String("filter-mime", "application/x-dosexec,application/x-msdos-program", "comma-separated list of filtered MIME types")
 	filterExt := flag.String("filter-ext", "exe,dll,msi,scr,com,pif", "comma-separated list of filtered file extensions")
+	whitelist := flag.Bool("whitelist", false, "use filter as a whitelist instead of blacklist")
 	grill := flag.Bool("grill", false, "enable grills")
 	idLength := flag.Int("id-length", 6, "length of uploaded file IDs")
 	idCharset := flag.String("id-charset", "", "charset for uploaded file IDs (default lowercase letters a-z)")
@@ -83,6 +84,7 @@ func main() {
 	storage = NewStorage("upload", *maxSize)
 	storage.FilterExt = strings.Split(*filterExt, ",")
 	storage.FilterMime = strings.Split(*filterMime, ",")
+	storage.Whitelist = *whitelist
 	storage.IdLength = *idLength
 	if *idCharset != "" {
 		storage.IdCharset = *idCharset
diff --git a/storage.go b/storage.go
index 1954c1a..11265b3 100644
--- a/storage.go
+++ b/storage.go
@@ -26,12 +26,13 @@ const (
 )
 
 type Storage struct {
-	Folder        string
-	IdCharset     string
-	IdLength      int
-	MaxSize       int64
+	Folder     string
+	IdCharset  string
+	IdLength   int
+	MaxSize    int64
 	FilterMime []string
 	FilterExt  []string
+	Whitelist  bool
 }
 
 type ErrForbidden struct{ Type string }
@@ -201,27 +202,45 @@ func (s *Storage) getMimeExt(fpath string, name string) (mimetype, ext string, e
 		}
 	}
 
-	// reject filtered MIME types and file extensions
-	if mimetype != "application/octet-stream" {
-		for _, e := range exts {
-			for _, fe := range s.FilterExt {
-				if e == fe {
-					err = ErrForbidden{fe}
-					return
+	filtered, ok := s.findFilter(exts, mimetype)
+	if !ok && s.Whitelist { // whitelist: reject if not on filters
+		err = ErrForbidden{mimetype}
+	} else if ok && !s.Whitelist { // blacklist: reject if filtered
+		forbid := true
+		// only block application/octet-stream if explicitly requested
+		if mimetype == "application/octet-stream" {
+			forbid = false
+			for _, fm := range s.FilterMime {
+				if mimetype == fm {
+					forbid = true
+					break
 				}
 			}
 		}
-		for _, fm := range s.FilterMime {
-			if mimetype == fm {
-				err = ErrForbidden{fm}
-				return
-			}
+		if forbid {
+			err = ErrForbidden{filtered}
 		}
 	}
 
 	return
 }
 
+func (s *Storage) findFilter(exts []string, mimetype string) (match string, ok bool) {
+	for _, fm := range s.FilterMime {
+		if mimetype == fm {
+			return mimetype, true
+		}
+	}
+	for _, ext := range exts {
+		for _, fe := range s.FilterExt {
+			if ext == "."+fe {
+				return ext, true
+			}
+		}
+	}
+	return "", false
+}
+
 func (s *Storage) storeFile(file *os.File, hash, ext string) (id string, err error) {
 	hfolder := s.idToFolder("files", hash)
 	hpath := path.Join(hfolder, "file")
-- 
cgit