zelda/backend/apiserver/auth/jwt/manager.go

package jwt

import (
	"fmt"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	gojwt "github.com/golang-jwt/jwt"
	"github.com/ycyxuehan/zelda/apiserver/auth/api"
)

type TokenClaim struct {
	Username string `json:"username"`
	Group    string `json:"group"`
	gojwt.StandardClaims
}

func (t *TokenClaim) Expired() bool {
	return t.ExpiresAt <= time.Now().Unix()
}

type tokenCache struct {
	kubernetesToken string
	claim           *TokenClaim
	cert            string
}

type manager struct {
	tokenStore map[string]tokenCache
	secret     string
	term       time.Duration
}

func NewManager(secret string, term time.Duration) api.AuthManager {
	m := manager{
		tokenStore: make(map[string]tokenCache),
		secret:     secret,
		term:       term,
	}
	return &m
}

//登录
func (m *manager) Login(f api.IdentifyFunc, data *api.AuthentitionRequest) *api.AuthentitionResponse {
	if _, ok := m.IsLogined(data.Username); ok {
		return &api.AuthentitionResponse{Error: fmt.Errorf("user is already logined"), Code: 1000}
	}
	res, err := f(data)
	if err != nil {
		return &api.AuthentitionResponse{Error: err, Code: 1001}
	}
	tokenClaim, err := m.GenerateTokenClaim(data)
	if err != nil {
		return &api.AuthentitionResponse{Error: err, Code: 1002}
	}
	tokenString, err := m.GenerateToken(tokenClaim)
	if err != nil {
		return &api.AuthentitionResponse{Error: err, Code: 1003}
	}
	//存储登录状态
	m.StoreLoginState(data.Username, res, tokenClaim)

	return &api.AuthentitionResponse{Token: tokenString}
}

func (m *manager) Logout(data *api.AuthentitionRequest) *api.AuthentitionResponse {
	//解析token，确认有logout权限
	claim, err := m.ParseTokenClaim(data.Token)
	if err != nil {
		return &api.AuthentitionResponse{Error: err, Code: 1004}

	}
	if c, ok := m.IsLogined(claim.Username); ok {
		if claim.ExpiresAt == c.ExpiresAt {
			delete(m.tokenStore, data.Username)
			return &api.AuthentitionResponse{Message: "success"}
		}
	}
	return &api.AuthentitionResponse{Error: fmt.Errorf("user not login or token invalid"), Code: 1005}
}

func (m *manager) Refresh(data *api.AuthentitionRequest) *api.AuthentitionResponse {
	cliam, err := m.ParseTokenClaim(data.Token)
	if err != nil {
		return &api.AuthentitionResponse{Error: err, Code: 1006}
	}
	//已过期
	if cliam.Valid() != nil {
		return &api.AuthentitionResponse{Error: fmt.Errorf("cannot use an invalid token to refresh: %v", err), Code: 1007}

	}
	//找不到登陆状态
	if _, ok := m.IsLogined(cliam.Username); !ok {
		return &api.AuthentitionResponse{Error: fmt.Errorf("token invalid, user maybe not login"), Code: 1005}

	}
	//重置有效期
	cliam.ExpiresAt = time.Now().Add(m.term).Unix()
	m.ReSetTokenClaim(cliam)

	tokenStr, err := m.GenerateToken(cliam)

	if err != nil {
		return &api.AuthentitionResponse{Error: err, Code: 1003}
	}

	return &api.AuthentitionResponse{Token: tokenStr}
}

func (m *manager) KubernetesToken(data *api.AuthentitionRequest) (string, error) {
	if t, ok := m.tokenStore[data.Username]; ok {
		return t.kubernetesToken, nil
	}
	return "", fmt.Errorf("not found")
}

func (m *manager) KubernetesCert(data *api.AuthentitionRequest) (string, error) {
	if t, ok := m.tokenStore[data.Username]; ok {
		return t.cert, nil
	}
	return "", fmt.Errorf("not found")
}

//生成token claim对象
func (m *manager) GenerateTokenClaim(data *api.AuthentitionRequest) (*TokenClaim, error) {
	now := time.Now()
	claim := TokenClaim{
		Username: data.Username,
		Group:    data.Group,
		StandardClaims: gojwt.StandardClaims{
			ExpiresAt: now.Add(m.term).Unix(),
		},
	}
	return &claim, nil
}

//从string解析token claim
func (m *manager) ParseTokenClaim(token string) (*TokenClaim, error) {
	tokenClaims, err := gojwt.ParseWithClaims(token, &TokenClaim{}, func(token *gojwt.Token) (interface{}, error) {
		return m.secret, nil
	})

	if tokenClaims != nil {
		if claims, ok := tokenClaims.Claims.(*TokenClaim); ok && tokenClaims.Valid {
			return claims, nil
		}
	}
	//
	return nil, err
}

//从token claim 生成 string token
func (m *manager) GenerateToken(claim *TokenClaim) (string, error) {
	tokenClaims := gojwt.NewWithClaims(gojwt.SigningMethodHS256, claim)
	token, err := tokenClaims.SignedString(m.secret)
	return token, err
}

//存储登录状态
func (m *manager) StoreLoginState(user string, result api.IdentifyResult, claim *TokenClaim) {
	//储存登录状态
	t := tokenCache{
		kubernetesToken: result.KubernetesToken,
		claim:           claim,
		cert:            result.Cert,
	}
	m.tokenStore[user] = t
}

//是否已登录
func (m *manager) IsLogined(user string) (*TokenClaim, bool) {
	t, ok := m.tokenStore[user]
	if ok {
		return t.claim, true
	}
	return nil, false
}

//重置token状态
func (m *manager) ReSetTokenClaim(claim *TokenClaim) {
	if t, ok := m.tokenStore[claim.Username]; ok {
		t.claim = claim
		m.tokenStore[claim.Username] = t
	}
}

//token是否有效
func (m *manager) TokenValid(claim *TokenClaim) bool {
	if t, ok := m.tokenStore[claim.Username]; ok {
		if t.claim.ExpiresAt == claim.ExpiresAt && claim.Valid() == nil {
			return true
		}
	}
	return false
}

//中间件
func (m *manager) MiddleWare() gin.HandlerFunc {
	return func(c *gin.Context) {
		tokenString := c.GetHeader("AccessToken")
		if tokenString == "" {
			c.AbortWithStatus(http.StatusUnauthorized)
			return
		}
		tokenClaim, err := m.ParseTokenClaim(tokenString)
		if err != nil {
			c.AbortWithError(http.StatusUnauthorized, err)
			return
		}
		if m.TokenValid(tokenClaim) {
			c.AbortWithError(http.StatusUnauthorized, fmt.Errorf("token invalid"))
			return
		}
		t, err := m.KubernetesToken(&api.AuthentitionRequest{Username: tokenClaim.Username})
		if err != nil || t == "" {
			c.AbortWithError(http.StatusUnauthorized, fmt.Errorf("k8s token not found"))
		}
		c.Request.Header.Set("Authorization", fmt.Sprintf("bearer %s", t))
		cert, _ := m.KubernetesCert(&api.AuthentitionRequest{Username: tokenClaim.Username})
		c.Request.Header.Set("cert", cert)
		c.Next()
	}
}

//初始化auth接口
func (m *manager) InitAuthRoute(identifyFunc api.IdentifyFunc, authGroup *gin.RouterGroup) {
	authGroup.POST("/login", m.HandlerLogin(identifyFunc))
	authGroup.POST("/logout", m.HandlerLogout())
	authGroup.POST("/refresh", m.HandlerRefreshToken())
}