recommender/internal/service/auth/auth.go

package auth

import (
	"context"
	"leafdev.top/Ecosystem/recommender/internal/base/conf"
	"leafdev.top/Ecosystem/recommender/internal/base/logger"
	"leafdev.top/Ecosystem/recommender/internal/dao"
	"leafdev.top/Ecosystem/recommender/internal/entity"
	"leafdev.top/Ecosystem/recommender/internal/schema"
	"leafdev.top/Ecosystem/recommender/internal/service/jwks"
	"leafdev.top/Ecosystem/recommender/pkg/consts"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/mitchellh/mapstructure"
)

type Service struct {
	config *conf.Config
	jwks   *jwks.JWKS
	logger *logger.Logger
	q      *dao.Query
}

func NewAuthService(config *conf.Config, jwks *jwks.JWKS, logger *logger.Logger, q *dao.Query) *Service {
	return &Service{
		config: config,
		jwks:   jwks,
		logger: logger,
		q:      q,
	}
}

func (a *Service) GetBearerToken(c *gin.Context) (string, error) {
	authorization := c.Request.Header.Get(consts.AuthHeader)

	if authorization == "" {
		return "", consts.ErrBearerTokenNotFound
	}

	authSplit := strings.Split(authorization, " ")
	if len(authSplit) != 2 {
		return "", consts.ErrNotBearerType
	}

	if authSplit[0] != consts.AuthPrefix {
		return "", consts.ErrNotBearerType
	}

	return authSplit[1], nil
}

func (a *Service) GinMiddlewareAuth(tokenType schema.JWTTokenTypes, c *gin.Context) (*schema.User, error) {
	if a.config.Debug.Enabled {
		return a.parseUserJWT(tokenType, "")
	}

	token, err := a.GetBearerToken(c)

	if err != nil {
		return nil, err
	}

	return a.parseUserJWT(tokenType, token)
}

func (a *Service) AuthFromToken(tokenType schema.JWTTokenTypes, token string) (*schema.User, error) {
	if a.config.Debug.Enabled {
		return a.parseUserJWT(tokenType, "")
	}

	return a.parseUserJWT(tokenType, token)
}

func (a *Service) GetUserFromIdToken(idToken string) (*schema.User, error) {
	return a.parseUserJWT(schema.JWTIDToken, idToken)
}

func (a *Service) GinUser(c *gin.Context) *schema.User {
	user, _ := c.Get(consts.AuthMiddlewareKeyUser.String())
	return user.(*schema.User)
}

func (a *Service) GetUserId(ctx context.Context) schema.UserId {
	user := a.GetUser(ctx)

	return user.Token.Sub
}

func (a *Service) GetUser(ctx context.Context) *schema.User {
	user := ctx.Value(consts.AuthMiddlewareKeyUser.String())

	user2 := user.(*schema.User)
	user2.ID = user2.Token.Sub

	return user2
}

func (a *Service) SetUser(ctx context.Context, user *schema.User) context.Context {
	return context.WithValue(ctx, consts.AuthMiddlewareKeyUser.String(), user)
}

func (a *Service) parseUserJWT(tokenType schema.JWTTokenTypes, jwtToken string) (*schema.User, error) {
	var sub = consts.AnonymousUser
	var jwtIdToken = &schema.User{}

	if a.config.Debug.Enabled {
		jwtIdToken.Token.Sub = sub
		jwtIdToken.Valid = true
		return jwtIdToken, nil
	} else {
		token, err := a.jwks.ParseJWT(jwtToken)
		if err != nil {
			return nil, consts.ErrNotValidToken
		}

		subStr, err := token.Claims.GetSubject()
		if err != nil {
			return nil, consts.ErrNotValidToken
		}

		//subInt, err := strconv.Atoi(subStr)
		//if err != nil {
		//	return nil, consts.ErrNotValidToken
		//}

		sub = schema.UserId(subStr)

		// 如果 token.Header 中没有 typ
		if token.Header["typ"] == "" {
			return nil, consts.ErrEmptyResponse
		}

		// 验证 token 类型
		if tokenType != "" && tokenType.String() != token.Header["typ"] {
			return nil, consts.ErrTokenError
		}

		jwtIdToken.Valid = true

		err = mapstructure.Decode(token.Claims, &jwtIdToken.Token)
		if err != nil {
			a.logger.Logger.Error("Failed to map token claims to JwtIDToken struct.\nError: " + err.Error())
			return nil, nil
		}

		// 手动指定，因为 mapstructure 无法转换 UserID 类型
		jwtIdToken.Token.Sub = sub
	}

	return jwtIdToken, nil
}

func (a *Service) GetAppByToken(ctx context.Context, token string) (*entity.Application, error) {
	r, err := a.q.ApplicationToken.WithContext(ctx).Where(a.q.ApplicationToken.Token.Eq(token)).
		Preload(a.q.ApplicationToken.Application).First()

	if err != nil {
		return nil, err
	}

	if r.Application == nil {
		return nil, consts.ErrApplicationNotFound
	}

	return r.Application, nil
}

func (a *Service) GetApplication(ctx context.Context) (*entity.Application, error) {
	app := ctx.Value(consts.AuthMiddlewareKeyApplication.String())

	app2, ok := app.(*entity.Application)

	if !ok {
		return nil, consts.ErrApplicationNotFound
	}

	return app2, nil
}

func (a *Service) SetApplicationScope(c *gin.Context, application *entity.Application) context.Context {
	c.Set(consts.AuthMiddlewareKeyApplication.String(), application)
	return context.WithValue(c, consts.AuthMiddlewareKeyApplication, application)
}