From eaad0a9054fad152c1b0df0c1d9a91e37d90c0f0 Mon Sep 17 00:00:00 2001
From: Jasper Lammens <lammensjasper@gmail.com>
Date: Sun, 26 Aug 2018 15:19:16 +0200
Subject: [PATCH] Allow a whitelist of email addresses

---
 README.md           |  1 +
 forwardauth.go      | 12 +++++++++++-
 forwardauth_test.go | 14 ++++++++++++++
 main.go             |  6 ++++++
 4 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index c0ef51b..2d5f680 100644
--- a/README.md
+++ b/README.md
@@ -33,6 +33,7 @@ The following configuration is supported:
 |-csrf-cookie-name|string|CSRF Cookie Name (default "_forward_auth_csrf")|
 |-direct|bool|Run in direct mode (use own hostname as oppose to <br>X-Forwarded-Host, used for testing/development)
 |-domain|string|Comma separated list of email domains to allow|
+|-whitelist|string|Comma separated list of email addresses to allow (Omit -domain)|
 |-lifetime|int|Session length in seconds (default 43200)|
 |-url-path|string|Callback URL (default "_oauth")|
 |-prompt|string|Space separated list of [OpenID prompt options](https://developers.google.com/identity/protocols/OpenIDConnect#prompt)|
diff --git a/forwardauth.go b/forwardauth.go
index 3bc90c7..975e305 100644
--- a/forwardauth.go
+++ b/forwardauth.go
@@ -38,6 +38,7 @@ type ForwardAuth struct {
   CookieSecure bool
 
   Domain []string
+  Whitelist []string
 
   Direct bool
 
@@ -86,12 +87,12 @@ func (f *ForwardAuth) ValidateCookie(r *http.Request, c *http.Cookie) (bool, str
 
 // Validate email
 func (f *ForwardAuth) ValidateEmail(email string) bool {
+  found := false
   if len(f.Domain) > 0 {
     parts := strings.Split(email, "@")
     if len(parts) < 2 {
       return false
     }
-    found := false
     for _, domain := range f.Domain {
       if domain == parts[1] {
         found = true
@@ -100,6 +101,15 @@ func (f *ForwardAuth) ValidateEmail(email string) bool {
     if !found {
       return false
     }
+  } else if len(f.Whitelist) > 0 {
+    for _, wlEmail := range f.Whitelist {
+      if wlEmail == email {
+        found = true
+      }
+    }
+    if !found {
+      return false
+    }
   }
 
   return true
diff --git a/forwardauth_test.go b/forwardauth_test.go
index a5fd09d..8cd0e50 100644
--- a/forwardauth_test.go
+++ b/forwardauth_test.go
@@ -81,6 +81,20 @@ func TestValidateEmail(t *testing.T) {
   if !fw.ValidateEmail("test@test.com") {
     t.Error("Should allow user from allowed domain")
   }
+
+  // Should block non whitelisted email address
+  fw.Domain = []string{}
+  fw.Whitelist = []string{"test@test.com"}
+  if fw.ValidateEmail("one@two.com") {
+    t.Error("Should not allow user not in whitelist.")
+  }
+
+  // Should allow matching whitelisted email address
+  fw.Domain = []string{}
+  fw.Whitelist = []string{"test@test.com"}
+  if !fw.ValidateEmail("test@test.com") {
+    t.Error("Should allow user in whitelist.")
+  }
 }
 
 func TestGetLoginURL(t *testing.T) {
diff --git a/main.go b/main.go
index fd96f52..d846106 100644
--- a/main.go
+++ b/main.go
@@ -141,6 +141,7 @@ func main() {
   cookieSecret := flag.String("cookie-secret", "", "depreciated")
   cookieSecure := flag.Bool("cookie-secure", true, "Use secure cookies")
   domainList := flag.String("domain", "", "Comma separated list of email domains to allow")
+  emailWhitelist := flag.String("whitelist", "", "Comma separated list of emails to allow")
   direct := flag.Bool("direct", false, "Run in direct mode (use own hostname as oppose to X-Forwarded-Host, used for testing/development)")
   prompt := flag.String("prompt", "", "Space separated list of OpenID prompt options")
 
@@ -182,6 +183,10 @@ func main() {
   if *domainList != "" {
     domain = strings.Split(*domainList, ",")
   }
+  var whitelist []string
+  if *emailWhitelist != "" {
+    whitelist = strings.Split(*emailWhitelist, ",")
+  }
 
   // Setup
   fw = &ForwardAuth{
@@ -215,6 +220,7 @@ func main() {
     CookieSecure: *cookieSecure,
 
     Domain: domain,
+    Whitelist: whitelist,
 
     Direct: *direct,