forked from nixcn/nixcn-cms
Add full refresh token and access token function
Signed-off-by: Asai Neko <sugar@sne.moe>
This commit is contained in:
@@ -1,77 +0,0 @@
|
||||
package jwt
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"github.com/google/uuid"
|
||||
"github.com/spf13/viper"
|
||||
)
|
||||
|
||||
type Claims struct {
|
||||
UserID uuid.UUID `json:"user_id"`
|
||||
jwt.RegisteredClaims
|
||||
}
|
||||
|
||||
func JWTAuth() gin.HandlerFunc {
|
||||
var JwtSecret = []byte(viper.GetString("secrets.jwt_secret"))
|
||||
return func(c *gin.Context) {
|
||||
auth := c.GetHeader("Authorization")
|
||||
if auth == "" {
|
||||
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
|
||||
"error": "missing Authorization header",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
parts := strings.SplitN(auth, " ", 2)
|
||||
if len(parts) != 2 || parts[0] != "Bearer" {
|
||||
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
|
||||
"error": "invalid Authorization header format",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
tokenStr := parts[1]
|
||||
|
||||
token, err := jwt.ParseWithClaims(tokenStr, &Claims{}, func(token *jwt.Token) (interface{}, error) {
|
||||
return JwtSecret, nil
|
||||
})
|
||||
|
||||
if err != nil || !token.Valid {
|
||||
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
|
||||
"error": "invalid or expired token",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
claims, ok := token.Claims.(*Claims)
|
||||
if !ok {
|
||||
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
|
||||
"error": "invalid token claims",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
c.Set("user_id", claims.UserID)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func GenerateToken(userID uuid.UUID, application string) (string, error) {
|
||||
var JwtSecret = []byte(viper.GetString("secrets.jwt_secret"))
|
||||
claims := Claims{
|
||||
UserID: userID,
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(viper.GetDuration("ttl.jwt_ttl"))),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Issuer: application,
|
||||
},
|
||||
}
|
||||
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
return token.SignedString(JwtSecret)
|
||||
}
|
||||
@@ -1,94 +0,0 @@
|
||||
package jwt
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"nixcn-cms/config"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"github.com/google/uuid"
|
||||
"github.com/spf13/viper"
|
||||
)
|
||||
|
||||
func init() {
|
||||
config.Init()
|
||||
}
|
||||
|
||||
func generateTestToken(userID uuid.UUID, expire time.Duration) string {
|
||||
var JwtSecret = []byte(viper.GetString("server.jwt_secret"))
|
||||
claims := Claims{
|
||||
UserID: userID,
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(expire)),
|
||||
},
|
||||
}
|
||||
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
tokenStr, _ := token.SignedString(JwtSecret)
|
||||
return tokenStr
|
||||
}
|
||||
func TestJWTAuth_MissingToken(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
r.Use(JWTAuth())
|
||||
r.GET("/test", func(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("expected 401, got %d", w.Code)
|
||||
}
|
||||
}
|
||||
func TestJWTAuth_InvalidToken(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
r.Use(JWTAuth())
|
||||
r.GET("/test", func(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"ok": true})
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.Header.Set("Authorization", "Bearer invalid.token.here")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("expected 401, got %d", w.Code)
|
||||
}
|
||||
}
|
||||
func TestJWTAuth_ValidToken(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
r.Use(JWTAuth())
|
||||
r.GET("/test", func(c *gin.Context) {
|
||||
userID := c.GetUint("user_id")
|
||||
c.JSON(200, gin.H{
|
||||
"user_id": userID,
|
||||
})
|
||||
})
|
||||
|
||||
uuid, _ := uuid.NewUUID()
|
||||
token := generateTestToken(uuid, time.Hour)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d", w.Code)
|
||||
}
|
||||
}
|
||||
145
internal/cryptography/token.go
Normal file
145
internal/cryptography/token.go
Normal file
@@ -0,0 +1,145 @@
|
||||
package cryptography
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"nixcn-cms/data"
|
||||
"time"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"github.com/google/uuid"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/spf13/viper"
|
||||
)
|
||||
|
||||
type Token struct {
|
||||
UserID uuid.UUID
|
||||
Application string
|
||||
}
|
||||
|
||||
type JwtClaims struct {
|
||||
UserID uuid.UUID `json:"user_id"`
|
||||
jwt.RegisteredClaims
|
||||
}
|
||||
|
||||
// Generate jwt clames
|
||||
func (self *Token) NewClaims() JwtClaims {
|
||||
return JwtClaims{
|
||||
UserID: self.UserID,
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
ExpiresAt: jwt.NewNumericDate(time.Now().Add(viper.GetDuration("ttl.jwt_ttl"))),
|
||||
IssuedAt: jwt.NewNumericDate(time.Now()),
|
||||
Issuer: self.Application,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Generate access token
|
||||
func (self *Token) GenerateAccessToken() (string, error) {
|
||||
claims := self.NewClaims()
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
secret := viper.GetString("secrets.jwt_secret")
|
||||
return token.SignedString(secret)
|
||||
}
|
||||
|
||||
// Generate refresh token
|
||||
func (self *Token) GenerateRefreshToken() (string, error) {
|
||||
b := make([]byte, 32)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return base64.RawURLEncoding.EncodeToString(b), nil
|
||||
}
|
||||
|
||||
// Issue both access and refresh token
|
||||
func (self *Token) IssueTokens() (string, string, error) {
|
||||
// Gen atk
|
||||
access, err := self.GenerateAccessToken()
|
||||
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,
|
||||
self.UserID.String(),
|
||||
ttl,
|
||||
).Err(); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
|
||||
// user -> refresh tokens
|
||||
userSetKey := "user:" + self.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 {
|
||||
if err == redis.Nil {
|
||||
return "", errors.New("invalid refresh token")
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
|
||||
userID, err := uuid.Parse(userIDStr)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
self.UserID = userID
|
||||
|
||||
// Generate access token
|
||||
return self.GenerateAccessToken()
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user