From 5004b4342b81f754fb13ad43c2c464b04c25ba30 Mon Sep 17 00:00:00 2001
From: Chongzhi Deng <chongzhi@hust.edu.cn>
Date: Fri, 14 Mar 2025 10:36:40 +0800
Subject: [PATCH] feat: enhance PAT validation in middleware (#3888)

Signed-off-by: BruceAko <chongzhi@hust.edu.cn>
---
 manager/middlewares/personal_access_token.go | 75 +++++++++++++++++++-
 1 file changed, 74 insertions(+), 1 deletion(-)

diff --git a/manager/middlewares/personal_access_token.go b/manager/middlewares/personal_access_token.go
index 71bdcee6e..b71fb7e65 100644
--- a/manager/middlewares/personal_access_token.go
+++ b/manager/middlewares/personal_access_token.go
@@ -17,14 +17,23 @@
 package middlewares
 
 import (
+	"fmt"
 	"net/http"
+	"regexp"
 	"strings"
+	"time"
 
 	"github.com/gin-gonic/gin"
 	"github.com/go-http-utils/headers"
 	"gorm.io/gorm"
 
+	logger "d7y.io/dragonfly/v2/internal/dflog"
 	"d7y.io/dragonfly/v2/manager/models"
+	"d7y.io/dragonfly/v2/manager/types"
+)
+
+var (
+	oapiResourceRegexp = regexp.MustCompile(`^/oapi/v[0-9]+/([-_a-zA-Z]*)[/.*]*`)
 )
 
 func PersonalAccessToken(gdb *gorm.DB) gin.HandlerFunc {
@@ -42,7 +51,9 @@ func PersonalAccessToken(gdb *gorm.DB) gin.HandlerFunc {
 
 		// Check if the personal access token is valid.
 		personalAccessToken := tokenFields[1]
-		if err := gdb.WithContext(c).Where("token = ?", personalAccessToken).First(&models.PersonalAccessToken{}).Error; err != nil {
+		var token models.PersonalAccessToken
+		if err := gdb.WithContext(c).Where("token = ?", personalAccessToken).First(&token).Error; err != nil {
+			logger.Errorf("Invalid personal access token attempt: %s, error: %v", c.Request.URL.Path, err)
 			c.JSON(http.StatusUnauthorized, ErrorResponse{
 				Message: http.StatusText(http.StatusUnauthorized),
 			})
@@ -50,6 +61,68 @@ func PersonalAccessToken(gdb *gorm.DB) gin.HandlerFunc {
 			return
 		}
 
+		// Check if the token is active.
+		if token.State != models.PersonalAccessTokenStateActive {
+			logger.Errorf("Inactive token used: %s, token name: %s, user_id: %d", c.Request.URL.Path, token.Name, token.UserID)
+			c.JSON(http.StatusForbidden, ErrorResponse{
+				Message: "Token is inactive",
+			})
+			c.Abort()
+			return
+		}
+
+		// Check if the token has expired.
+		if time.Now().After(token.ExpiredAt) {
+			logger.Errorf("Expired token used: %s, token name: %s, user_id: %d, expired: %v",
+				c.Request.URL.Path, token.Name, token.UserID, token.ExpiredAt)
+			c.JSON(http.StatusForbidden, ErrorResponse{
+				Message: "Token has expired",
+			})
+			c.Abort()
+			return
+		}
+
+		// Check if the token's scopes include the required resource type.
+		hasScope := false
+		resourceType := getAPIResourceType(c.Request.URL.Path)
+		for _, scope := range token.Scopes {
+			if scope == resourceType {
+				hasScope = true
+				break
+			}
+		}
+
+		if !hasScope {
+			logger.Errorf("Insufficient scope token used: %s, token name: %s, user_id: %d, required: %s, available: %v",
+				c.Request.URL.Path, token.Name, token.UserID, resourceType, token.Scopes)
+			c.JSON(http.StatusForbidden, ErrorResponse{
+				Message: fmt.Sprintf("Token doesn't have permission to access this resource. Required scope: %s", resourceType),
+			})
+			c.Abort()
+			return
+		}
+
 		c.Next()
 	}
 }
+
+// getAPIResourceType extracts the resource type from the path.
+// For example: /oapi/v1/jobs -> job, /oapi/v1/clusters -> cluster.
+func getAPIResourceType(path string) string {
+	matches := oapiResourceRegexp.FindStringSubmatch(path)
+	if len(matches) != 2 {
+		return ""
+	}
+
+	resource := strings.ToLower(matches[1])
+	switch resource {
+	case "jobs":
+		return types.PersonalAccessTokenScopeJob
+	case "clusters":
+		return types.PersonalAccessTokenScopeCluster
+	case "preheats":
+		return types.PersonalAccessTokenScopePreheat
+	default:
+		return resource
+	}
+}