jwtmdl.go

package jwtmdl

import (
	"strings"
	"sync"
	"time"

	"corelab.mkcl.org/MKCLOS/coredevelopmentplatform/corepkgv2/authmdl/sessionmdl"
	"corelab.mkcl.org/MKCLOS/coredevelopmentplatform/corepkgv2/errormdl"
	jwt "github.com/golang-jwt/jwt/v5"
	"github.com/tidwall/gjson"
)

// map to store kid
type Jwtkid struct {
	jwtKidMap map[string]string
	mu        sync.Mutex
}

func (j *Jwtkid) Get(kid string) (string, bool) {
	j.mu.Lock()
	defer j.mu.Unlock()
	key, ok := j.jwtKidMap[kid]
	return key, ok
}
func (j *Jwtkid) Set(kid, key string) {
	j.mu.Lock()
	defer j.mu.Unlock()
	j.jwtKidMap[kid] = key
}

// jwtKidStore - store to keep kid and jwtKey
var jwtKidStore = Jwtkid{jwtKidMap: make(map[string]string)}

// GlobalJWTKey - key to decode and encode token
var GlobalJWTKey string

var keyFunc = func(key string) jwt.Keyfunc {
	return func(token *jwt.Token) (interface{}, error) {
		if kid, ok := token.Header["kid"].(string); ok {
			if jwtKey, ok := jwtKidStore.Get(kid); ok && jwtKey == "DEFAULTKEY" {
				return ([]byte(key)), nil
			}
			if jwtKey, exists := jwtKidStore.Get(kid); exists {
				return ([]byte(jwtKey)), nil
			}
		}
		b := ([]byte(key))
		return b, nil
	}
}

func InitJwtKidStore(kidConfig gjson.Result) {
	for _, kid := range kidConfig.Array() {
		jwtKidStore.Set(kid.Get("kid").String(), kid.Get("jwtKey").String())
	}
}

type jwtCustomClaim struct {
	UserID    string   `json:"userId"`
	SessionId string   `json:"sessionId,omitempty"`
	Groups    []string `json:"groups"`
	ClientIP  string   `json:"clientIP"`
	HitsCount int      `json:"hitsCount"`
	Token     string   `json:"token"`
	Metadata  string   `json:"metadata"`
	jwt.RegisteredClaims
}

func generate(claims jwtCustomClaim, key string) (string, error) {
	return jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString([]byte(key))
}

// extract return token from header string
func extract(tokenReq string) (string, error) {
	tokenArray := strings.Split(tokenReq, "Bearer")
	if len(tokenArray) <= 1 {
		return "", errormdl.Wrap("Provided JWT token is nil or invalid ")
	}

	return strings.Trim(tokenArray[1], " "), nil
}

// decode accepts a parsed token and error from parse operation.
func decode(token *jwt.Token, err error) (jwt.MapClaims, error) {
	if err != nil {
		// loggermdl.LogError("Error while parsing JWT Token: ", err)
		return nil, err
	}

	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		// loggermdl.LogError("Error while parsing claims to MapClaims")
		return nil, errormdl.Wrap("Error while getting claims")
	}

	// validate user session from session id present in token
	if err := sessionmdl.ValidateSessionFromToken(claims); err != nil {
		// loggermdl.LogError("session validation failed with err:", err)
		return nil, sessionmdl.ErrSessionValidationFailed
	}

	return claims, nil
}

func GenerateTokenWithOptions(args ...Option) (string, error) {
	options := new(Options)

	options.Key = GlobalJWTKey

	for i := range args {
		args[i](options)
	}

	claims := jwtCustomClaim{
		ClientIP:  options.ClientIP,
		Groups:    options.Groups,
		Metadata:  options.Metadata,
		SessionId: options.Session.SessionId,
		UserID:    options.UserID,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(time.Unix(options.ExpiresAt, 0)),
		},
	}

	t, err := generate(claims, options.Key)
	if err != nil {
		return "", err
	}

	if len(options.Session.SessionId) > 0 {
		sessionmdl.Set(options.UserID, options.Session)
	}

	return t, nil
}

// GenerateToken generates JWT token from Login object
func GenerateToken(loginID string, groups []string, clientIP string, metadata gjson.Result, expirationTime time.Duration) (string, error) {
	return GenerateTokenWithJWTKey(loginID, groups, clientIP, metadata, expirationTime, GlobalJWTKey)
}

// GenerateTokenWithJWTKey generates JWT token from Login object
func GenerateTokenWithJWTKey(loginID string, groups []string, clientIP string, metadata gjson.Result, expirationTime time.Duration, JWTKey string) (string, error) {
	claims := jwtCustomClaim{
		UserID:   loginID,
		Groups:   groups,
		ClientIP: clientIP,
		Metadata: metadata.String(),
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(expirationTime)),
		},
	}

	return generate(claims, JWTKey)
}

// GeneratePricipleObjUsingToken GeneratePricipleObjUsingToken
func GeneratePricipleObjUsingToken(tokenReq string, jwtKey string) (jwt.MapClaims, error) {

	token, err := extract(tokenReq)
	if err != nil {
		return nil, err
	}

	return decode(jwt.Parse(token, keyFunc(jwtKey)))
}