feat: Support multi urls matching. (#125)

* feat: Support multi urls matching.

Signed-off-by: IRONICBo <47499836+IRONICBo@users.noreply.github.com>

* fix: Update comments.

Signed-off-by: IRONICBo <47499836+IRONICBo@users.noreply.github.com>

* fix: Update misspelling.

Signed-off-by: IRONICBo <47499836+IRONICBo@users.noreply.github.com>

---------

Signed-off-by: IRONICBo <47499836+IRONICBo@users.noreply.github.com>

server/cmd/dbmigration/main.go

@@ -61,7 +61,7 @@ func main() {
 				log.Panicf("OpenKF Table Migration", "Drop table %T... failed", tables[i])
 			}
 			log.Infof("OpenKF Table Migration", "Drop table %T... ok", tables[i])
-		}	
+		}
 	}
 
 	// migrate tables.

server/cmd/genhooks/pkg/template.go

@@ -51,22 +51,22 @@ type {{.HookName}} struct {
 	urltrie.Hook
 }
 
-// EDIT THIS TO YOUR OWN HOOK PATTERN
-func (h {{.HookName}}) Pattern() string {
+// Patterns EDIT THIS TO YOUR OWN HOOK PATTERN
+func (h {{.HookName}}) Patterns() string {
 	return "{{.UrlPattern}}"
 }
 
-// EDIT THIS TO YOUR OWN HOOK PRIORITY
+// Priority EDIT THIS TO YOUR OWN HOOK PRIORITY
 func (h GlobalHook) Priority() int64 {
 	return {{.Prority}}
 }
 
-// EDIT THIS TO YOUR OWN HOOK BEFORE RUN
+// BeforeRun EDIT THIS TO YOUR OWN HOOK BEFORE RUN
 func (h {{.HookName}}) BeforeRun(c *gin.Context) {
 	c.Next()
 }
 
-// EDIT THIS TO YOUR OWN HOOK AFTER RUN
+// AfterRun EDIT THIS TO YOUR OWN HOOK AFTER RUN
 func (h {{.HookName}}) AfterRun(c *gin.Context) {
 	c.Next()
 }

server/internal/middleware/hooks/gen_cors_hook.go

@@ -0,0 +1,67 @@
+// Copyright © 2023 OpenIM open source community. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package hooks
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+
+	urltrie "github.com/OpenIMSDK/OpenKF/server/internal/middleware/hooks/url_trie"
+)
+
+var _ urltrie.Hook = (*CORS)(nil)
+
+func init() {
+	urltrie.RegisterHook(&CORS{})
+	fmt.Println("RegisterHook", "Register Hook[CORS] success...")
+}
+
+// CORS implement urltrie.Hook.
+type CORS struct {
+	urltrie.Hook
+}
+
+// Patterns EDIT THIS TO YOUR OWN HOOK PATTERN.
+func (h *CORS) Patterns() []string {
+	return []string{
+		"/*",
+	}
+}
+
+// Priority EDIT THIS TO YOUR OWN HOOK PRIORITY.
+func (h *CORS) Priority() int64 {
+	return 0
+}
+
+// BeforeRun EDIT THIS TO YOUR OWN HOOK BEFORE RUN, DO NOT NEED USE Next() FUNCTION.
+func (h *CORS) BeforeRun(c *gin.Context) {
+	// c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
+	c.Writer.Header().Set("Access-Control-Allow-Origin", c.GetHeader("origin"))
+	c.Writer.Header().Set("Access-Control-Max-Age", "70000")
+	c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, UPDATE, PATCH")
+	c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length,token")
+	c.Writer.Header().Set("Access-Control-Expose-Headers", "Content-Length, Access-Control-Allow-Headers,token")
+	c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
+
+	if c.Request.Method == "OPTIONS" {
+		c.AbortWithStatus(http.StatusNoContent)
+	}
+}
+
+// AfterRun EDIT THIS TO YOUR OWN HOOK AFTER RUN, DO NOT NEED USE Next() FUNCTION.
+func (h *CORS) AfterRun(c *gin.Context) {
+}

server/internal/middleware/hooks/gen_jwt_hook.go

@@ -0,0 +1,75 @@
+// Copyright © 2023 OpenIM open source community. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package hooks
+
+import (
+	"fmt"
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+
+	urltrie "github.com/OpenIMSDK/OpenKF/server/internal/middleware/hooks/url_trie"
+	"github.com/OpenIMSDK/OpenKF/server/internal/utils"
+)
+
+var _ urltrie.Hook = (*JWT)(nil)
+
+func init() {
+	urltrie.RegisterHook(&JWT{})
+	fmt.Println("RegisterHook", "Register Hook[JWT] success...")
+}
+
+// JWT implement urltrie.Hook.
+type JWT struct {
+	urltrie.Hook
+}
+
+// Patterns EDIT THIS TO YOUR OWN HOOK PATTERN.
+func (h *JWT) Patterns() []string {
+	return []string{
+		"/api/v1/user/*",
+		"/api/v1/community/*",
+	}
+}
+
+// Priority EDIT THIS TO YOUR OWN HOOK PRIORITY.
+func (h *JWT) Priority() int64 {
+	return 0
+}
+
+// BeforeRun EDIT THIS TO YOUR OWN HOOK BEFORE RUN, DO NOT NEED USE Next() FUNCTION.
+func (h *JWT) BeforeRun(c *gin.Context) {
+	token := c.GetHeader("Authorization")
+	if token == "" || strings.Fields(token)[0] != "Bearer" {
+		c.AbortWithStatus(http.StatusUnauthorized)
+
+		return
+	}
+
+	claims, err := utils.ParseJwtToken(strings.Fields(token)[1])
+	if err != nil {
+		c.AbortWithStatus(http.StatusUnauthorized)
+
+		return
+	}
+
+	// Set claims to context.
+	c.Set("claims", claims)
+}
+
+// AfterRun EDIT THIS TO YOUR OWN HOOK AFTER RUN, DO NOT NEED USE Next() FUNCTION.
+func (h *JWT) AfterRun(c *gin.Context) {
+}

server/internal/middleware/hooks/gen_rate_hook.go

@@ -0,0 +1,106 @@
+// Copyright © 2023 OpenIM open source community. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package hooks
+
+import (
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/gin-gonic/gin"
+
+	urltrie "github.com/OpenIMSDK/OpenKF/server/internal/middleware/hooks/url_trie"
+)
+
+var _ urltrie.Hook = (*RATE)(nil)
+
+func init() {
+	urltrie.RegisterHook(&RATE{})
+	fmt.Println("RegisterHook", "Register Hook[RATE] success...")
+}
+
+// RATE implement urltrie.Hook.
+type RATE struct {
+	urltrie.Hook
+}
+
+// Patterns EDIT THIS TO YOUR OWN HOOK PATTERN.
+func (h *RATE) Patterns() []string {
+	return []string{
+		"",
+	}
+}
+
+// Priority EDIT THIS TO YOUR OWN HOOK PRIORITY.
+func (h *RATE) Priority() int64 {
+	return 0
+}
+
+// BeforeRun EDIT THIS TO YOUR OWN HOOK BEFORE RUN, DO NOT NEED USE Next() FUNCTION.
+func (h *RATE) BeforeRun(c *gin.Context) {
+}
+
+// AfterRun EDIT THIS TO YOUR OWN HOOK AFTER RUN, DO NOT NEED USE Next() FUNCTION.
+func (h *RATE) AfterRun(c *gin.Context) {
+}
+
+type frequencyControlByTokenBucket struct {
+	refreshRate float64    // token refresh rate
+	capacity    int64      // bucket capacity
+	tokens      float64    // token count
+	lastToken   time.Time  // latest time token stored
+	mtx         sync.Mutex // mutex
+}
+
+// allow frequency.
+func (tb *frequencyControlByTokenBucket) Allow() bool {
+	tb.mtx.Lock()
+	defer tb.mtx.Unlock()
+	now := time.Now()
+	// compute tokens which needs
+	tb.tokens = tb.tokens + tb.refreshRate*now.Sub(tb.lastToken).Seconds()
+	if tb.tokens > float64(tb.capacity) {
+		tb.tokens = float64(tb.capacity)
+	}
+	// judge weather to pass through
+	if tb.tokens >= 1 {
+		tb.tokens--
+		tb.lastToken = now
+
+		return true
+	}
+
+	return false
+}
+
+// LimitHandler registried a middle ware to use.
+func LimitHandler(maxConn int, refreshRate float64) gin.HandlerFunc {
+	tb := &frequencyControlByTokenBucket{
+		capacity:    int64(maxConn),
+		refreshRate: refreshRate,
+		tokens:      0,
+		lastToken:   time.Now(),
+	}
+
+	return func(c *gin.Context) {
+		if !tb.Allow() {
+			c.String(503, "Too many request")
+			c.Abort()
+
+			return
+		}
+		c.Next()
+	}
+}

server/internal/middleware/hooks/global_hook.go

@@ -22,8 +22,10 @@ import (
 	urltrie "github.com/OpenIMSDK/OpenKF/server/internal/middleware/hooks/url_trie"
 )
 
+var _ urltrie.Hook = (*GlobalHook)(nil)
+
 func init() {
-	urltrie.RegisterHook(GlobalHook{})
+	urltrie.RegisterHook(&GlobalHook{})
 	fmt.Println("RegisterHook", "Register Hook[GlobalHook] success...")
 }
 
@@ -32,22 +34,22 @@ type GlobalHook struct {
 	urltrie.Hook
 }
 
-// Pattern return pattern.
-func (h GlobalHook) Pattern() string {
-	return "/*"
+// Patterns return pattern.
+func (h *GlobalHook) Patterns() []string {
+	return []string{
+		"/*",
+	}
 }
 
 // Priority return priority.
-func (h GlobalHook) Priority() int64 {
+func (h *GlobalHook) Priority() int64 {
 	return 0
 }
 
 // BeforeRun do something before controller run.
-func (h GlobalHook) BeforeRun(c *gin.Context) {
-	c.Next()
+func (h *GlobalHook) BeforeRun(c *gin.Context) {
 }
 
 // AfterRun do something after controller run.
-func (h GlobalHook) AfterRun(c *gin.Context) {
-	c.Next()
+func (h *GlobalHook) AfterRun(c *gin.Context) {
 }

server/internal/middleware/hooks/mail_hook.go

@@ -23,8 +23,10 @@ import (
 	"github.com/OpenIMSDK/OpenKF/server/pkg/log"
 )
 
+var _ urltrie.Hook = (*MailHook)(nil)
+
 func init() {
-	urltrie.RegisterHook(MailHook{})
+	urltrie.RegisterHook(&MailHook{})
 	fmt.Println("RegisterHook", "Register Hook[MailHook] success...")
 }
 
@@ -33,23 +35,23 @@ type MailHook struct {
 	urltrie.Hook
 }
 
-// Pattern return pattern.
-func (h MailHook) Pattern() string {
-	return "/api/v1/register/email/code"
-}
-
-// BeforeRun do something before controller run.
-func (h MailHook) BeforeRun(c *gin.Context) {
-	log.Debugf("GlobalHook", "path: %v", c.Request.URL.Path)
-	c.Next()
+// Patterns return pattern.
+func (h *MailHook) Patterns() []string {
+	return []string{
+		"/api/v1/register/email/code",
+	}
 }
 
 // Priority return priority.
-func (h MailHook) Priority() int64 {
+func (h *MailHook) Priority() int64 {
 	return 0
 }
 
+// BeforeRun do something before controller run.
+func (h *MailHook) BeforeRun(c *gin.Context) {
+	log.Debugf("GlobalHook", "path: %v", c.Request.URL.Path)
+}
+
 // AfterRun do something after controller run.
-func (h MailHook) AfterRun(c *gin.Context) {
-	c.Next()
+func (h *MailHook) AfterRun(c *gin.Context) {
 }

server/internal/middleware/hooks/url_trie/hook.go

@@ -31,7 +31,7 @@ func init() {
 
 // RegisterHook register url & hook to trie.
 func RegisterHook(hook Hook) {
-	hookTrie.Insert(hook.Pattern(), hook)
+	hookTrie.InsertBatch(hook.Patterns(), hook)
 }
 
 // RunHook enable hook for interceptor.
@@ -39,7 +39,7 @@ func RunHook() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		raw := c.Request.URL.Path
 
-		// get path from url
+		// Get path from url
 		p, err := url.Parse(raw)
 		if err != nil {
 			log.Errorf("Hook", "parse url error: %v", err)
@@ -53,10 +53,16 @@ func RunHook() gin.HandlerFunc {
 			return
 		}
 
-		// run hooks
+		// Run all before hooks
 		for _, hook := range hooks {
 			hook.BeforeRun(c)
-			c.Next()
+		}
+
+		// Run controllers
+		c.Next()
+
+		// Run all after hooks
+		for _, hook := range hooks {
 			hook.AfterRun(c)
 		}
 	}

server/internal/middleware/hooks/url_trie/trie.go

@@ -26,7 +26,9 @@ type Hook interface {
 	// Register with url pattern
 	// Support * wildcard, you can use it like this:
 	// /api/v1/*, and the url like /api/v1/123 will be matched
-	Pattern() string
+	//
+	// Now we support multi urls.
+	Patterns() []string
 
 	// Priority will set the priority of the hook
 	Priority() int64
@@ -72,8 +74,15 @@ func NewTrie() *Trie {
 	}
 }
 
+// InsertBatch insert urls with hooks.
+func (t *Trie) InsertBatch(urls []string, hooks ...Hook) {
+	for _, url := range urls {
+		t.insert(url, hooks...)
+	}
+}
+
 // Insert insert url with hooks.
-func (t *Trie) Insert(url string, hooks ...Hook) {
+func (t *Trie) insert(url string, hooks ...Hook) {
 	current := t.root
 
 	// split url with '/'

server/internal/middleware/hooks/url_trie/trie_test.go

@@ -23,28 +23,38 @@ import (
 type testHook struct {
 	Hook
 	priority int64
-	url      string
+	urls     []string
 }
 
-func (h testHook) Priority() int64 {
+func (h *testHook) Priority() int64 {
 	return h.priority
 }
 
-func (h testHook) Pattern() string {
-	return h.url
+func (h *testHook) Patterns() []string {
+	return h.urls
 }
 
 func TestUrlTrie(t *testing.T) {
-	hooks := []Hook{
-		testHook{priority: 1, url: "/gin"},
-		testHook{priority: 1, url: "/api/v1/123"},
-		testHook{priority: 1, url: "/openkf/*"},
-		testHook{priority: 2, url: "/gin"},
-		testHook{priority: 3, url: "/gin/1"},
+	// test case
+	testData := []struct {
+		priority int64
+		url      []string
+	}{
+		{1, []string{"/gin"}},
+		{1, []string{"/api/v1/123"}},
+		{1, []string{"/openkf/*"}},
+		{2, []string{"/gin"}},
+		{3, []string{"/gin/1"}},
 	}
+
 	trie := NewTrie()
-	for _, h := range hooks {
-		trie.Insert(h.Pattern(), h)
+
+	// range test data
+	for _, data := range testData {
+		trie.InsertBatch(data.url, &testHook{
+			urls:     data.url,
+			priority: data.priority,
+		})
 	}
 
 	values, matched := trie.Match("/gin/1")