Taobin-Recipe-Manager/server/routers/auth.go

package routers

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"recipe-manager/config"
	"recipe-manager/services/logger"
	"recipe-manager/services/oauth"
	"recipe-manager/services/user"
	"strconv"
	"time"

	"github.com/go-chi/chi/v5"
	"go.uber.org/zap"
	"golang.org/x/oauth2"
)

type AuthRouter struct {
	cfg         *config.ServerConfig
	oauth       oauth.OAuthService
	userService user.UserService
	taoLogger   *logger.TaoLogger
}

func (ar *AuthRouter) Route(r chi.Router) {
	r.Route("/auth", func(r chi.Router) {
		r.Get("/google", func(w http.ResponseWriter, r *http.Request) {

			// generate state and nonce
			bytes := make([]byte, 32)
			_, err := rand.Read(bytes)
			if err != nil {
				return
			}
			state := base64.URLEncoding.EncodeToString(bytes)

			stateMap := map[string]string{}

			if r.URL.Query().Get("redirect_to") != "" {
				stateMap["redirect_to"] = r.URL.Query().Get("redirect_to")
			}

			if r.URL.Query().Get("kind") != "" {
				stateMap["kind"] = r.URL.Query().Get("kind")
			}

			authURL := ar.oauth.AuthURL(state, stateMap)
			http.Redirect(w, r, authURL, http.StatusTemporaryRedirect)
		})

		r.Get("/google/callback", func(w http.ResponseWriter, r *http.Request) {
			// check state
			ctx, cancel := context.WithTimeout(r.Context(), time.Second*5)
			defer cancel()

			var redirectTo string
			var kind string
			state := r.URL.Query().Get("state")

			fmt.Println("url query", r.URL.Query())

			if state == "" {
				http.Error(w, "State not found", http.StatusBadRequest)
				return
			} else {
				val, ok := ar.oauth.GetState(state)

				if !ok {
					http.Error(w, "Invalid state", http.StatusBadRequest)
					return
				}

				fmt.Println("val", val)

				redirectTo = val["redirect_to"]
				kind = val["kind"]

				ar.oauth.RemoveState(state)
			}

			// exchange code for token
			token, err := ar.oauth.Exchange(r.Context(), r.FormValue("code"))
			if err != nil {
				http.Error(w, "Error exchanging code for token", http.StatusBadRequest)
				return
			}

			// get userInfo info
			userInfo, err := ar.oauth.GetUserInfo(r.Context(), token)

			if err != nil {
				http.Error(w, "Error getting userInfo info", http.StatusBadRequest)
				return
			}

			// map with database
			userFromDb, err := ar.userService.GetUserByEmail(ctx, userInfo.Email)
			if err != nil {
				http.Error(w, "Error while getting user data from database.", http.StatusInternalServerError)
				return
			}

			if userFromDb == nil {
				http.Error(w, "Unauthorized, We not found your email, Please contact admin.", http.StatusUnauthorized)
				return
			}

			picture := userInfo.Picture
			if userFromDb.Picture != "" {
				picture = userFromDb.Picture
			}

			value := url.Values{
				"id":          {userFromDb.ID},
				"name":        {userFromDb.Name},
				"email":       {userInfo.Email},
				"picture":     {picture},
				"permissions": {strconv.Itoa(int(userFromDb.Permissions))},
			}

			if redirectTo != "" {
				value.Add("redirect_to", redirectTo)
			}

			ar.taoLogger.Log.Info("User Log-In Success", zap.String("userInfo", userInfo.Name), zap.String("email", userInfo.Email))

			if kind == "electron" {
				value.Add("access_token", token.AccessToken)
				value.Add("max_age", "3600")
				value.Add("refresh_token", token.RefreshToken)
				http.Redirect(w, r, ar.cfg.ClientElectronRedirectURL+"login?"+value.Encode(), http.StatusTemporaryRedirect)
				return
			}
			// redirect to frontend with token and refresh token
			w.Header().Add("set-cookie", "access_token="+token.AccessToken+"; Path=/; HttpOnly; SameSite=None; Secure; Max-Age=3600")
			w.Header().Add("set-cookie", "refresh_token="+token.RefreshToken+"; Path=/; HttpOnly; SameSite=None; Secure")
			http.Redirect(w, r, ar.cfg.ClientRedirectURL+"/?"+value.Encode(), http.StatusTemporaryRedirect)
		})

		r.Get("/me", func(w http.ResponseWriter, r *http.Request) {
			// get access token from cookie
			cookie, err := r.Cookie("access_token")
			if err != nil {
				http.Error(w, "Access token not found", http.StatusBadRequest)
				return
			}

			// get userInfo info
			userInfo, err := ar.oauth.GetUserInfo(r.Context(), &oauth2.Token{AccessToken: cookie.Value})
			if err != nil {
				http.Error(w, "Error getting userInfo info", http.StatusBadRequest)
				return
			}

			// map with database
			userFromDb, err := ar.userService.GetUserByEmail(r.Context(), userInfo.Email)
			if err != nil {
				http.Error(w, "Error while getting user data from database.", http.StatusInternalServerError)
				return
			}

			if userFromDb == nil {
				http.Error(w, "Unauthorized, We not found your email, Please contact admin.", http.StatusUnauthorized)
				return
			}

			picture := userInfo.Picture
			if userFromDb.Picture != "" {
				picture = userFromDb.Picture
			}

			value := url.Values{
				"id":          {userFromDb.ID},
				"name":        {userFromDb.Name},
				"email":       {userInfo.Email},
				"picture":     {picture},
				"permissions": {strconv.Itoa(int(userFromDb.Permissions))},
			}

			json.NewEncoder(w).Encode(value)
		})

		r.Get("/refresh", func(w http.ResponseWriter, r *http.Request) {
			// get refresh token from query string
			refreshToken := r.URL.Query().Get("refresh_token")
			redirectTo := r.URL.Query().Get("redirect_to")
			if refreshToken == "" {
				http.Error(w, "Refresh token not found", http.StatusBadRequest)
				return
			}

			// exchange refresh token for new token
			token, err := ar.oauth.RefreshToken(r.Context(), &oauth2.Token{RefreshToken: refreshToken})
			if err != nil {
				http.Error(w, "Error exchanging refresh token for token", http.StatusBadRequest)
				return
			}

			// redirect to frontend with token and refresh token
			http.Redirect(w, r, ar.cfg.ClientRedirectURL+"/?token="+token.AccessToken+"&redirect_to="+redirectTo, http.StatusTemporaryRedirect)
		})

		r.Get("/revoke", func(w http.ResponseWriter, r *http.Request) {
			// get access token and refresh token from cookie
			if cookie, err := r.Cookie("access_token"); err == nil {
				// request to revoke token at oauth2.googleapis.com/revoke
				_, err := http.PostForm("https://oauth2.googleapis.com/revoke", url.Values{"token": {cookie.Value}})
				if err != nil {
					http.Error(w, "Error revoking token", http.StatusBadRequest)
					return
				}
			}

			// remove cookie with expire from frontend and response no content
			w.Header().Add("set-cookie", "access_token=; Path=/; HttpOnly; SameSite=None; Secure; Max-Age=0")
			w.Header().Add("set-cookie", "refresh_token=; Path=/; HttpOnly; SameSite=None; Secure; Max-Age=0")
			w.WriteHeader(http.StatusNoContent)
		})
	})
}

func NewAuthRouter(cfg *config.ServerConfig, oauth oauth.OAuthService, userService user.UserService, taoLogger *logger.TaoLogger) *AuthRouter {
	return &AuthRouter{cfg, oauth, userService, taoLogger}
}