ZFile/backend/oauth_google.go

package main

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

const googleOAuthStateCookie = "google_oauth_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, googleOAuthStateCookie, state, 600, 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"))
	stateCookie, err := r.Cookie(googleOAuthStateCookie)
	if err != nil || stateCookie == nil || stateCookie.Value == "" || subtleConstantTimeEq(stateCookie.Value, state) == 0 {
		writeErr(w, http.StatusUnauthorized, "invalid oauth state")
		return
	}
	clearCookie(w, googleOAuthStateCookie, s.config.CookieSecure)

	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
	}

	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) 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)
}