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

package routers

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

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

type AuthRouter struct {
	cfg     *config.ServerConfig
	gConfig *oauth2.Config
	nonce   map[string]map[string]string
}

func NewAuthRouter(cfg *config.ServerConfig) *AuthRouter {

	file, err := os.Open("client_secret.json")
	if err != nil {
		panic(err)
	}
	defer file.Close()

	var clientSecret map[string]interface{}
	json.NewDecoder(file).Decode(&clientSecret)

	return &AuthRouter{
		cfg: cfg,
		gConfig: &oauth2.Config{
			ClientID:     clientSecret["web"].(map[string]interface{})["client_id"].(string),
			ClientSecret: clientSecret["web"].(map[string]interface{})["client_secret"].(string),
			RedirectURL:  cfg.ServerDomain + "/auth/google/callback",
			Scopes:       []string{"https://www.googleapis.com/auth/userinfo.email", "https://www.googleapis.com/auth/userinfo.profile"},
			Endpoint: oauth2.Endpoint{
				AuthURL:  clientSecret["web"].(map[string]interface{})["auth_uri"].(string),
				TokenURL: clientSecret["web"].(map[string]interface{})["token_uri"].(string),
			},
		},
		nonce: make(map[string]map[string]string),
	}
}

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

			ar.nonce[state] = stateMap

			url := ar.gConfig.AuthCodeURL(state, oauth2.SetAuthURLParam("hd", "forth.co.th"), oauth2.SetAuthURLParam("include_granted_scopes", "true"), oauth2.AccessTypeOffline)
			http.Redirect(w, r, url, http.StatusTemporaryRedirect)
		})

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

			var redirect_to string
			if r.URL.Query().Get("state") == "" {
				http.Error(w, "State not found", http.StatusBadRequest)
				return
			} else {
				val, ok := ar.nonce[r.URL.Query().Get("state")]

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

				redirect_to = val["redirect_to"]

				// remove state from nonce
				delete(ar.nonce, r.URL.Query().Get("state"))
			}

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

			// get user info
			client := ar.gConfig.Client(r.Context(), token)
			resp, err := client.Get("https://www.googleapis.com/oauth2/v3/userinfo")
			if err != nil {
				http.Error(w, "Error getting user info", http.StatusBadRequest)
				return
			}
			defer resp.Body.Close()

			var userInfo map[string]interface{}
			json.NewDecoder(resp.Body).Decode(&userInfo)

			value := url.Values{}

			value.Add("email", userInfo["email"].(string))
			value.Add("name", userInfo["name"].(string))
			value.Add("picture", userInfo["picture"].(string))
			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")
			w.Header().Add("set-cookie", "refresh_token="+token.RefreshToken+"; Path=/; HttpOnly; SameSite=None; Secure")
			http.Redirect(w, r, ar.cfg.ClientRedirectURL+"/callback?"+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.gConfig.TokenSource(r.Context(), &oauth2.Token{RefreshToken: refreshToken}).Token()
			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+"/callback?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
			client := ar.gConfig.Client(r.Context(), token)
			resp, err := client.Get("https://www.googleapis.com/oauth2/v3/userinfo")
			if err != nil {
				http.Error(w, "Error getting user info", http.StatusBadRequest)
				return
			}
			defer resp.Body.Close()

			var userInfo map[string]interface{}
			if err = json.NewDecoder(resp.Body).Decode(&userInfo); err != nil {
				http.Error(w, "Error getting user info", http.StatusBadRequest)
				return
			}

			if userInfo["error"] != nil {
				w.Header().Add("Content-Type", "application/json")
				w.WriteHeader(http.StatusUnauthorized)
				json.NewEncoder(w).Encode(userInfo)
				return
			}

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

	})
}