tests: fix tests

internal/bootstrap/app_bootstrap.go

@@ -166,7 +166,7 @@ func (app *BootstrapApp) Setup() error {
 		app.log.App.Warn().Msg("Subdomains are disabled, cookies will be set for the current domain only")
 	}
 
-	cookieDomain, err := utils.GetCookieDomain(app.runtime.AppURL)
+	cookieDomain, err := utils.GetCookieDomain(app.runtime.AppURL, app.config.Auth.SubdomainsEnabled)
 
 	if err != nil {
 		return fmt.Errorf("failed to get cookie domain: %w", err)
@@ -288,7 +288,7 @@ func (app *BootstrapApp) Setup() error {
 			app.log.App.Info().Msg("Listening on tailscale, replacing app url with tailscale hostname")
 			app.runtime.AppURL = tailscaleUrl
 			// also update cookie domain
-			cookieDomain, err := utils.GetCookieDomain(tailscaleUrl)
+			cookieDomain, err := utils.GetCookieDomain(tailscaleUrl, app.config.Auth.SubdomainsEnabled)
 
 			if err != nil {
 				return fmt.Errorf("failed to get cookie domain: %w", err)

internal/controller/context_controller_test.go

@@ -48,9 +48,9 @@ func TestContextController(t *testing.T) {
 						WarningsEnabled:       cfg.UI.WarningsEnabled,
 					},
 					App: ACRApp{
-						AppURL:         runtime.AppURL,
-						CookieDomain:   runtime.CookieDomain,
-						TrustedDomains: runtime.TrustedDomains,
+						AppURL:            runtime.AppURL,
+						CookieDomain:      runtime.CookieDomain,
+						SubdomainsEnabled: cfg.Auth.SubdomainsEnabled,
 					},
 				}
 				bytes, err := json.Marshal(expectedAppContextResponse)

internal/controller/oauth_controller_test.go

@@ -9,7 +9,7 @@ import (
 	"github.com/tinyauthapp/tinyauth/internal/utils/logger"
 )
 
-func TestOAuthController(t *testing.T) {
+func TestOAuthControllerIsRedirectSafe(t *testing.T) {
 	log := logger.NewLogger().WithTestConfig()
 	log.Init()
 
@@ -17,145 +17,171 @@ func TestOAuthController(t *testing.T) {
 
 	type testCase struct {
 		description       string
-		run               func(ctrl *OAuthController)
-		trustedDomains    []string
+		appURL            string
+		cookieDomain      string
 		subdomainsEnabled bool
+		redirectURI       string
+		expected          bool
 	}
 
 	tests := []testCase{
 		{
-			description:       "Test exact match of redirect URI",
-			trustedDomains:    []string{"https://tinyauth.example.com"},
+			description:       "Exact host match returns true",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
 			subdomainsEnabled: true,
-			run: func(ctrl *OAuthController) {
-				redirectUri := "https://tinyauth.example.com"
-				assert.True(t, ctrl.isRedirectSafe(redirectUri))
-			},
+			redirectURI:       "https://tinyauth.example.com",
+			expected:          true,
 		},
 		{
-			description:       "Test subdomain match of redirect URI",
-			trustedDomains:    []string{"https://tinyauth.example.com"},
+			description:       "Exact host match is case insensitive",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
 			subdomainsEnabled: true,
-			run: func(ctrl *OAuthController) {
-				redirectUri := "https://sub.example.com"
-				assert.True(t, ctrl.isRedirectSafe(redirectUri))
-			},
+			redirectURI:       "https://TinyAuth.Example.com",
+			expected:          true,
 		},
 		{
-			description:       "Test different trusted domain",
-			trustedDomains:    []string{"https://tinyauth.example.com", "https://tinyauth.foo.com"},
-			subdomainsEnabled: true,
-			run: func(ctrl *OAuthController) {
-				redirectUri := "https://app.foo.com"
-				assert.True(t, ctrl.isRedirectSafe(redirectUri))
-			},
-		},
-		{
-			description: "Test invalid redirect URI",
-			run: func(ctrl *OAuthController) {
-				redirectUri := "https:/malicious"
-				assert.False(t, ctrl.isRedirectSafe(redirectUri))
-			},
-		},
-		{
-			description: "Test empty redirect URI",
-			run: func(ctrl *OAuthController) {
-				redirectUri := ""
-				assert.False(t, ctrl.isRedirectSafe(redirectUri))
-			},
-		},
-		{
-			description:       "Test redirect URI with different scheme",
-			trustedDomains:    []string{"https://tinyauth.example.com"},
-			subdomainsEnabled: true,
-			run: func(ctrl *OAuthController) {
-				redirectUri := "http://tinyauth.example.com"
-				assert.False(t, ctrl.isRedirectSafe(redirectUri))
-			},
-		},
-		{
-			description:       "Test redirect URI with different port",
-			trustedDomains:    []string{"https://tinyauth.example.com"},
-			subdomainsEnabled: true,
-			run: func(ctrl *OAuthController) {
-				redirectUri := "https://tinyauth.example.com:8080"
-				assert.False(t, ctrl.isRedirectSafe(redirectUri))
-			},
-		},
-		{
-			// weird case, subdomains enabled and domain without subdomain can't happen
-			description:    "Test with trusted domain that's in PSL when split",
-			trustedDomains: []string{"https://example.com"}, // will become .com which we
-			// obviously don't want to allow
-			subdomainsEnabled: true,
-			run: func(ctrl *OAuthController) {
-				redirectUri := "https://sub.example.com"
-				assert.False(t, ctrl.isRedirectSafe(redirectUri))
-			},
-		},
-		{
-			description:       "Test subdomain redirect URI when subdomains are disabled",
-			trustedDomains:    []string{"https://tinyauth.example.com"},
+			description:       "Exact host match with subdomains disabled returns true",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
 			subdomainsEnabled: false,
-			run: func(ctrl *OAuthController) {
-				redirectUri := "https://sub.tinyauth.example.com"
-				assert.False(t, ctrl.isRedirectSafe(redirectUri))
-			},
+			redirectURI:       "https://tinyauth.example.com",
+			expected:          true,
 		},
 		{
-			description:       "Test domain like the .co.uk",
-			trustedDomains:    []string{"https://example.co.uk"},
+			description:       "Subdomain of cookie domain returns true when subdomains enabled",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
 			subdomainsEnabled: true,
-			run: func(ctrl *OAuthController) {
-				redirectUri := "https://sub.example.co.uk"
-				assert.False(t, ctrl.isRedirectSafe(redirectUri))
-			},
+			redirectURI:       "https://sub.example.com",
+			expected:          true,
 		},
 		{
-			description:       "Test domain like the .co.uk with subdomains disabled",
-			trustedDomains:    []string{"https://example.co.uk"},
+			description:       "Subdomain of cookie domain is case insensitive",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "Example.COM",
+			subdomainsEnabled: true,
+			redirectURI:       "https://SUB.example.com",
+			expected:          true,
+		},
+		{
+			description:       "Subdomain not matching cookie domain returns false",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
+			subdomainsEnabled: true,
+			redirectURI:       "https://sub.evil.com",
+			expected:          false,
+		},
+		{
+			description:       "Subdomain returns false when subdomains disabled",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
 			subdomainsEnabled: false,
-			run: func(ctrl *OAuthController) {
-				redirectUri := "https://example.co.uk"
-				assert.True(t, ctrl.isRedirectSafe(redirectUri))
-			},
+			redirectURI:       "https://sub.example.com",
+			expected:          false,
 		},
 		{
-			description:       "Test caps domain",
-			trustedDomains:    []string{"https://TINYAUTH.ExAmpLe.com"},
+			description:       "Cookie domain itself is not a subdomain match",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
 			subdomainsEnabled: true,
-			run: func(ctrl *OAuthController) {
-				redirectUri := "https://sUb.ExAmPle.com"
-				assert.True(t, ctrl.isRedirectSafe(redirectUri))
-			},
+			redirectURI:       "https://example.com",
+			expected:          false,
 		},
 		{
-			description:       "Test edge case with @",
-			trustedDomains:    []string{"https://tinyauth.example.com"},
+			description:       "Different scheme returns false",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
 			subdomainsEnabled: true,
-			run: func(ctrl *OAuthController) {
-				redirectUri := "https://malicious.example.com@evil.com"
-				assert.False(t, ctrl.isRedirectSafe(redirectUri))
-			},
+			redirectURI:       "http://tinyauth.example.com",
+			expected:          false,
+		},
+		{
+			description:       "Different port returns false",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
+			subdomainsEnabled: true,
+			redirectURI:       "https://tinyauth.example.com:8080",
+			expected:          false,
+		},
+		{
+			description:       "Empty redirect URI returns false",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
+			subdomainsEnabled: true,
+			redirectURI:       "",
+			expected:          false,
+		},
+		{
+			description:       "Redirect URI without host returns false",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
+			subdomainsEnabled: true,
+			redirectURI:       "https:/malicious",
+			expected:          false,
+		},
+		{
+			description:       "Redirect URI without scheme returns false",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
+			subdomainsEnabled: true,
+			redirectURI:       "tinyauth.example.com",
+			expected:          false,
+		},
+		{
+			description:       "Relative redirect URI returns false",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
+			subdomainsEnabled: true,
+			redirectURI:       "/some/path",
+			expected:          false,
+		},
+		{
+			description:       "Userinfo trick with malicious host returns false",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
+			subdomainsEnabled: true,
+			redirectURI:       "https://malicious.example.com@evil.com",
+			expected:          false,
+		},
+		{
+			description:       "Unparseable redirect URI returns false",
+			appURL:            "https://tinyauth.example.com",
+			cookieDomain:      "example.com",
+			subdomainsEnabled: true,
+			redirectURI:       "https://exa\x7fmple.com",
+			expected:          false,
+		},
+		{
+			description:       "Unparseable app URL returns false",
+			appURL:            "https://tinyauth.\x7fexample.com",
+			cookieDomain:      "example.com",
+			subdomainsEnabled: true,
+			redirectURI:       "https://tinyauth.example.com",
+			expected:          false,
 		},
 	}
 
-	// TODO: add auth service
 	for _, tc := range tests {
 		t.Run(tc.description, func(t *testing.T) {
 			router := gin.Default()
 			group := router.Group("/api")
 			gin.SetMode(gin.TestMode)
-			// overwrite the trusted domains and subdomain setting for each test case
-			runtime.TrustedDomains = tc.trustedDomains
+
+			// Overwrite the app URL, cookie domain and subdomain setting for each test case
+			runtime.AppURL = tc.appURL
+			runtime.CookieDomain = tc.cookieDomain
 			cfg.Auth.SubdomainsEnabled = tc.subdomainsEnabled
+
 			ctrl := NewOAuthController(OAuthControllerInput{
 				Log:           log,
 				Config:        &cfg,
 				RuntimeConfig: &runtime,
 				RouterGroup:   group,
 			})
-			tc.run(ctrl)
+
+			assert.Equal(t, tc.expected, ctrl.isRedirectSafe(tc.redirectURI))
 		})
 	}
 }

internal/service/auth_service.go

@@ -380,33 +380,11 @@ func (auth *AuthService) CreateSession(ctx context.Context, data repository.Sess
 		return nil, fmt.Errorf("failed to create session entry: %w", err)
 	}
 
-	if data.Provider == "tailscale" {
-		auth.log.App.Trace().Str("url", fmt.Sprintf("https://%s", auth.tailscale.GetHostname())).Msg("Extracting root domain from Tailscale hostname")
-
-		tsCookieDomain, err := utils.GetCookieDomain(fmt.Sprintf("https://%s", auth.tailscale.GetHostname()))
-
-		if err != nil {
-			return nil, fmt.Errorf("failed to get cookie domain for tailscale user: %w", err)
-		}
-
-		return &http.Cookie{
-			Name:     auth.runtime.SessionCookieName,
-			Value:    session.UUID,
-			Path:     "/",
-			Domain:   fmt.Sprintf(".%s", tsCookieDomain),
-			Expires:  expiresAt,
-			MaxAge:   int(time.Until(expiresAt).Seconds()),
-			Secure:   auth.config.Auth.SecureCookie,
-			HttpOnly: true,
-			SameSite: http.SameSiteLaxMode,
-		}, nil
-	}
-
 	return &http.Cookie{
 		Name:     auth.runtime.SessionCookieName,
 		Value:    session.UUID,
 		Path:     "/",
-		Domain:   fmt.Sprintf(".%s", auth.runtime.CookieDomain),
+		Domain:   auth.getCookieDomain(),
 		Expires:  expiresAt,
 		MaxAge:   int(time.Until(expiresAt).Seconds()),
 		Secure:   auth.config.Auth.SecureCookie,
@@ -459,7 +437,7 @@ func (auth *AuthService) RefreshSession(ctx context.Context, uuid string) (*http
 		Name:     auth.runtime.SessionCookieName,
 		Value:    session.UUID,
 		Path:     "/",
-		Domain:   fmt.Sprintf(".%s", auth.runtime.CookieDomain),
+		Domain:   auth.getCookieDomain(),
 		Expires:  time.Now().Add(time.Duration(newExpiry-currentTime) * time.Second),
 		MaxAge:   int(newExpiry - currentTime),
 		Secure:   auth.config.Auth.SecureCookie,
@@ -480,7 +458,7 @@ func (auth *AuthService) DeleteSession(ctx context.Context, uuid string) (*http.
 		Name:     auth.runtime.SessionCookieName,
 		Value:    "",
 		Path:     "/",
-		Domain:   fmt.Sprintf(".%s", auth.runtime.CookieDomain),
+		Domain:   auth.getCookieDomain(),
 		Expires:  time.Now(),
 		MaxAge:   -1,
 		Secure:   auth.config.Auth.SecureCookie,
@@ -726,3 +704,10 @@ func (auth *AuthService) calculateLockdownLimit() int {
 
 	return limit
 }
+
+func (auth *AuthService) getCookieDomain() string {
+	if auth.config.Auth.SubdomainsEnabled {
+		return "." + auth.runtime.CookieDomain
+	}
+	return auth.runtime.CookieDomain
+}

internal/test/test.go

@@ -43,6 +43,7 @@ func CreateTestConfigs(t *testing.T) (model.Config, model.RuntimeConfig) {
 			ACLs: model.ACLsConfig{
 				Policy: "allow",
 			},
+			SubdomainsEnabled: true,
 		},
 		Database: model.DatabaseConfig{
 			Path: filepath.Join(tempDir, "test.db"),
@@ -165,10 +166,6 @@ func CreateTestConfigs(t *testing.T) (model.Config, model.RuntimeConfig) {
 		CookieDomain:      "example.com",
 		AppURL:            "https://tinyauth.example.com",
 		SessionCookieName: "tinyauth-session",
-		TrustedDomains: []string{
-			"https://tinyauth.example.com",
-			"https://tinyauth.foo.com",
-		},
 	}
 
 	return config, runtime

internal/utils/app_utils.go

@@ -10,7 +10,7 @@ import (
 )
 
 // Get cookie domain parses a hostname and returns the upper domain (e.g. sub1.sub2.domain.com -> sub2.domain.com)
-func GetCookieDomain(appUrl string) (string, error) {
+func GetCookieDomain(appUrl string, subdomainsEnabled bool) (string, error) {
 	u, err := url.Parse(appUrl)
 
 	if err != nil {
@@ -29,12 +29,16 @@ func GetCookieDomain(appUrl string) (string, error) {
 		return "", fmt.Errorf("invalid app url, must be in format subdomain.domain.tld or domain.tld")
 	}
 
-	if len(parts) == 2 {
+	if !subdomainsEnabled || len(parts) == 2 {
+		_, err = publicsuffix.DomainFromListWithOptions(publicsuffix.DefaultList, hostname, nil)
+
+		if err != nil {
+			return "", fmt.Errorf("domain in public suffix list, cannot set cookies: %w", err)
+		}
+
 		return strings.ToLower(u.Host), nil
 	}
 
-	// parts > 3
-
 	domain := strings.Join(parts[1:], ".")
 
 	_, err = publicsuffix.DomainFromListWithOptions(publicsuffix.DefaultList, domain, nil)
@@ -44,10 +48,10 @@ func GetCookieDomain(appUrl string) (string, error) {
 	}
 
 	// now that we validated the domain, return with the port
-	parts = strings.Split(strings.ToLower(u.Host), ":")
-	domainWithPort := strings.Join(parts[1:], ":")
+	parts = strings.Split(strings.ToLower(u.Host), ".")
+	host := strings.Join(parts[1:], ".")
 
-	return domainWithPort, nil
+	return host, nil
 }
 
 func ParseFileToLine(content string) string {

internal/utils/app_utils_test.go

@@ -11,50 +11,71 @@ func TestGetRootDomain(t *testing.T) {
 	// Normal case
 	domain := "http://sub.tinyauth.app"
 	expected := "tinyauth.app"
-	result, err := utils.GetCookieDomain(domain)
+	result, err := utils.GetCookieDomain(domain, true)
 	assert.NoError(t, err)
 	assert.Equal(t, expected, result)
 
 	// Domain with multiple subdomains
 	domain = "http://b.c.tinyauth.app"
 	expected = "c.tinyauth.app"
-	result, err = utils.GetCookieDomain(domain)
+	result, err = utils.GetCookieDomain(domain, true)
 	assert.NoError(t, err)
 	assert.Equal(t, expected, result)
 
 	// Invalid domain (only TLD)
 	domain = "com"
-	_, err = utils.GetCookieDomain(domain)
-	assert.ErrorContains(t, err, "invalid app url, must be at least second level domain")
+	_, err = utils.GetCookieDomain(domain, true)
+	assert.EqualError(t, err, "invalid app url, must be in format subdomain.domain.tld or domain.tld")
 
 	// IP address
 	domain = "http://10.10.10.10"
-	_, err = utils.GetCookieDomain(domain)
+	_, err = utils.GetCookieDomain(domain, true)
 	assert.ErrorContains(t, err, "ip addresses not allowed")
 
 	// Invalid URL
 	domain = "http://[::1]:namedport"
-	_, err = utils.GetCookieDomain(domain)
+	_, err = utils.GetCookieDomain(domain, true)
 	assert.ErrorContains(t, err, "parse \"http://[::1]:namedport\": invalid port \":namedport\" after host")
 
 	// URL with scheme and path
 	domain = "https://sub.tinyauth.app/path"
 	expected = "tinyauth.app"
-	result, err = utils.GetCookieDomain(domain)
+	result, err = utils.GetCookieDomain(domain, true)
 	assert.NoError(t, err)
 	assert.Equal(t, expected, result)
 
 	// URL with port
 	domain = "http://sub.tinyauth.app:8080"
-	expected = "tinyauth.app"
-	result, err = utils.GetCookieDomain(domain)
+	expected = "tinyauth.app:8080"
+	result, err = utils.GetCookieDomain(domain, true)
 	assert.NoError(t, err)
 	assert.Equal(t, expected, result)
 
 	// Domain managed by ICANN
 	domain = "http://example.co.uk"
-	_, err = utils.GetCookieDomain(domain)
+	_, err = utils.GetCookieDomain(domain, true)
 	assert.Error(t, err, "domain in public suffix list, cannot set cookies")
+
+	// Domain without subdomain
+	domain = "http://tinyauth.app"
+	expected = "tinyauth.app"
+	result, err = utils.GetCookieDomain(domain, true)
+	assert.NoError(t, err)
+	assert.Equal(t, expected, result)
+
+	// Case insensitivity
+	domain = "http://Sub.Tinyauth.App"
+	expected = "tinyauth.app"
+	result, err = utils.GetCookieDomain(domain, true)
+	assert.NoError(t, err)
+	assert.Equal(t, expected, result)
+
+	// Subdomains disabled
+	domain = "http://sub.tinyauth.app"
+	expected = "sub.tinyauth.app"
+	result, err = utils.GetCookieDomain(domain, false)
+	assert.NoError(t, err)
+	assert.Equal(t, expected, result)
 }
 
 func TestParseFileToLine(t *testing.T) {