Merge pull request #2745 from mehunk/feat/custom-stripe-url

feat: Support customizing the success and cancel url of Stripe.

common/init.go

@@ -159,4 +159,17 @@ func initConstantEnv() {
 		}
 		constant.TaskPricePatches = taskPricePatches
 	}
+
+	// Initialize trusted redirect domains for URL validation
+	trustedDomainsStr := GetEnvOrDefaultString("TRUSTED_REDIRECT_DOMAINS", "")
+	var trustedDomains []string
+	domains := strings.Split(trustedDomainsStr, ",")
+	for _, domain := range domains {
+		trimmedDomain := strings.TrimSpace(domain)
+		if trimmedDomain != "" {
+			// Normalize domain to lowercase
+			trustedDomains = append(trustedDomains, strings.ToLower(trimmedDomain))
+		}
+	}
+	constant.TrustedRedirectDomains = trustedDomains
 }

common/url_validator.go

@@ -0,0 +1,39 @@
+package common
+
+import (
+	"fmt"
+	"net/url"
+	"strings"
+
+	"github.com/QuantumNous/new-api/constant"
+)
+
+// ValidateRedirectURL validates that a redirect URL is safe to use.
+// It checks that:
+//   - The URL is properly formatted
+//   - The scheme is either http or https
+//   - The domain is in the trusted domains list (exact match or subdomain)
+//
+// Returns nil if the URL is valid and trusted, otherwise returns an error
+// describing why the validation failed.
+func ValidateRedirectURL(rawURL string) error {
+	// Parse the URL
+	parsedURL, err := url.Parse(rawURL)
+	if err != nil {
+		return fmt.Errorf("invalid URL format: %s", err.Error())
+	}
+
+	if parsedURL.Scheme != "http" && parsedURL.Scheme != "https" {
+		return fmt.Errorf("invalid URL scheme: only http and https are allowed")
+	}
+
+	domain := strings.ToLower(parsedURL.Hostname())
+
+	for _, trustedDomain := range constant.TrustedRedirectDomains {
+		if domain == trustedDomain || strings.HasSuffix(domain, "."+trustedDomain) {
+			return nil
+		}
+	}
+
+	return fmt.Errorf("domain %s is not in the trusted domains list", domain)
+}

common/url_validator_test.go

@@ -0,0 +1,134 @@
+package common
+
+import (
+	"testing"
+
+	"github.com/QuantumNous/new-api/constant"
+)
+
+func TestValidateRedirectURL(t *testing.T) {
+	// Save original trusted domains and restore after test
+	originalDomains := constant.TrustedRedirectDomains
+	defer func() {
+		constant.TrustedRedirectDomains = originalDomains
+	}()
+
+	tests := []struct {
+		name           string
+		url            string
+		trustedDomains []string
+		wantErr        bool
+		errContains    string
+	}{
+		// Valid cases
+		{
+			name:           "exact domain match with https",
+			url:            "https://example.com/success",
+			trustedDomains: []string{"example.com"},
+			wantErr:        false,
+		},
+		{
+			name:           "exact domain match with http",
+			url:            "http://example.com/callback",
+			trustedDomains: []string{"example.com"},
+			wantErr:        false,
+		},
+		{
+			name:           "subdomain match",
+			url:            "https://sub.example.com/success",
+			trustedDomains: []string{"example.com"},
+			wantErr:        false,
+		},
+		{
+			name:           "case insensitive domain",
+			url:            "https://EXAMPLE.COM/success",
+			trustedDomains: []string{"example.com"},
+			wantErr:        false,
+		},
+
+		// Invalid cases - untrusted domain
+		{
+			name:           "untrusted domain",
+			url:            "https://evil.com/phishing",
+			trustedDomains: []string{"example.com"},
+			wantErr:        true,
+			errContains:    "not in the trusted domains list",
+		},
+		{
+			name:           "suffix attack - fakeexample.com",
+			url:            "https://fakeexample.com/success",
+			trustedDomains: []string{"example.com"},
+			wantErr:        true,
+			errContains:    "not in the trusted domains list",
+		},
+		{
+			name:           "empty trusted domains list",
+			url:            "https://example.com/success",
+			trustedDomains: []string{},
+			wantErr:        true,
+			errContains:    "not in the trusted domains list",
+		},
+
+		// Invalid cases - scheme
+		{
+			name:           "javascript scheme",
+			url:            "javascript:alert('xss')",
+			trustedDomains: []string{"example.com"},
+			wantErr:        true,
+			errContains:    "invalid URL scheme",
+		},
+		{
+			name:           "data scheme",
+			url:            "data:text/html,<script>alert('xss')</script>",
+			trustedDomains: []string{"example.com"},
+			wantErr:        true,
+			errContains:    "invalid URL scheme",
+		},
+
+		// Edge cases
+		{
+			name:           "empty URL",
+			url:            "",
+			trustedDomains: []string{"example.com"},
+			wantErr:        true,
+			errContains:    "invalid URL scheme",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Set up trusted domains for this test case
+			constant.TrustedRedirectDomains = tt.trustedDomains
+
+			err := ValidateRedirectURL(tt.url)
+
+			if tt.wantErr {
+				if err == nil {
+					t.Errorf("ValidateRedirectURL(%q) expected error containing %q, got nil", tt.url, tt.errContains)
+					return
+				}
+				if tt.errContains != "" && !contains(err.Error(), tt.errContains) {
+					t.Errorf("ValidateRedirectURL(%q) error = %q, want error containing %q", tt.url, err.Error(), tt.errContains)
+				}
+			} else {
+				if err != nil {
+					t.Errorf("ValidateRedirectURL(%q) unexpected error: %v", tt.url, err)
+				}
+			}
+		})
+	}
+}
+
+func contains(s, substr string) bool {
+	return len(s) >= len(substr) && (s == substr || len(substr) == 0 ||
+		(len(s) > 0 && len(substr) > 0 && findSubstring(s, substr)))
+}
+
+func findSubstring(s, substr string) bool {
+	for i := 0; i <= len(s)-len(substr); i++ {
+		if s[i:i+len(substr)] == substr {
+			return true
+		}
+	}
+	return false
+}