refactor: remove concurrent listeners and rework cookie logic (#950)

2026-06-23 20:00:24 +00:00 · 2026-06-23 13:35:29 +03:00
parent 2572376686
commit 69f4206f65
20 changed files with 466 additions and 414 deletions
@@ -1,6 +1,8 @@
 package controller

 import (
+	"errors"
+
 	"github.com/tinyauthapp/tinyauth/internal/model"
 	"github.com/tinyauthapp/tinyauth/internal/utils/logger"
 	"go.uber.org/dig"
@@ -58,9 +60,9 @@ type ACRUI struct {
 }

 type ACRApp struct {
-	AppURL         string   `json:"appUrl"`
-	CookieDomain   string   `json:"cookieDomain"`
-	TrustedDomains []string `json:"trustedDomains"`
+	AppURL            string `json:"appUrl"`
+	CookieDomain      string `json:"cookieDomain"`
+	SubdomainsEnabled bool   `json:"subdomainsEnabled"`
 }

 type AppContextResponse struct {
@@ -109,7 +111,9 @@ func (controller *ContextController) userContextHandler(c *gin.Context) {
 	context, err := new(model.UserContext).NewFromGin(c)

 	if err != nil {
-		controller.log.App.Error().Err(err).Msg("Failed to create user context from request")
+		if !errors.Is(err, model.ErrUserContextNotFound) {
+			controller.log.App.Error().Err(err).Msg("Failed to create user context from request")
+		}
 		c.JSON(200, UserContextResponse{
 			Status:  401,
 			Message: "Unauthorized",
@@ -160,9 +164,9 @@ func (controller *ContextController) appContextHandler(c *gin.Context) {
 			WarningsEnabled:       controller.config.UI.WarningsEnabled,
 		},
 		App: ACRApp{
-			AppURL:         controller.runtime.AppURL,
-			CookieDomain:   controller.runtime.CookieDomain,
-			TrustedDomains: controller.runtime.TrustedDomains,
+			AppURL:            controller.runtime.AppURL,
+			CookieDomain:      controller.runtime.CookieDomain,
+			SubdomainsEnabled: controller.config.Auth.SubdomainsEnabled,
 		},
 	})
 }
@@ -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)
@@ -12,7 +12,6 @@ import (
 	"github.com/tinyauthapp/tinyauth/internal/service"
 	"github.com/tinyauthapp/tinyauth/internal/utils"
 	"github.com/tinyauthapp/tinyauth/internal/utils/logger"
-	"github.com/weppos/publicsuffix-go/publicsuffix"
 	"go.uber.org/dig"

 	"github.com/gin-gonic/gin"
@@ -305,8 +304,8 @@ func (controller *OAuthController) isOidcRequest(params service.OAuthCallbackPar
 }

 func (controller *OAuthController) getCookieDomain() string {
-	if controller.config.Auth.SubdomainsEnabled {
-		return "." + controller.runtime.CookieDomain
+	if !controller.config.Auth.SubdomainsEnabled {
+		return ""
 	}
 	return controller.runtime.CookieDomain
 }
@@ -314,51 +313,53 @@ func (controller *OAuthController) getCookieDomain() string {
 func (controller *OAuthController) isRedirectSafe(redirectURI string) bool {
 	u, err := url.Parse(redirectURI)

-	if err != nil || u.Host == "" || u.Scheme == "" {
+	if err != nil {
+		controller.log.App.Error().Err(err).Msg("Failed to parse redirect URI")
 		return false
 	}

-	for _, allowed := range controller.runtime.TrustedDomains {
-		tu, err := url.Parse(allowed)
-		if err != nil {
-			controller.log.App.Error().Err(err).Str("allowed", allowed).Msg("Failed to parse trusted domain")
-			continue
+	if u.Scheme == "" || u.Host == "" {
+		controller.log.App.Warn().Msg("Redirect URI has invalid scheme or host")
+		return false
+	}
+
+	au, err := url.Parse(controller.runtime.AppURL)
+
+	if err != nil {
+		controller.log.App.Error().Err(err).Msg("Failed to parse app URL")
+		return false
+	}
+
+	if u.Scheme != au.Scheme {
+		controller.log.App.Warn().Msg("Redirect URI scheme does not match app URL scheme")
+		return false
+	}
+
+	getEffectivePort := func(u *url.URL) string {
+		if u.Port() != "" {
+			return u.Port()
 		}
-
-		if tu.Scheme != u.Scheme {
-			continue
+		if u.Scheme == "https" {
+			return "443"
 		}
+		return "80"
+	}

-		// exact match
-		if strings.EqualFold(u.Host, tu.Host) {
-			return true
-		}
+	if getEffectivePort(u) != getEffectivePort(au) {
+		controller.log.App.Warn().Msg("Redirect URI port does not match app URL port")
+		return false
+	}

-		// if subdomains are disabled, end here
-		if !controller.config.Auth.SubdomainsEnabled {
-			continue
-		}
+	if strings.EqualFold(u.Hostname(), au.Hostname()) {
+		return true
+	}

-		// get the root domain (e.g. tinyauth.example.com -> example.com or
-		// tinyauth.sub.example.com -> sub.example.com)
-		_, root, ok := strings.Cut(tu.Host, ".")
-		if !ok {
-			continue
-		}
+	if !controller.config.Auth.SubdomainsEnabled {
+		return false
+	}

-		root = strings.ToLower(root)
-
-		// check if the root domain is in the psl
-		_, err = publicsuffix.DomainFromListWithOptions(publicsuffix.DefaultList, root, nil)
-
-		if err != nil {
-			continue
-		}
-
-		// subdomain match
-		if strings.HasSuffix(strings.ToLower(u.Host), "."+root) {
-			return true
-		}
+	if strings.HasSuffix(strings.ToLower(u.Hostname()), "."+strings.ToLower(controller.runtime.CookieDomain)) {
+		return true
 	}

 	return false
@@ -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))
 		})
 	}
 }
@@ -295,6 +295,14 @@ func (controller *UserController) totpHandler(c *gin.Context) {
 	context, err := new(model.UserContext).NewFromGin(c)

 	if err != nil {
+		if errors.Is(err, model.ErrUserContextNotFound) {
+			controller.log.App.Warn().Msg("TOTP verification attempt without user context")
+			c.JSON(401, gin.H{
+				"status":  401,
+				"message": "Unauthorized",
+			})
+			return
+		}
 		controller.log.App.Error().Err(err).Msg("Failed to create user context from request for TOTP verification")
 		c.JSON(500, gin.H{
 			"status":  500,
@@ -405,6 +413,14 @@ func (controller *UserController) tailscaleHandler(c *gin.Context) {
 	context, err := new(model.UserContext).NewFromGin(c)

 	if err != nil {
+		if errors.Is(err, model.ErrUserContextNotFound) {
+			controller.log.App.Warn().Msg("Tailscale login attempt without user context")
+			c.JSON(401, gin.H{
+				"status":  401,
+				"message": "Unauthorized",
+			})
+			return
+		}
 		controller.log.App.Error().Err(err).Msg("Failed to create user context from request")
 		c.JSON(401, gin.H{
 			"status":  401,