ZFile/backend/oauth_google.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strings"

	"github.com/golang-jwt/jwt/v5"
)

const (
	googleOAuthLoginStateCookie = "google_oauth_state"
	googleOAuthLinkStateCookie  = "google_oauth_link_state"
)

type googleTokenResponse struct {
	AccessToken string `json:"access_token"`
}

type googleUserInfo struct {
	Sub           string `json:"sub"`
	Email         string `json:"email"`
	EmailVerified bool   `json:"email_verified"`
}

func (s *Server) handleGoogleAuthStart(w http.ResponseWriter, r *http.Request) {
	if !s.config.GoogleAuthEnabled {
		writeErr(w, http.StatusNotFound, "google auth is disabled")
		return
	}

	state, err := randomToken()
	if err != nil {
		writeErr(w, http.StatusInternalServerError, "failed to initialize oauth")
		return
	}
	setCookie(w, googleOAuthLoginStateCookie, state, 600, s.config.CookieSecure)
	clearCookie(w, googleOAuthLinkStateCookie, s.config.CookieSecure)

	u, err := url.Parse(s.config.GoogleAuthURL)
	if err != nil {
		writeErr(w, http.StatusInternalServerError, "invalid google auth config")
		return
	}
	q := u.Query()
	q.Set("client_id", strings.TrimSpace(s.config.GoogleClientID))
	q.Set("redirect_uri", s.googleRedirectURL(r))
	q.Set("response_type", "code")
	q.Set("scope", "openid email profile")
	q.Set("state", state)
	q.Set("prompt", "select_account")
	u.RawQuery = q.Encode()

	http.Redirect(w, r, u.String(), http.StatusFound)
}

func (s *Server) handleGoogleLinkStart(w http.ResponseWriter, r *http.Request) {
	if !s.config.GoogleAuthEnabled {
		writeErr(w, http.StatusNotFound, "google auth is disabled")
		return
	}

	state, err := randomToken()
	if err != nil {
		writeErr(w, http.StatusInternalServerError, "failed to initialize oauth")
		return
	}
	setCookie(w, googleOAuthLinkStateCookie, state, 600, s.config.CookieSecure)
	clearCookie(w, googleOAuthLoginStateCookie, s.config.CookieSecure)

	u, err := url.Parse(s.config.GoogleAuthURL)
	if err != nil {
		writeErr(w, http.StatusInternalServerError, "invalid google auth config")
		return
	}
	q := u.Query()
	q.Set("client_id", strings.TrimSpace(s.config.GoogleClientID))
	q.Set("redirect_uri", s.googleRedirectURL(r))
	q.Set("response_type", "code")
	q.Set("scope", "openid email profile")
	q.Set("state", state)
	q.Set("prompt", "select_account")
	u.RawQuery = q.Encode()

	http.Redirect(w, r, u.String(), http.StatusFound)
}

func (s *Server) handleGoogleAuthCallback(w http.ResponseWriter, r *http.Request) {
	if !s.config.GoogleAuthEnabled {
		writeErr(w, http.StatusNotFound, "google auth is disabled")
		return
	}

	if oauthErr := strings.TrimSpace(r.URL.Query().Get("error")); oauthErr != "" {
		writeErr(w, http.StatusUnauthorized, "google login was denied")
		return
	}

	code := strings.TrimSpace(r.URL.Query().Get("code"))
	if code == "" {
		writeErr(w, http.StatusBadRequest, "missing oauth code")
		return
	}

	state := strings.TrimSpace(r.URL.Query().Get("state"))
	flow, uid, ok := s.resolveGoogleOAuthFlow(w, r, state)
	if !ok {
		writeErr(w, http.StatusUnauthorized, "invalid oauth state")
		return
	}

	token, err := s.exchangeGoogleCode(r.Context(), code, s.googleRedirectURL(r))
	if err != nil {
		log.Printf("auth.google.failed ip=%q reason=%q", clientIP(r), "token_exchange_failed")
		writeErr(w, http.StatusUnauthorized, "google auth failed")
		return
	}

	info, err := s.fetchGoogleUserInfo(r.Context(), token)
	if err != nil {
		log.Printf("auth.google.failed ip=%q reason=%q", clientIP(r), "userinfo_failed")
		writeErr(w, http.StatusUnauthorized, "google auth failed")
		return
	}

	if flow == "link" {
		if err := s.linkGoogleSubToUser(uid, info.Sub); err != nil {
			log.Printf("auth.google.failed ip=%q reason=%q", clientIP(r), "user_link_failed")
			writeErr(w, http.StatusUnauthorized, "google account link failed")
			return
		}
		log.Printf("auth.google.linked user_id=%d ip=%q", uid, clientIP(r))
		http.Redirect(w, r, "/drive", http.StatusFound)
		return
	}

	user, err := s.findOrCreateGoogleUser(info.Sub, info.Email)
	if err != nil {
		log.Printf("auth.google.failed ip=%q reason=%q", clientIP(r), "user_provision_failed")
		writeErr(w, http.StatusUnauthorized, "google auth failed")
		return
	}

	if err := s.issueUserSession(w, r, user.ID); err != nil {
		log.Printf("auth.google.failed ip=%q reason=%q", clientIP(r), "session_issue_failed")
		writeErr(w, http.StatusInternalServerError, "failed to create session")
		return
	}

	log.Printf("auth.google.success user_id=%d username=%q ip=%q", user.ID, user.Username, clientIP(r))
	http.Redirect(w, r, "/drive", http.StatusFound)
}

func (s *Server) resolveGoogleOAuthFlow(w http.ResponseWriter, r *http.Request, state string) (string, int64, bool) {
	state = strings.TrimSpace(state)
	if state == "" {
		return "", 0, false
	}

	if linkCookie, err := r.Cookie(googleOAuthLinkStateCookie); err == nil && linkCookie != nil && linkCookie.Value != "" {
		if subtleConstantTimeEq(linkCookie.Value, state) == 1 {
			clearCookie(w, googleOAuthLinkStateCookie, s.config.CookieSecure)
			uid, uidErr := s.userIDFromAccessCookie(r)
			if uidErr != nil {
				return "", 0, false
			}
			return "link", uid, true
		}
	}

	if loginCookie, err := r.Cookie(googleOAuthLoginStateCookie); err == nil && loginCookie != nil && loginCookie.Value != "" {
		if subtleConstantTimeEq(loginCookie.Value, state) == 1 {
			clearCookie(w, googleOAuthLoginStateCookie, s.config.CookieSecure)
			return "login", 0, true
		}
	}

	return "", 0, false
}

func (s *Server) userIDFromAccessCookie(r *http.Request) (int64, error) {
	cookie, err := r.Cookie("access_token")
	if err != nil || cookie == nil || cookie.Value == "" {
		return 0, fmt.Errorf("missing access token")
	}

	claims := &AccessClaims{}
	tkn, err := jwt.ParseWithClaims(cookie.Value, claims, func(token *jwt.Token) (any, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method")
		}
		return []byte(s.config.JWTSecret), nil
	})
	if err != nil || !tkn.Valid {
		return 0, fmt.Errorf("invalid access token")
	}
	if claims.UserID <= 0 {
		return 0, fmt.Errorf("invalid access token claims")
	}
	return claims.UserID, nil
}

func (s *Server) linkGoogleSubToUser(userID int64, googleSub string) error {
	googleSub = strings.TrimSpace(googleSub)
	if userID <= 0 || googleSub == "" {
		return fmt.Errorf("invalid link request")
	}

	if _, err := s.findUser(userID); err != nil {
		return err
	}

	if existing, err := s.findUserByGoogleSub(googleSub); err == nil {
		if existing.ID != userID {
			return fmt.Errorf("google account already linked")
		}
		return nil
	} else if !isNoRows(err) {
		return err
	}

	if err := s.orm.updateGoogleSub(userID, googleSub); err != nil {
		if strings.Contains(strings.ToLower(err.Error()), "unique") {
			return fmt.Errorf("google account already linked")
		}
		return err
	}
	return nil
}

func (s *Server) googleRedirectURL(r *http.Request) string {
	if v := strings.TrimSpace(s.config.GoogleRedirectURL); v != "" {
		return v
	}
	return fmt.Sprintf("%s://%s/api/auth/google/callback", schemeOf(r), r.Host)
}

func (s *Server) exchangeGoogleCode(ctx context.Context, code, redirectURI string) (string, error) {
	values := url.Values{}
	values.Set("code", code)
	values.Set("client_id", strings.TrimSpace(s.config.GoogleClientID))
	values.Set("client_secret", s.config.GoogleClientSecret)
	values.Set("redirect_uri", redirectURI)
	values.Set("grant_type", "authorization_code")

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.config.GoogleTokenURL, strings.NewReader(values.Encode()))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("google token endpoint returned %d", resp.StatusCode)
	}

	var out googleTokenResponse
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return "", err
	}
	if strings.TrimSpace(out.AccessToken) == "" {
		return "", fmt.Errorf("google token response missing access_token")
	}
	return out.AccessToken, nil
}

func (s *Server) fetchGoogleUserInfo(ctx context.Context, accessToken string) (googleUserInfo, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.config.GoogleUserInfoURL, nil)
	if err != nil {
		return googleUserInfo{}, err
	}
	req.Header.Set("Authorization", "Bearer "+accessToken)

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return googleUserInfo{}, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return googleUserInfo{}, fmt.Errorf("google userinfo endpoint returned %d", resp.StatusCode)
	}

	var out googleUserInfo
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return googleUserInfo{}, err
	}
	out.Sub = strings.TrimSpace(out.Sub)
	out.Email = strings.ToLower(strings.TrimSpace(out.Email))
	if out.Sub == "" || out.Email == "" || !out.EmailVerified {
		return googleUserInfo{}, fmt.Errorf("google account data is incomplete")
	}
	return out, nil
}

func (s *Server) findOrCreateGoogleUser(googleSub, email string) (User, error) {
	googleSub = strings.TrimSpace(googleSub)
	email = strings.ToLower(strings.TrimSpace(email))
	if googleSub == "" || email == "" {
		return User{}, fmt.Errorf("invalid google identity")
	}

	if user, err := s.findUserByGoogleSub(googleSub); err == nil {
		return user, nil
	} else if !isNoRows(err) {
		return User{}, err
	}

	if user, err := s.findUserByEmail(email); err == nil {
		if err := s.orm.updateGoogleSub(user.ID, googleSub); err != nil {
			return User{}, err
		}
		return user, nil
	} else if !isNoRows(err) {
		return User{}, err
	}

	hash, err := hashPasswordArgon2ID(mustRandomPassword())
	if err != nil {
		return User{}, err
	}

	id, err := s.orm.createUser(email, hash, "dracula", "auto", "zip", &googleSub)
	if err != nil {
		if strings.Contains(strings.ToLower(err.Error()), "unique") {
			if user, findErr := s.findUserByEmail(email); findErr == nil {
				if linkErr := s.orm.updateGoogleSub(user.ID, googleSub); linkErr != nil {
					return User{}, linkErr
				}
				return s.findUser(user.ID)
			}
		}
		return User{}, err
	}
	if err := s.storage.Mkdir(id, "/"); err != nil {
		return User{}, fmt.Errorf("failed to provision user storage: %w", err)
	}
	return User{ID: id, Username: email, Theme: "dracula", ColorMode: "auto", Archive: "zip"}, nil
}

func (s *Server) findUserByGoogleSub(googleSub string) (User, error) {
	return s.orm.findUserByGoogleSub(googleSub)
}

func (s *Server) findUserByEmail(email string) (User, error) {
	return s.orm.findUserByEmail(email)
}

func mustRandomPassword() string {
	tok, err := randomToken()
	if err != nil {
		return "google-oauth-password-fallback"
	}
	if len(tok) < 16 {
		return tok + "-google-oauth"
	}
	return tok
}

func isNoRows(err error) bool {
	return isORMNotFound(err)
}