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

package routers

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/url"
	"recipe-manager/config"
	"recipe-manager/services/oauth"

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

type AuthRouter struct {
	cfg   *config.ServerConfig
	oauth oauth.OAuthService
}

func NewAuthRouter(cfg *config.ServerConfig, oauth oauth.OAuthService) *AuthRouter {
	return &AuthRouter{cfg, oauth}
}

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)
			rand.Read(bytes)
			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")
			}

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

		r.Get("/google/callback", func(w http.ResponseWriter, r *http.Request) {
			// check state

			var redirect_to string
			state := r.URL.Query().Get("state")
			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
				}

				redirect_to = val["redirect_to"]

				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 user info
			user, err := ar.oauth.GetUserInfo(r.Context(), token)

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

			value := url.Values{
				"name":    {user.Name},
				"email":   {user.Email},
				"picture": {user.Picture},
			}

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

			// 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("/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)
		})

		r.Get("/user", func(w http.ResponseWriter, r *http.Request) {

			token := &oauth2.Token{}
			if cookie, err := r.Cookie("access_token"); err == nil {
				token.AccessToken = cookie.Value
			}

			// if have refresh token, set refresh token to token
			if cookie, err := r.Cookie("refresh_token"); err == nil {
				token.RefreshToken = cookie.Value
			}

			if token.AccessToken == "" && token.RefreshToken == "" {
				http.Error(w, "Unauthorized", http.StatusUnauthorized)
				return
			}

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

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

			// return user info
			w.Header().Add("Content-Type", "application/json")
			json.NewEncoder(w).Encode(map[string]interface{}{
				"user": user,
			})
		})

	})
}