Use the gin framework

utils/csrf.go

@@ -5,6 +5,8 @@ import (
 	"crypto/subtle"
 	"encoding/base64"
 	"net/http"
+
+	"github.com/gin-gonic/gin"
 )
 
 const (
@@ -34,55 +36,51 @@ func GenerateCSRFToken() (string, error) {
 }
 
 // SetCSRFToken 设置CSRF令牌到Cookie和响应头
-func SetCSRFToken(w http.ResponseWriter, token string) {
-	// 设置CSRF令牌到Cookie
-	cookie := CreateSecureCookie(CSRFCookieName, token, 3600) // 1小时过期
-	http.SetCookie(w, cookie)
-
-	// 设置CSRF令牌到响应头，方便JavaScript获取
-	w.Header().Set("X-CSRF-Token", token)
+func SetCSRFToken(c *gin.Context, token string) {
+	c.SetCookie(CSRFCookieName, token, 3600*24, "/", "", false, true)
+	c.Header(CSRFHeaderName, token)
 }
 
-// GetCSRFTokenFromRequest 从请求中获取CSRF令牌
+// GetCSRFTokenFromRequest 从Gin请求中获取CSRF令牌
 // 优先级：Header > Form > Cookie
-func GetCSRFTokenFromRequest(r *http.Request) string {
+func GetCSRFTokenFromRequest(c *gin.Context) string {
 	// 1. 从Header获取
-	if token := r.Header.Get(CSRFHeaderName); token != "" {
+	if token := c.GetHeader(CSRFHeaderName); token != "" {
 		return token
 	}
 
 	// 2. 从Form获取
-	if token := r.FormValue(CSRFFormField); token != "" {
+	if token := c.PostForm(CSRFFormField); token != "" {
 		return token
 	}
 
 	// 3. 从Cookie获取（作为备选）
-	if cookie, err := r.Cookie(CSRFCookieName); err == nil {
-		return cookie.Value
+	if cookie, err := c.Cookie(CSRFCookieName); err == nil {
+		return cookie
 	}
 
 	return ""
 }
 
 // GetCSRFTokenFromCookie 从Cookie中获取CSRF令牌
-func GetCSRFTokenFromCookie(r *http.Request) string {
-	cookie, err := r.Cookie(CSRFCookieName)
+func GetCSRFTokenFromCookie(c *gin.Context) string {
+	cookie, err := c.Cookie(CSRFCookieName)
 	if err != nil {
 		return ""
 	}
-	return cookie.Value
+	return cookie
 }
 
 // ValidateCSRFToken 验证CSRF令牌
-func ValidateCSRFToken(r *http.Request) bool {
+func ValidateCSRFToken(c *gin.Context) bool {
 	// 获取Cookie中的令牌（服务器端存储的）
-	cookieToken := GetCSRFTokenFromCookie(r)
+	cookieToken := GetCSRFTokenFromCookie(c)
 	if cookieToken == "" {
 		return false
 	}
 
 	// 获取请求中的令牌（客户端提交的）
-	requestToken := GetCSRFTokenFromRequest(r)
+	requestToken := GetCSRFTokenFromRequest(c)
 	if requestToken == "" {
 		return false
 	}
@@ -92,47 +90,62 @@ func ValidateCSRFToken(r *http.Request) bool {
 }
 
 // CSRFProtection CSRF保护中间件
-func CSRFProtection(next http.HandlerFunc) http.HandlerFunc {
-	return func(w http.ResponseWriter, r *http.Request) {
+func CSRFProtection() gin.HandlerFunc {
+	return func(c *gin.Context) {
 		// 对于GET、HEAD、OPTIONS请求，只生成令牌，不验证
-		if r.Method == http.MethodGet || r.Method == http.MethodHead || r.Method == http.MethodOptions {
+		if c.Request.Method == http.MethodGet || c.Request.Method == http.MethodHead || c.Request.Method == http.MethodOptions {
 			// 生成新的CSRF令牌
 			token, err := GenerateCSRFToken()
 			if err != nil {
-				http.Error(w, "Internal Server Error", http.StatusInternalServerError)
+				c.JSON(http.StatusInternalServerError, gin.H{
+					"code": 1,
+					"msg":  "Internal Server Error",
+					"data": nil,
+				})
+				c.Abort()
 				return
 			}
-			SetCSRFToken(w, token)
-			next(w, r)
+			SetCSRFToken(c, token)
+			c.Next()
 			return
 		}
 
 		// 对于POST、PUT、DELETE等修改性请求，验证CSRF令牌
-		if !ValidateCSRFToken(r) {
-			JsonResponse(w, http.StatusForbidden, false, "CSRF令牌验证失败", nil)
+		if !ValidateCSRFToken(c) {
+			c.JSON(http.StatusForbidden, gin.H{
+				"code": 1,
+				"msg":  "CSRF令牌验证失败",
+				"data": nil,
+			})
+			c.Abort()
 			return
 		}
 
 		// 验证通过，继续处理请求
-		next(w, r)
+		c.Next()
 	}
 }
 
 // RequireCSRFToken 要求CSRF令牌的中间件（用于特定路由）
-func RequireCSRFToken(next http.HandlerFunc) http.HandlerFunc {
-	return func(w http.ResponseWriter, r *http.Request) {
-		if !ValidateCSRFToken(r) {
-			JsonResponse(w, http.StatusForbidden, false, "CSRF令牌验证失败", nil)
+func RequireCSRFToken() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		if !ValidateCSRFToken(c) {
+			c.JSON(http.StatusForbidden, gin.H{
+				"code": 1,
+				"msg":  "CSRF令牌验证失败",
+				"data": nil,
+			})
+			c.Abort()
 			return
 		}
-		next(w, r)
+		c.Next()
 	}
 }
 
 // GetCSRFTokenForTemplate 获取用于模板的CSRF令牌
-func GetCSRFTokenForTemplate(r *http.Request) string {
+func GetCSRFTokenForTemplate(c *gin.Context) string {
 	// 尝试从Cookie获取现有令牌
-	if token := GetCSRFTokenFromCookie(r); token != "" {
+	if token := GetCSRFTokenFromCookie(c); token != "" {
 		return token
 	}
 
@@ -145,24 +158,36 @@ func GetCSRFTokenForTemplate(r *http.Request) string {
 }
 
 // CSRFTokenHandler 专门用于获取CSRF令牌的API端点
-func CSRFTokenHandler(w http.ResponseWriter, r *http.Request) {
-	if r.Method != http.MethodGet {
-		JsonResponse(w, http.StatusMethodNotAllowed, false, "只支持GET请求", nil)
+func CSRFTokenHandler(c *gin.Context) {
+	if c.Request.Method != http.MethodGet {
+		c.JSON(http.StatusMethodNotAllowed, gin.H{
+			"code": 1,
+			"msg":  "只支持GET请求",
+			"data": nil,
+		})
 		return
 	}
 
 	// 生成新的CSRF令牌
 	token, err := GenerateCSRFToken()
 	if err != nil {
-		JsonResponse(w, http.StatusInternalServerError, false, "生成CSRF令牌失败", nil)
+		c.JSON(http.StatusInternalServerError, gin.H{
+			"code": 1,
+			"msg":  "生成CSRF令牌失败",
+			"data": nil,
+		})
 		return
 	}
 
 	// 设置令牌到Cookie和响应头
-	SetCSRFToken(w, token)
+	SetCSRFToken(c, token)
 
 	// 返回令牌给前端
-	JsonResponse(w, http.StatusOK, true, "CSRF令牌获取成功", map[string]interface{}{
-		"csrf_token": token,
+	c.JSON(http.StatusOK, gin.H{
+		"code": 0,
+		"msg":  "CSRF令牌生成成功",
+		"data": gin.H{
+			"csrf_token": token,
+		},
 	})
 }