added fetching of origin in auth login

2026-03-05 10:09:12 +08:00
parent 8f51faeb12
commit 30c91cf5c8
6 changed files with 319 additions and 76 deletions
@@ -27,9 +27,12 @@ import (
 )

 var googleOauthConfig oauth2.Config
-var oauthStateString = generateRandomState()
 var AuthorizationURL string
-var FetchedRedirectURI *string
+
+const (
+	oauthStateCookieName       = "oauth_state"
+	oauthRedirectURICookieName = "oauth_redirect_uri"
+)

 func isTestEnvironment() bool {
 	return flag.Lookup("test.v") != nil || strings.Contains(os.Args[0], ".test")
@@ -106,29 +109,41 @@ func generateRandomState() string {
 }

 func GoogleLogin(w http.ResponseWriter, r *http.Request) {
-
-	helper.LogInfo(fmt.Sprintf("Generated oauth_state: %s", oauthStateString))
-
 	isSecure := strings.HasPrefix(os.Getenv("BACKEND_URL"), HTTPS)
+	state := generateRandomState()
+	helper.LogInfo(fmt.Sprintf("Generated oauth_state: %s", state))
+
+	redirectURI := strings.TrimSpace(r.URL.Query().Get("redirect_uri"))
+	if redirectURI == "" {
+		helper.RespondWithError(w, http.StatusBadRequest, "redirect_uri is required")
+		return
+	}
+
+	if !IsAllowedRedirectURI(redirectURI) {
+		helper.RespondWithError(w, http.StatusUnauthorized, "Unauthorized RedirectURI")
+		return
+	}

 	http.SetCookie(w, &http.Cookie{
-		Name:     "oauth_state",
-		Value:    oauthStateString,
+		Name:     oauthStateCookieName,
+		Value:    state,
+		Path:     "/",
+		HttpOnly: true,
+		Secure:   isSecure,
+		SameSite: http.SameSiteLaxMode,
+		Expires:  time.Now().Add(5 * time.Minute),
+	})
+	http.SetCookie(w, &http.Cookie{
+		Name:     oauthRedirectURICookieName,
+		Value:    redirectURI,
 		Path:     "/",
 		HttpOnly: true,
 		Secure:   isSecure,
 		SameSite: http.SameSiteLaxMode,
 		Expires:  time.Now().Add(5 * time.Minute),
 	})
-	redirectURI := r.URL.Query().Get("redirect_uri")
-	if redirectURI != "" {
-		FetchedRedirectURI = &redirectURI
-		log.Print("FetchedRedirectURI set to: ", *FetchedRedirectURI)
-	} else {
-		FetchedRedirectURI = nil
-	}

-	url := googleOauthConfig.AuthCodeURL(oauthStateString, oauth2.AccessTypeOffline, oauth2.ApprovalForce)
+	url := googleOauthConfig.AuthCodeURL(state, oauth2.AccessTypeOffline, oauth2.ApprovalForce)
 	http.Redirect(w, r, url, http.StatusFound)
 }

@@ -184,6 +199,11 @@ func GoogleCallback(w http.ResponseWriter, r *http.Request) {
 	}
 	helper.LogInfo(fmt.Sprintf("[oauth-debug] state validation ok duration_ms=%d", time.Since(stateStart).Milliseconds()))

+	redirectURI, ok := callbackRedirectURI(w, r)
+	if !ok {
+		return
+	}
+
 	googleUserInfoStart := time.Now()
 	userInfo, err := FetchGoogleUserInfo(w, r)
 	if err != nil {
@@ -232,24 +252,8 @@ func GoogleCallback(w http.ResponseWriter, r *http.Request) {

 	if !emailExists {
 		helper.LogError(errors.New("unregistered email"), "Google login attempt with unregistered email: "+email)
-		if FetchedRedirectURI != nil && *FetchedRedirectURI != "" {
-			RedirectURI := *FetchedRedirectURI
-			log.Print("RedirectURI from query param: ", RedirectURI)
-			if !IsAllowedRedirectURI(RedirectURI) {
-				helper.LogError(errors.New("unauthorized redirect uri"), "Blocked redirect URI for unregistered email: "+RedirectURI)
-				helper.RespondWithError(w, http.StatusUnauthorized, "Unauthorized RedirectURI")
-				log.Print("Unauthorized RedirectURI: ", RedirectURI)
-				return
-			}
-			log.Print("Valid redirect_uri: ", RedirectURI)
-			RedirectURL := fmt.Sprintf("%s/callback?error=%s=", RedirectURI, "unregistered_email")
-			http.Redirect(w, r, RedirectURL, http.StatusSeeOther)
-			return
-		}
-
-		log.Print("No redirect_uri provided, returning JSON response")
-		// No redirect_uri provided, return JSON response
-		helper.RespondWithError(w, http.StatusUnauthorized, "Your email is not registered in the system. Please contact your administrator to request access.")
+		RedirectURL := fmt.Sprintf("%s/callback?error=%s=", redirectURI, "unregistered_email")
+		http.Redirect(w, r, RedirectURL, http.StatusSeeOther)
 		return
 	}

@@ -322,33 +326,12 @@ func GoogleCallback(w http.ResponseWriter, r *http.Request) {

 	helper.LogInfo("Copy this access token: " + accessToken)

-	if FetchedRedirectURI != nil && *FetchedRedirectURI != "" {
-		RedirectURI := *FetchedRedirectURI
-		log.Print("RedirectURI from query param: ", RedirectURI)
-		if !IsAllowedRedirectURI(RedirectURI) {
-			helper.LogError(errors.New("unauthorized redirect uri"), "Blocked redirect URI after successful auth: "+RedirectURI)
-			helper.RespondWithError(w, http.StatusUnauthorized, "Unauthorized RedirectURI")
-			log.Print("Unauthorized RedirectURI: ", RedirectURI)
-			return
-		}
-		log.Print("Valid redirect_uri: ", RedirectURI)
-		RedirectURL := fmt.Sprintf("%s/callback?token=%s&user_id=%s", RedirectURI, accessToken, userID)
-		helper.LogInfo(fmt.Sprintf("[oauth-debug] callback complete redirect=true total_ms=%d", time.Since(callbackStart).Milliseconds()))
-		http.Redirect(w, r, RedirectURL, http.StatusSeeOther)
-		return
-	}
-
-	log.Print("No redirect_uri provided, returning JSON response")
-	// No redirect_uri provided, return JSON response
-	helper.LogInfo(fmt.Sprintf("[oauth-debug] callback complete redirect=false total_ms=%d", time.Since(callbackStart).Milliseconds()))
-	helper.RespondWithJSON(w, http.StatusOK, map[string]string{
-		"message":      "Authentication successful",
-		"access_token": accessToken,
-	})
+	RedirectURL := fmt.Sprintf("%s/callback?token=%s&user_id=%s", redirectURI, accessToken, userID)
+	http.Redirect(w, r, RedirectURL, http.StatusSeeOther)
 }

 func validateState(w http.ResponseWriter, r *http.Request) bool {
-	cookie, err := r.Cookie("oauth_state")
+	cookie, err := r.Cookie(oauthStateCookieName)
 	callbackState := r.URL.Query().Get("state")
 	if err != nil {
 		helper.LogError(err, "oauth_state cookie missing or unreadable during callback")
@@ -357,6 +340,12 @@ func validateState(w http.ResponseWriter, r *http.Request) bool {
 		return false
 	}

+	if strings.TrimSpace(callbackState) == "" {
+		helper.LogWarn(errorInvalidState)
+		helper.RespondWithError(w, http.StatusUnauthorized, errorInvalidState)
+		return false
+	}
+
 	if callbackState != cookie.Value {
 		helper.LogError(errors.New("oauth state mismatch"), fmt.Sprintf("OAuth state mismatch. cookie_state=%s callback_state=%s", cookie.Value, callbackState))
 		helper.LogWarn(errorInvalidState)
@@ -367,6 +356,23 @@ func validateState(w http.ResponseWriter, r *http.Request) bool {
 	return true
 }

+func callbackRedirectURI(w http.ResponseWriter, r *http.Request) (string, bool) {
+	cookie, err := r.Cookie(oauthRedirectURICookieName)
+	if err != nil {
+		helper.LogError(err, "oauth redirect_uri cookie missing or unreadable during callback")
+		helper.RespondWithError(w, http.StatusUnauthorized, "Unauthorized RedirectURI")
+		return "", false
+	}
+
+	redirectURI := strings.TrimSpace(cookie.Value)
+	if redirectURI == "" || !IsAllowedRedirectURI(redirectURI) {
+		helper.RespondWithError(w, http.StatusUnauthorized, "Unauthorized RedirectURI")
+		return "", false
+	}
+
+	return redirectURI, true
+}
+
 func FetchGoogleUserInfo(w http.ResponseWriter, r *http.Request) (models.UserGoogleInfo, error) {
 	fetchStart := time.Now()
 	code := r.URL.Query().Get("code")
@@ -0,0 +1,97 @@
+package handlers
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"testing"
+)
+
+func TestGoogleLogin_RequiresRedirectURI(t *testing.T) {
+	original := os.Getenv("ALLOWED_REDIRECT_URIS")
+	os.Setenv("ALLOWED_REDIRECT_URIS", "http://localhost:5173")
+	defer os.Setenv("ALLOWED_REDIRECT_URIS", original)
+
+	req := httptest.NewRequest(http.MethodGet, "/v1/auth/login", nil)
+	recorder := httptest.NewRecorder()
+
+	GoogleLogin(recorder, req)
+
+	if recorder.Code != http.StatusBadRequest {
+		t.Fatalf("expected status %d, got %d", http.StatusBadRequest, recorder.Code)
+	}
+}
+
+func TestGoogleLogin_RejectsUnauthorizedRedirectURI(t *testing.T) {
+	original := os.Getenv("ALLOWED_REDIRECT_URIS")
+	os.Setenv("ALLOWED_REDIRECT_URIS", "http://localhost:5173")
+	defer os.Setenv("ALLOWED_REDIRECT_URIS", original)
+
+	req := httptest.NewRequest(http.MethodGet, "/v1/auth/login?redirect_uri=http://malicious.example", nil)
+	recorder := httptest.NewRecorder()
+
+	GoogleLogin(recorder, req)
+
+	if recorder.Code != http.StatusUnauthorized {
+		t.Fatalf("expected status %d, got %d", http.StatusUnauthorized, recorder.Code)
+	}
+}
+
+func TestValidateState_MissingCookie(t *testing.T) {
+	req := httptest.NewRequest(http.MethodGet, "/v1/auth/callback?state=test-state", nil)
+	recorder := httptest.NewRecorder()
+
+	ok := validateState(recorder, req)
+	if ok {
+		t.Fatal("expected validateState to return false when oauth_state cookie is missing")
+	}
+	if recorder.Code != http.StatusUnauthorized {
+		t.Fatalf("expected status %d, got %d", http.StatusUnauthorized, recorder.Code)
+	}
+}
+
+func TestValidateState_Success(t *testing.T) {
+	req := httptest.NewRequest(http.MethodGet, "/v1/auth/callback?state=test-state", nil)
+	req.AddCookie(&http.Cookie{Name: oauthStateCookieName, Value: "test-state"})
+	recorder := httptest.NewRecorder()
+
+	ok := validateState(recorder, req)
+	if !ok {
+		t.Fatal("expected validateState to return true for matching state")
+	}
+}
+
+func TestCallbackRedirectURI_MissingCookie(t *testing.T) {
+	original := os.Getenv("ALLOWED_REDIRECT_URIS")
+	os.Setenv("ALLOWED_REDIRECT_URIS", "http://localhost:5173")
+	defer os.Setenv("ALLOWED_REDIRECT_URIS", original)
+
+	req := httptest.NewRequest(http.MethodGet, "/v1/auth/callback?state=test-state", nil)
+	recorder := httptest.NewRecorder()
+
+	_, ok := callbackRedirectURI(recorder, req)
+	if ok {
+		t.Fatal("expected callbackRedirectURI to return false when redirect cookie is missing")
+	}
+	if recorder.Code != http.StatusUnauthorized {
+		t.Fatalf("expected status %d, got %d", http.StatusUnauthorized, recorder.Code)
+	}
+}
+
+func TestCallbackRedirectURI_Success(t *testing.T) {
+	original := os.Getenv("ALLOWED_REDIRECT_URIS")
+	os.Setenv("ALLOWED_REDIRECT_URIS", "http://localhost:5173")
+	defer os.Setenv("ALLOWED_REDIRECT_URIS", original)
+
+	req := httptest.NewRequest(http.MethodGet, "/v1/auth/callback?state=test-state", nil)
+	req.AddCookie(&http.Cookie{Name: oauthRedirectURICookieName, Value: "http://localhost:5173"})
+	recorder := httptest.NewRecorder()
+
+	uri, ok := callbackRedirectURI(recorder, req)
+	if !ok {
+		t.Fatal("expected callbackRedirectURI to return true for allowed redirect URI")
+	}
+	if uri != "http://localhost:5173" {
+		t.Fatalf("expected redirect URI %q, got %q", "http://localhost:5173", uri)
+	}
+}