nixcn-cms/pkgs/authtoken/authtoken.go

package authtoken

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"fmt"
	"nixcn-cms/data"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"
	"github.com/spf13/viper"
)

type Token struct {
	Application string
}

type JwtClaims struct {
	UserID uuid.UUID `json:"user_id"`
	jwt.RegisteredClaims
}

// Generate jwt clames
func (self *Token) NewClaims(userId uuid.UUID) JwtClaims {
	return JwtClaims{
		UserID: userId,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(viper.GetDuration("ttl.access_ttl"))),
			IssuedAt:  jwt.NewNumericDate(time.Now()),
			Issuer:    self.Application,
		},
	}
}

// Generate access token
func (self *Token) GenerateAccessToken(userId uuid.UUID) (string, error) {
	claims := self.NewClaims(userId)
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	secret := viper.GetString("secrets.jwt_secret")
	signedToken, err := token.SignedString([]byte(secret))
	if err != nil {
		return "", fmt.Errorf("error signing token: %v", err)
	}
	return signedToken, nil
}

// Generate refresh token
func (self *Token) GenerateRefreshToken() (string, error) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return base64.URLEncoding.EncodeToString(b), nil
}

// Issue both access and refresh token
func (self *Token) IssueTokens(userId uuid.UUID) (string, string, error) {
	// Gen atk
	access, err := self.GenerateAccessToken(userId)
	if err != nil {
		return "", "", err
	}

	// Gen rtk
	refresh, err := self.GenerateRefreshToken()
	if err != nil {
		return "", "", err
	}

	// Store to redis
	ctx := context.Background()
	ttl := viper.GetDuration("ttl.refresh_ttl")

	// refresh -> user
	if err := data.Redis.Set(
		ctx,
		"refresh:"+refresh,
		userId.String(),
		ttl,
	).Err(); err != nil {
		return "", "", err
	}

	// user -> refresh tokens
	userSetKey := "user:" + userId.String() + ":refresh_tokens"

	if err := data.Redis.SAdd(
		ctx,
		userSetKey,
		refresh,
	).Err(); err != nil {
		return "", "", err
	}

	// set user ttl >= all refresh token
	_ = data.Redis.Expire(ctx, userSetKey, ttl).Err()

	return access, refresh, nil
}

// Refresh access token
func (self *Token) RefreshAccessToken(refreshToken string) (string, error) {
	// Read rtk:userid from redis
	ctx := context.Background()
	key := "refresh:" + refreshToken

	userIdStr, err := data.Redis.Get(ctx, key).Result()
	if err != nil {
		return "", err
	}

	userId, err := uuid.Parse(userIdStr)
	if err != nil {
		return "", err
	}

	// Generate access token
	return self.GenerateAccessToken(userId)
}

func (self *Token) RenewRefreshToken(refreshToken string) (string, error) {
	ctx := context.Background()
	ttl := viper.GetDuration("ttl.refresh_ttl")

	key := "refresh:" + refreshToken
	userIdStr, err := data.Redis.Get(ctx, key).Result()
	if err != nil {
		return "", err
	}

	refresh, err := self.GenerateRefreshToken()
	if err != nil {
		return "", err
	}

	err = self.RevokeRefreshToken(refreshToken)
	if err != nil {
		return "", err
	}

	// refresh -> user
	if err := data.Redis.Set(
		ctx,
		"refresh:"+refresh,
		userIdStr,
		ttl,
	).Err(); err != nil {
		return "", err
	}

	// user -> refresh tokens
	userSetKey := "user:" + userIdStr + ":refresh_tokens"

	if err := data.Redis.SAdd(
		ctx,
		userSetKey,
		refresh,
	).Err(); err != nil {
		return "", err
	}

	// set user ttl >= all refresh token
	_ = data.Redis.Expire(ctx, userSetKey, ttl).Err()

	return refresh, nil
}

func (self *Token) RevokeRefreshToken(refreshToken string) error {
	ctx := context.Background()

	key := "refresh:" + refreshToken

	userIDStr, err := data.Redis.Get(ctx, key).Result()
	if err != nil {
		return nil
	}

	userSetKey := "user:" + userIDStr + ":refresh_tokens"

	// Delete rtk from redis
	pipe := data.Redis.TxPipeline()
	pipe.Del(ctx, key)                       // rtk:userid index
	pipe.SRem(ctx, userSetKey, refreshToken) // userid:rtk index
	_, err = pipe.Exec(ctx)

	return err
}

func (self *Token) HeaderVerify(header string) (string, error) {
	if header == "" {
		return "", nil
	}

	jwtSecret := []byte(viper.GetString("secrets.jwt_secret"))
	// Split header to 2
	parts := strings.SplitN(header, " ", 2)
	if len(parts) != 2 || parts[0] != "Bearer" {
		return "", errors.New("invalid Authorization header format")
	}

	tokenStr := parts[1]

	// Verify access token
	claims := &JwtClaims{}
	token, err := jwt.ParseWithClaims(
		tokenStr,
		claims,
		func(token *jwt.Token) (any, error) {
			return jwtSecret, nil
		},
	)

	if err != nil || !token.Valid {
		return "", errors.New("invalid or expired token")
	}

	return claims.UserID.String(), nil
}