Add full refresh token and access token function

Signed-off-by: Asai Neko <sugar@sne.moe>

config.default.yaml

@@ -1,4 +1,5 @@
 server:
+    application: nixcn-cms
     address: :8000
     external_url: https://example.com
     debug_mode: false
@@ -22,5 +23,6 @@ secrets:
     jwt_secret: something
     turnstile_secret: something
 ttl:
-    magic_link_ttl: 15000
+    magic_link_ttl: 10m
-    jwt_ttl: 86400000
+    jwt_ttl: 15s
+    refresh_ttl: 48h

config/types.go

@@ -10,6 +10,7 @@ type config struct {
 }
 
 type server struct {
+	Application string `yaml:"application"`
 	Address     string `yaml:"address"`
 	DebugMode   string `yaml:"debug_mode"`
 	FileLogger  string `yaml:"file_logger"`
@@ -38,11 +39,12 @@ type email struct {
 }
 
 type secrets struct {
-	jwt_secret       string `yaml:"jwt_secret"`
+	JwtSecret       string `yaml:"jwt_secret"`
-	turnstile_secret string `yaml:"turnstile_secret"`
+	TurnstileSecret string `yaml:"turnstile_secret"`
 }
 
 type ttl struct {
-	magic_link_ttl string `yaml:"magic_link_ttl"`
+	MagicLinkTTL string `yaml:"magic_link_ttl"`
-	jwt_ttl        string `yaml:"jwt_ttl"`
+	JwtTTL       string `yaml:"jwt_ttl"`
+	RefreshTTL   string `yaml:"refresh_ttl"`
 }

data/data.go

@@ -9,7 +9,7 @@ import (
 )
 
 var Database *drivers.DBClient
-var Redis *redis.UniversalClient
+var Redis redis.UniversalClient
 
 func Init() {
 	// Init database

data/drivers/redis.go

@@ -6,7 +6,7 @@ import (
 	"github.com/redis/go-redis/v9"
 )
 
-func Redis(dsn RedisDSN) (*redis.UniversalClient, error) {
+func Redis(dsn RedisDSN) (redis.UniversalClient, error) {
 	// Connect to Redis
 	rdb := redis.NewUniversalClient(&redis.UniversalOptions{
 		Addrs:      dsn.Hosts,
@@ -18,5 +18,5 @@ func Redis(dsn RedisDSN) (*redis.UniversalClient, error) {
 	ctx := context.Background()
 	// Ping Redis
 	_, err := rdb.Ping(ctx).Result()
-	return &rdb, err
+	return rdb, err
 }

internal/crypto/jwt/jwt.go

@@ -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)
-}

internal/crypto/jwt/jwt_test.go

@@ -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)
-	}
-}

internal/cryptography/token.go

@@ -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
+}

middleware/jwt.go

@@ -0,0 +1,57 @@
+package middleware
+
+import (
+	"net/http"
+	"strings"
+
+	"nixcn-cms/internal/cryptography"
+
+	"github.com/gin-gonic/gin"
+	"github.com/golang-jwt/jwt/v5"
+	"github.com/spf13/viper"
+)
+
+func JWTAuth() gin.HandlerFunc {
+	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
+		}
+
+		// Split header to 2
+		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]
+
+		// Verify access token
+		claims := &cryptography.JwtClaims{}
+		token, err := jwt.ParseWithClaims(
+			tokenStr,
+			claims,
+			func(token *jwt.Token) (any, error) {
+				return jwtSecret, nil
+			},
+		)
+
+		if err != nil || !token.Valid {
+			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
+				"error": "invalid or expired token",
+			})
+			return
+		}
+
+		c.Set("user_id", claims.UserID)
+		c.Next()
+	}
+}

service/auth/handler.go

@@ -5,4 +5,5 @@ import "github.com/gin-gonic/gin"
 func Handler(r *gin.RouterGroup) {
 	r.POST("/magic", RequestMagicLink)
 	r.GET("/magic/verify", VerifyMagicLink)
+	r.POST("/refresh", Refresh)
 }

service/auth/magic.go

@@ -3,7 +3,7 @@ package auth
 import (
 	"net/http"
 	"nixcn-cms/data"
-	"nixcn-cms/internal/crypto/jwt"
+	"nixcn-cms/internal/cryptography"
 	"nixcn-cms/pkgs/email"
 	"nixcn-cms/pkgs/magiclink"
 	"nixcn-cms/pkgs/turnstile"
@@ -61,14 +61,14 @@ func RequestMagicLink(c *gin.Context) {
 
 func VerifyMagicLink(c *gin.Context) {
 	// Get token from url
-	token := c.Query("token")
+	magicToken := c.Query("token")
-	if token == "" {
+	if magicToken == "" {
 		c.JSON(400, gin.H{"error": "missing token"})
 		return
 	}
 
 	// Verify email token
-	email, ok := magiclink.VerifyMagicToken(token)
+	email, ok := magiclink.VerifyMagicToken(magicToken)
 	if !ok {
 		c.JSON(401, gin.H{"error": "invalid or expired token"})
 		return
@@ -80,10 +80,19 @@ func VerifyMagicLink(c *gin.Context) {
 	if err != nil {
 		c.JSON(http.StatusUnauthorized, gin.H{"status": "user not found"})
 	}
-	jwtToken, _ := jwt.GenerateToken(userInfo.UserId, "application")
+	JwtTool := cryptography.Token{
+		UserID:      userInfo.UserId,
+		Application: viper.GetString("server.application"),
+	}
+	accessToken, refreshToken, err := JwtTool.IssueTokens()
+	if err != nil {
+		c.JSON(500, gin.H{
+			"status": "error generating tokens",
+		})
+	}
 
 	c.JSON(200, gin.H{
-		"jwt_token": jwtToken,
+		"access_token":  accessToken,
-		"email":     email,
+		"refresh_token": refreshToken,
 	})
 }

service/auth/refresh.go

@@ -0,0 +1,34 @@
+package auth
+
+import (
+	"net/http"
+	"nixcn-cms/internal/cryptography"
+
+	"github.com/gin-gonic/gin"
+	"github.com/spf13/viper"
+)
+
+func Refresh(c *gin.Context) {
+	var req struct {
+		RefreshToken string `json:"refresh_token"`
+	}
+
+	if err := c.ShouldBindJSON(&req); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request"})
+		return
+	}
+
+	JwtTool := cryptography.Token{
+		Application: viper.GetString("server.application"),
+	}
+
+	access, err := JwtTool.RefreshAccessToken(req.RefreshToken)
+	if err != nil {
+		c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid refresh token"})
+		return
+	}
+
+	c.JSON(http.StatusOK, gin.H{
+		"access_token": access,
+	})
+}

service/checkin/handler.go

@@ -1,12 +1,12 @@
 package checkin
 
 import (
-	"nixcn-cms/internal/crypto/jwt"
+	"nixcn-cms/middleware"
 
 	"github.com/gin-gonic/gin"
 )
 
 func Handler(r *gin.RouterGroup) {
-	r.Use(jwt.JWTAuth())
+	r.Use(middleware.JWTAuth())
 	r.POST("", Checkin)
 }

service/user/handler.go

@@ -1,12 +1,12 @@
 package user
 
 import (
-	"nixcn-cms/internal/crypto/jwt"
+	"nixcn-cms/middleware"
 
 	"github.com/gin-gonic/gin"
 )
 
 func Handler(r *gin.RouterGroup) {
-	r.Use(jwt.JWTAuth())
+	r.Use(middleware.JWTAuth())
 	r.GET("/info", UserInfo)
 }