Add multiple auth hosts (#59)

* 🧪 Modify test cases * 🔧 Config Change, should be backwards compatible * ✨ Logic to correctly use multiple authHosts * ✅ Add test cases * fix unused var * fix test failures --------- Co-authored-by: Adyanth H <[email protected]>
mkska · Aug 19, 2024 · a45c261 · a45c261
1 parent 7b720d3
commit a45c261
Show file tree

Hide file tree

Showing 4 changed files with 98 additions and 18 deletions.
diff --git a/internal/auth.go b/internal/auth.go
@@ -137,10 +137,10 @@ func ValidateRedirect(r *http.Request, redirect string) (*url.URL, error) {
 	}
 
 	// If we're using an auth domain?
-	if use, base := useAuthDomain(r); use {
+	if use, _, reqHost := useAuthDomain(r); use {
 		// If we are using an auth domain, they redirect must share a common
 		// suffix with the requested redirect
-		if !strings.HasSuffix(redirectURL.Host, base) {
+		if !strings.HasSuffix(redirectURL.Host, reqHost) {
 			return nil, errors.New("Redirect host does not match any expected hosts (should match cookie domain when using auth host)")
 		}
 	} else {
@@ -167,28 +167,43 @@ func returnUrl(r *http.Request) string {
 
 // Get oauth redirect uri
 func redirectUri(r *http.Request) string {
-	if use, _ := useAuthDomain(r); use {
+	if use, authHost, _ := useAuthDomain(r); use {
 		p := r.Header.Get("X-Forwarded-Proto")
-		return fmt.Sprintf("%s://%s%s", p, config.AuthHost, config.Path)
+		return fmt.Sprintf("%s://%s%s", p, authHost, config.Path)
 	}
 
 	return fmt.Sprintf("%s%s", redirectBase(r), config.Path)
 }
 
 // Should we use auth host + what it is
-func useAuthDomain(r *http.Request) (bool, string) {
-	if config.AuthHost == "" {
-		return false, ""
+func useAuthDomain(r *http.Request) (bool, string, string) {
+	if len(config.AuthHosts) == 0 {
+		return false, "", ""
 	}
 
 	// Does the request match a given cookie domain?
 	reqMatch, reqHost := matchCookieDomains(r.Host)
 
 	// Do any of the auth hosts match a cookie domain?
-	authMatch, authHost := matchCookieDomains(config.AuthHost)
+	authMatch, authHost := matchAuthHosts(reqHost)
 
 	// We need both to match the same domain
-	return reqMatch && authMatch && reqHost == authHost, reqHost
+	return reqMatch && authMatch, authHost, reqHost
+}
+
+// Return matching auth host domain if exists
+func matchAuthHosts(domain string) (bool, string) {
+	// Remove port
+	p := strings.Split(domain, ":")
+
+	for _, d := range config.AuthHosts {
+		// Subdomain match?
+		if len(d) >= len(domain) && d[len(d)-len(domain):] == domain {
+			return true, d
+		}
+	}
+
+	return false, p[0]
 }
 
 // Cookie methods
@@ -320,7 +335,7 @@ func cookieDomain(r *http.Request) string {
 // Cookie domain
 func csrfCookieDomain(r *http.Request) string {
 	var host string
-	if use, domain := useAuthDomain(r); use {
+	if use, _, domain := useAuthDomain(r); use {
 		host = domain
 	} else {
 		host = r.Host

diff --git a/internal/auth_test.go b/internal/auth_test.go
@@ -266,7 +266,7 @@ func TestAuthValidateRedirect(t *testing.T) {
 	//
 	// With Auth Host
 	//
-	config.AuthHost = "auth.example.com"
+	config.AuthHosts = CommaSeparatedList{"auth.example.com"}
 	config.CookieDomains = []CookieDomain{*NewCookieDomain("example.com")}
 	errStr = "Redirect host does not match any expected hosts (should match cookie domain when using auth host)"
 
@@ -326,7 +326,7 @@ func TestRedirectUri(t *testing.T) {
 	// With Auth URL but no matching cookie domain
 	// - will not use auth host
 	//
-	config.AuthHost = "auth.example.com"
+	config.AuthHosts = CommaSeparatedList{"auth.example.com"}
 
 	uri, err = url.Parse(redirectUri(r))
 	assert.Nil(err)
@@ -337,7 +337,7 @@ func TestRedirectUri(t *testing.T) {
 	//
 	// With correct Auth URL + cookie domain
 	//
-	config.AuthHost = "auth.example.com"
+	config.AuthHosts = CommaSeparatedList{"auth.example.com"}
 	config.CookieDomains = []CookieDomain{*NewCookieDomain("example.com")}
 
 	// Check url
@@ -354,7 +354,7 @@ func TestRedirectUri(t *testing.T) {
 	r = httptest.NewRequest("GET", "https://another.com/hello", nil)
 	r.Header.Add("X-Forwarded-Proto", "https")
 
-	config.AuthHost = "auth.example.com"
+	config.AuthHosts = CommaSeparatedList{"auth.example.com"}
 	config.CookieDomains = []CookieDomain{*NewCookieDomain("example.com")}
 
 	// Check url
@@ -363,6 +363,71 @@ func TestRedirectUri(t *testing.T) {
 	assert.Equal("https", uri.Scheme)
 	assert.Equal("another.com", uri.Host)
 	assert.Equal("/_oauth", uri.Path)
+
+	//
+	// With correct Auth URL + cookie domain, multiple AuthHosts and cookie domains
+	// - will use matching authHost
+	//
+	config.AuthHosts = CommaSeparatedList{"auth.example.com", "auth.another.com"}
+	config.CookieDomains = []CookieDomain{*NewCookieDomain("example.com"), *NewCookieDomain("another.com")}
+
+	uri, err = url.Parse(redirectUri(r))
+	assert.Nil(err)
+	assert.Equal("https", uri.Scheme)
+	assert.Equal("auth.another.com", uri.Host)
+	assert.Equal("/_oauth", uri.Path)
+
+	//
+	// With correct Auth URL + no cookie domains
+	// - will not use authHost
+	//
+	config.AuthHosts = CommaSeparatedList{"auth.example.com", "auth.another.com"}
+	config.CookieDomains = []CookieDomain{}
+
+	uri, err = url.Parse(redirectUri(r))
+	assert.Nil(err)
+	assert.Equal("https", uri.Scheme)
+	assert.Equal("another.com", uri.Host)
+	assert.Equal("/_oauth", uri.Path)
+
+	//
+	// With correct Auth URL + no matching cookie domains
+	// - will not use authHost
+	//
+	config.AuthHosts = CommaSeparatedList{"auth.example.com", "auth.another.com"}
+	config.CookieDomains = []CookieDomain{*NewCookieDomain("example.com"), *NewCookieDomain("another.example")}
+
+	uri, err = url.Parse(redirectUri(r))
+	assert.Nil(err)
+	assert.Equal("https", uri.Scheme)
+	assert.Equal("another.com", uri.Host)
+	assert.Equal("/_oauth", uri.Path)
+
+	//
+	// With no matching Auth Host + matching cookie domains
+	// - will not use authHost
+	//
+	config.AuthHosts = CommaSeparatedList{"auth.example.com", "auth.another.example"}
+	config.CookieDomains = []CookieDomain{*NewCookieDomain("example.com"), *NewCookieDomain("another.com")}
+
+	uri, err = url.Parse(redirectUri(r))
+	assert.Nil(err)
+	assert.Equal("https", uri.Scheme)
+	assert.Equal("another.com", uri.Host)
+	assert.Equal("/_oauth", uri.Path)
+
+	//
+	// With no matching Auth Host + no matching cookie domains
+	// - will not use authHost
+	//
+	config.AuthHosts = CommaSeparatedList{"auth.example.com", "auth.another.example"}
+	config.CookieDomains = []CookieDomain{*NewCookieDomain("example.com"), *NewCookieDomain("another.example")}
+
+	uri, err = url.Parse(redirectUri(r))
+	assert.Nil(err)
+	assert.Equal("https", uri.Scheme)
+	assert.Equal("another.com", uri.Host)
+	assert.Equal("/_oauth", uri.Path)
 }
 
 func TestAuthMakeCookie(t *testing.T) {
@@ -416,7 +481,7 @@ func TestAuthMakeCSRFCookie(t *testing.T) {
 	assert.Equal("app.example.com", c.Domain)
 
 	// With cookie domain and auth url
-	config.AuthHost = "auth.example.com"
+	config.AuthHosts = CommaSeparatedList{"auth.example.com"}
 	config.CookieDomains = []CookieDomain{*NewCookieDomain("example.com")}
 	c = MakeCSRFCookie(r, "12333378901234567890123456789012")
 	assert.Equal("_forward_auth_csrf_123333", c.Name)

diff --git a/internal/config.go b/internal/config.go
@@ -25,7 +25,7 @@ type Config struct {
 	LogLevel  string `long:"log-level" env:"LOG_LEVEL" default:"warn" choice:"trace" choice:"debug" choice:"info" choice:"warn" choice:"error" choice:"fatal" choice:"panic" description:"Log level"`
 	LogFormat string `long:"log-format"  env:"LOG_FORMAT" default:"text" choice:"text" choice:"json" choice:"pretty" description:"Log format"`
 
-	AuthHost               string               `long:"auth-host" env:"AUTH_HOST" description:"Single host to use when returning from 3rd party auth"`
+	AuthHosts              CommaSeparatedList   `long:"auth-host" env:"AUTH_HOST" env-delim:"," description:"Single host to use when returning from 3rd party auth"`
 	Config                 func(s string) error `long:"config" env:"CONFIG" description:"Path to config file" json:"-"`
 	CookieDomains          []CookieDomain       `long:"cookie-domain" env:"COOKIE_DOMAIN" env-delim:"," description:"Domain to set auth cookie on, can be set multiple times"`
 	InsecureCookie         bool                 `long:"insecure-cookie" env:"INSECURE_COOKIE" description:"Use insecure cookies"`

diff --git a/internal/config_test.go b/internal/config_test.go
@@ -24,7 +24,7 @@ func TestConfigDefaults(t *testing.T) {
 	assert.Equal("warn", c.LogLevel)
 	assert.Equal("text", c.LogFormat)
 
-	assert.Equal("", c.AuthHost)
+	assert.Len(c.AuthHosts, 0)
 	assert.Len(c.CookieDomains, 0)
 	assert.False(c.InsecureCookie)
 	assert.Equal("_forward_auth", c.CookieName)
@@ -210,7 +210,7 @@ func TestConfigFileBackwardsCompatability(t *testing.T) {
 	require.Nil(t, err)
 
 	assert.Equal("/two", c.Path, "variable in legacy config file should be read")
-	assert.Equal("auth.legacy.com", c.AuthHost, "variable in legacy config file should be read")
+	assert.Equal(CommaSeparatedList{"auth.legacy.com"}, c.AuthHosts, "variable in legacy config file should be read")
 }
 
 func TestConfigParseEnvironment(t *testing.T) {