feat: add psl checks to the oauth controller is safe redirect check

internal/controller/oauth_controller.go

@@ -12,6 +12,7 @@ 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"
@@ -329,16 +330,33 @@ func (controller *OAuthController) isRedirectSafe(redirectURI string) bool {
 		}
 
 		// exact match
-		if u.Host == tu.Host {
+		if strings.EqualFold(u.Host, tu.Host) {
 			return true
 		}
 
-		// subdomain match (trim the tinyauth part)
+		// if subdomains are disabled, end here
+		if !controller.config.Auth.SubdomainsEnabled {
+			continue
+		}
+
+		// 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 strings.HasSuffix(u.Host, "."+root) {
+
+		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
 		}
 	}

internal/controller/oauth_controller_test.go

@@ -16,27 +16,35 @@ func TestOAuthController(t *testing.T) {
 	cfg, runtime := test.CreateTestConfigs(t)
 
 	type testCase struct {
-		description string
+		description       string
-		run         func(ctrl *OAuthController)
+		run               func(ctrl *OAuthController)
+		trustedDomains    []string
+		subdomainsEnabled bool
 	}
 
 	tests := []testCase{
 		{
-			description: "Test exact match of redirect URI",
+			description:       "Test exact match of redirect URI",
+			trustedDomains:    []string{"https://tinyauth.example.com"},
+			subdomainsEnabled: true,
 			run: func(ctrl *OAuthController) {
 				redirectUri := "https://tinyauth.example.com"
 				assert.True(t, ctrl.isRedirectSafe(redirectUri))
 			},
 		},
 		{
-			description: "Test subdomain match of redirect URI",
+			description:       "Test subdomain match of redirect URI",
+			trustedDomains:    []string{"https://tinyauth.example.com"},
+			subdomainsEnabled: true,
 			run: func(ctrl *OAuthController) {
 				redirectUri := "https://sub.example.com"
 				assert.True(t, ctrl.isRedirectSafe(redirectUri))
 			},
 		},
 		{
-			description: "Test different trusted domain",
+			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))
@@ -45,7 +53,7 @@ func TestOAuthController(t *testing.T) {
 		{
 			description: "Test invalid redirect URI",
 			run: func(ctrl *OAuthController) {
-				redirectUri := "https://malicious.com"
+				redirectUri := "https:/malicious"
 				assert.False(t, ctrl.isRedirectSafe(redirectUri))
 			},
 		},
@@ -57,12 +65,79 @@ func TestOAuthController(t *testing.T) {
 			},
 		},
 		{
-			description: "Test redirect URI with different scheme",
+			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"},
+			subdomainsEnabled: false,
+			run: func(ctrl *OAuthController) {
+				redirectUri := "https://sub.tinyauth.example.com"
+				assert.False(t, ctrl.isRedirectSafe(redirectUri))
+			},
+		},
+		{
+			description:       "Test domain like the .co.uk",
+			trustedDomains:    []string{"https://example.co.uk"},
+			subdomainsEnabled: true,
+			run: func(ctrl *OAuthController) {
+				redirectUri := "https://sub.example.co.uk"
+				assert.False(t, ctrl.isRedirectSafe(redirectUri))
+			},
+		},
+		{
+			description:       "Test domain like the .co.uk with subdomains disabled",
+			trustedDomains:    []string{"https://example.co.uk"},
+			subdomainsEnabled: false,
+			run: func(ctrl *OAuthController) {
+				redirectUri := "https://example.co.uk"
+				assert.True(t, ctrl.isRedirectSafe(redirectUri))
+			},
+		},
+		{
+			description:       "Test caps domain",
+			trustedDomains:    []string{"https://TINYAUTH.ExAmpLe.com"},
+			subdomainsEnabled: true,
+			run: func(ctrl *OAuthController) {
+				redirectUri := "https://sUb.ExAmPle.com"
+				assert.True(t, ctrl.isRedirectSafe(redirectUri))
+			},
+		},
+		{
+			description:       "Test edge case with @",
+			trustedDomains:    []string{"https://tinyauth.example.com"},
+			subdomainsEnabled: true,
+			run: func(ctrl *OAuthController) {
+				redirectUri := "https://malicious.example.com@evil.com"
+				assert.False(t, ctrl.isRedirectSafe(redirectUri))
+			},
+		},
 	}
 
 	// TODO: add auth service
@@ -71,6 +146,9 @@ func TestOAuthController(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
+			cfg.Auth.SubdomainsEnabled = tc.subdomainsEnabled
 			ctrl := NewOAuthController(OAuthControllerInput{
 				Log:           log,
 				Config:        &cfg,