feat: 绑定 oauth2

2026-05-06 13:48:52 +00:00 · 2024-12-28 23:50:59 +08:00
parent 8554f3eba7
commit 18020939da
15 changed files with 360 additions and 24 deletions
@@ -0,0 +1,240 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"strconv"
+	"strings"
+	"time"
+
+	jwt "github.com/appleboy/gin-jwt/v2"
+	"github.com/gin-gonic/gin"
+	"github.com/patrickmn/go-cache"
+	"github.com/tidwall/gjson"
+	"golang.org/x/oauth2"
+	"gorm.io/gorm"
+
+	"github.com/nezhahq/nezha/model"
+	"github.com/nezhahq/nezha/pkg/utils"
+	"github.com/nezhahq/nezha/service/singleton"
+)
+
+type Oauth2LoginType uint8
+
+const (
+	_ Oauth2LoginType = iota
+	rTypeLogin
+	rTypeBind
+)
+
+func getRedirectURL(c *gin.Context, provider string, rType Oauth2LoginType) string {
+	scheme := "http://"
+	referer := c.Request.Referer()
+	if forwardedProto := c.Request.Header.Get("X-Forwarded-Proto"); forwardedProto == "https" || strings.HasPrefix(referer, "https://") {
+		scheme = "https://"
+	}
+	var suffix string
+	if rType == rTypeLogin {
+		suffix = "/dashboard/login?provider=" + provider
+	} else if rType == rTypeBind {
+		suffix = "/dashboard/profile?provider=" + provider
+	}
+	return scheme + c.Request.Host + suffix
+}
+
+// @Summary Get Oauth2 Redirect URL
+// @Description Get Oauth2 Redirect URL
+// @Produce json
+// @Param provider path string true "provider"
+// @Param type query int false "type" Enums(1, 2) default(1)
+// @Success 200 {object} model.Oauth2LoginResponse
+// @Router /api/v1/oauth2/{provider} [get]
+func oauth2redirect(c *gin.Context) (*model.Oauth2LoginResponse, error) {
+	provider := c.Param("provider")
+	if provider == "" {
+		return nil, singleton.Localizer.ErrorT("provider is required")
+	}
+
+	rTypeInt, err := strconv.Atoi(c.Query("type"))
+	if err != nil {
+		return nil, err
+	}
+
+	o2confRaw, has := singleton.Conf.Oauth2[provider]
+	if !has {
+		return nil, singleton.Localizer.ErrorT("provider not found")
+	}
+	o2conf := o2confRaw.Setup(getRedirectURL(c, provider, Oauth2LoginType(rTypeInt)))
+
+	randomString, err := utils.GenerateRandomString(32)
+	if err != nil {
+		return nil, err
+	}
+	state, stateKey := randomString[:16], randomString[16:]
+	singleton.Cache.Set(fmt.Sprintf("%s%s", model.CacheKeyOauth2State, stateKey), state, cache.DefaultExpiration)
+
+	url := o2conf.AuthCodeURL(state, oauth2.AccessTypeOnline)
+	c.SetCookie("nz-o2s", stateKey, 60*5, "", "", false, false)
+
+	return &model.Oauth2LoginResponse{Redirect: url}, nil
+}
+
+func exchangeOpenId(c *gin.Context, o2confRaw *model.Oauth2Config, provider string, callbackData model.Oauth2Callback) (string, error) {
+	// 验证登录跳转时的 State
+	stateKey, err := c.Cookie("nz-o2s")
+	if err != nil {
+		return "", singleton.Localizer.ErrorT("invalid state key")
+	}
+	state, ok := singleton.Cache.Get(fmt.Sprintf("%s%s", model.CacheKeyOauth2State, stateKey))
+	if !ok || state.(string) != callbackData.State {
+		return "", singleton.Localizer.ErrorT("invalid state key")
+	}
+
+	o2conf := o2confRaw.Setup(getRedirectURL(c, provider, rTypeLogin))
+
+	ctx := context.Background()
+
+	otk, err := o2conf.Exchange(ctx, callbackData.Code)
+	if err != nil {
+		return "", err
+	}
+	oauth2client := o2conf.Client(ctx, otk)
+	resp, err := oauth2client.Get(o2confRaw.UserInfoURL)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return "", err
+	}
+
+	return gjson.Get(string(body), o2confRaw.UserIDPath).String(), nil
+}
+
+// @Summary Oauth2 Callback
+// @Description Oauth2 Callback
+// @Accept json
+// @Produce json
+// @Param provider path string true "provider"
+// @Param body body model.Oauth2Callback true "body"
+// @Success 200 {object} model.LoginResponse
+// @Router /api/v1/oauth2/{provider}/callback [post]
+func oauth2callback(jwtConfig *jwt.GinJWTMiddleware) func(c *gin.Context) (*model.LoginResponse, error) {
+	return func(c *gin.Context) (*model.LoginResponse, error) {
+		provider := c.Param("provider")
+		if provider == "" {
+			return nil, singleton.Localizer.ErrorT("provider is required")
+		}
+
+		o2confRaw, has := singleton.Conf.Oauth2[provider]
+		if !has {
+			return nil, singleton.Localizer.ErrorT("provider not found")
+		}
+		provider = strings.ToLower(provider)
+
+		var callbackData model.Oauth2Callback
+		if err := c.ShouldBind(&callbackData); err != nil {
+			return nil, err
+		}
+
+		realip := c.GetString(model.CtxKeyRealIPStr)
+
+		if callbackData.Code == "" {
+			model.BlockIP(singleton.DB, realip, model.WAFBlockReasonTypeBruteForceOauth2, model.BlockIDToken)
+			return nil, singleton.Localizer.ErrorT("code is required")
+		}
+
+		openId, err := exchangeOpenId(c, o2confRaw, provider, callbackData)
+		if err != nil {
+			model.BlockIP(singleton.DB, realip, model.WAFBlockReasonTypeBruteForceOauth2, model.BlockIDToken)
+			return nil, err
+		}
+
+		var bind model.Oauth2Bind
+		if err := singleton.DB.Where("provider = ? AND open_id = ?", provider, openId).First(&bind).Error; err != nil {
+			return nil, singleton.Localizer.ErrorT("oauth2 user not binded yet")
+		}
+
+		tokenString, expire, err := jwtConfig.TokenGenerator(fmt.Sprintf("%d", bind.UserID))
+		if err != nil {
+			return nil, err
+		}
+
+		jwtConfig.SetCookie(c, tokenString)
+
+		return &model.LoginResponse{Token: tokenString, Expire: expire.Format(time.RFC3339)}, nil
+	}
+}
+
+// @Summary Bind Oauth2
+// @Description Bind Oauth2
+// @Accept json
+// @Produce json
+// @Param provider path string true "provider"
+// @Param body body model.Oauth2Callback true "body"
+// @Success 200 {object} any
+// @Router /api/v1/oauth2/{provider}/bind [post]
+func bindOauth2(c *gin.Context) (any, error) {
+	var bindData model.Oauth2Callback
+	if err := c.ShouldBind(&bindData); err != nil {
+		return nil, err
+	}
+
+	provider := c.Param("provider")
+	o2conf, has := singleton.Conf.Oauth2[provider]
+	if !has {
+		return nil, singleton.Localizer.ErrorT("provider not found")
+	}
+	provider = strings.ToLower(provider)
+
+	openId, err := exchangeOpenId(c, o2conf, provider, bindData)
+	if err != nil {
+		return nil, err
+	}
+
+	u := c.MustGet(model.CtxKeyAuthorizedUser).(*model.User)
+
+	var bind model.Oauth2Bind
+	result := singleton.DB.Where("provider = ? AND open_id = ?", provider, openId).Limit(1).Find(&bind)
+	if result.Error != nil && result.Error != gorm.ErrRecordNotFound {
+		return nil, newGormError("%v", result.Error)
+	}
+	bind.UserID = u.ID
+	bind.Provider = provider
+	bind.OpenID = openId
+	if result.Error == gorm.ErrRecordNotFound {
+		result = singleton.DB.Create(&bind)
+	} else {
+		result = singleton.DB.Save(&bind)
+	}
+	if result.Error != nil {
+		return nil, newGormError("%v", result.Error)
+	}
+	return nil, nil
+}
+
+// @Summary Unbind Oauth2
+// @Description Unbind Oauth2
+// @Accept json
+// @Produce json
+// @Param provider path string true "provider"
+// @Success 200 {object} any
+// @Router /api/v1/oauth2/{provider}/unbind [post]
+func unbindOauth2(c *gin.Context) (any, error) {
+	provider := c.Param("provider")
+	if provider == "" {
+		return nil, singleton.Localizer.ErrorT("provider is required")
+	}
+	_, has := singleton.Conf.Oauth2[provider]
+	if !has {
+		return nil, singleton.Localizer.ErrorT("provider not found")
+	}
+	provider = strings.ToLower(provider)
+	u := c.MustGet(model.CtxKeyAuthorizedUser).(*model.User)
+	if err := singleton.DB.Where("provider = ? AND user_id = ?", provider, u.ID).Delete(&model.Oauth2Bind{}).Error; err != nil {
+		return nil, newGormError("%v", err)
+	}
+	return nil, nil
+}