Merge pull request #574 from trheyi/main

Add cookie-trace and  cookie-jwt guards

service/guards.go

@@ -4,6 +4,7 @@ import (
 	"strings"
 
 	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
 	"github.com/yaoapp/yao/helper"
 
 	"github.com/yaoapp/yao/widgets/chart"
@@ -18,6 +19,8 @@ var Guards = map[string]gin.HandlerFunc{
 	"bearer-jwt":       guardBearerJWT,   // Bearer JWT
 	"query-jwt":        guardQueryJWT,    // Get JWT Token from query string  "__tk"
 	"cross-origin":     guardCrossOrigin, // Cross-Origin Resource Sharing
+	"cookie-trace":     guardCookieTrace, // Set sid cookie
+	"cookie-jwt":       guardCookieJWT,   // Get JWT Token from cookie "__tk"
 	"widget-table":     table.Guard,      // Widget Table Guard
 	"widget-list":      list.Guard,       // Widget List Guard
 	"widget-form":      form.Guard,       // Widget Form Guard
@@ -25,6 +28,40 @@ var Guards = map[string]gin.HandlerFunc{
 	"widget-dashboard": dashboard.Guard,  // Widget Dashboard Guard
 }
 
+// guardCookieTrace set sid cookie
+func guardCookieTrace(c *gin.Context) {
+	sid, err := c.Cookie("sid")
+	if err != nil {
+		sid = uuid.New().String()
+		c.SetCookie("sid", sid, 0, "/", "", false, true)
+		c.Set("__sid", sid)
+		c.Next()
+		return
+	}
+	c.Set("__sid", sid)
+	return
+}
+
+// Cookie Cookie JWT
+func guardCookieJWT(c *gin.Context) {
+	tokenString, err := c.Cookie("__tk")
+	if err != nil {
+		c.JSON(403, gin.H{"code": 403, "message": "No permission"})
+		c.Abort()
+		return
+	}
+
+	if tokenString == "" {
+		c.JSON(403, gin.H{"code": 403, "message": "No permission"})
+		c.Abort()
+		return
+	}
+
+	claims := helper.JwtValidate(tokenString)
+	c.Set("__sid", claims.SID)
+	return
+}
+
 // JWT Bearer JWT
 func guardBearerJWT(c *gin.Context) {
 	tokenString := c.Request.Header.Get("Authorization")

sui/api/guards.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
 	jsoniter "github.com/json-iterator/go"
 	"github.com/yaoapp/gou/process"
 	v8 "github.com/yaoapp/gou/runtime/v8"
@@ -18,10 +19,10 @@ import (
 
 // Guards middlewares
 var Guards = map[string]func(c *Request) error{
-	"bearer-jwt": guardBearerJWT, // Bearer JWT
-	"query-jwt":  guardQueryJWT,  // Get JWT Token from query string  "__tk"
-	"cookie-jwt": guardCookieJWT, // Get JWT Token from cookie "__tk"
-
+	"bearer-jwt":   guardBearerJWT,   // Bearer JWT
+	"query-jwt":    guardQueryJWT,    // Get JWT Token from query string  "__tk"
+	"cookie-jwt":   guardCookieJWT,   // Get JWT Token from cookie "__tk"
+	"cookie-trace": guardCookieTrace, // Set sid cookie
 }
 
 // JWT Bearer JWT
@@ -47,7 +48,7 @@ func guardBearerJWT(r *Request) error {
 // JWT Bearer JWT
 func guardCookieJWT(r *Request) error {
 	if r.context == nil {
-		return fmt.Errorf("No permission")
+		return fmt.Errorf("Context is nil")
 	}
 	c := r.context
 
@@ -70,6 +71,25 @@ func guardCookieJWT(r *Request) error {
 	return nil
 }
 
+func guardCookieTrace(r *Request) error {
+	if r.context == nil {
+		return fmt.Errorf("Context is nil")
+	}
+
+	c := r.context
+	sid, err := c.Cookie("sid")
+	if err != nil {
+		sid = uuid.New().String()
+		c.SetCookie("sid", sid, 0, "/", "", false, true)
+		c.Set("__sid", sid)
+		r.Sid = sid
+		return nil
+	}
+	c.Set("__sid", sid)
+	r.Sid = sid
+	return nil
+}
+
 // JWT Bearer JWT
 func guardQueryJWT(r *Request) error {
 	if r.context == nil {