From 4b053e43e4611427dc1bb6a08bbfb7c9fe0e0f77 Mon Sep 17 00:00:00 2001 From: Asklv <47499836+IRONICBo@users.noreply.github.com> Date: Wed, 19 Jul 2023 21:29:20 +0800 Subject: [PATCH] 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 | 2 +- server/cmd/genhooks/pkg/template.go | 10 +- .../middleware/hooks/gen_cors_hook.go | 67 +++++++++++ .../internal/middleware/hooks/gen_jwt_hook.go | 75 +++++++++++++ .../middleware/hooks/gen_rate_hook.go | 106 ++++++++++++++++++ .../internal/middleware/hooks/global_hook.go | 20 ++-- server/internal/middleware/hooks/mail_hook.go | 28 ++--- .../middleware/hooks/url_trie/hook.go | 14 ++- .../middleware/hooks/url_trie/trie.go | 13 ++- .../middleware/hooks/url_trie/trie_test.go | 34 ++++-- 10 files changed, 323 insertions(+), 46 deletions(-) create mode 100644 server/internal/middleware/hooks/gen_cors_hook.go create mode 100644 server/internal/middleware/hooks/gen_jwt_hook.go create mode 100644 server/internal/middleware/hooks/gen_rate_hook.go diff --git a/server/cmd/dbmigration/main.go b/server/cmd/dbmigration/main.go index 457092f..3253657 100644 --- a/server/cmd/dbmigration/main.go +++ b/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. diff --git a/server/cmd/genhooks/pkg/template.go b/server/cmd/genhooks/pkg/template.go index ee5886e..04ca747 100644 --- a/server/cmd/genhooks/pkg/template.go +++ b/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() } diff --git a/server/internal/middleware/hooks/gen_cors_hook.go b/server/internal/middleware/hooks/gen_cors_hook.go new file mode 100644 index 0000000..64caae9 --- /dev/null +++ b/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) { +} diff --git a/server/internal/middleware/hooks/gen_jwt_hook.go b/server/internal/middleware/hooks/gen_jwt_hook.go new file mode 100644 index 0000000..6eab98c --- /dev/null +++ b/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) { +} diff --git a/server/internal/middleware/hooks/gen_rate_hook.go b/server/internal/middleware/hooks/gen_rate_hook.go new file mode 100644 index 0000000..ff874d3 --- /dev/null +++ b/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() + } +} diff --git a/server/internal/middleware/hooks/global_hook.go b/server/internal/middleware/hooks/global_hook.go index d2f67ed..f2a3bc3 100644 --- a/server/internal/middleware/hooks/global_hook.go +++ b/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) { } diff --git a/server/internal/middleware/hooks/mail_hook.go b/server/internal/middleware/hooks/mail_hook.go index 5b42a64..63f3320 100644 --- a/server/internal/middleware/hooks/mail_hook.go +++ b/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) { } diff --git a/server/internal/middleware/hooks/url_trie/hook.go b/server/internal/middleware/hooks/url_trie/hook.go index 2bcd7a9..75fb017 100644 --- a/server/internal/middleware/hooks/url_trie/hook.go +++ b/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) } } diff --git a/server/internal/middleware/hooks/url_trie/trie.go b/server/internal/middleware/hooks/url_trie/trie.go index cf464ac..cd51ee7 100644 --- a/server/internal/middleware/hooks/url_trie/trie.go +++ b/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 '/' diff --git a/server/internal/middleware/hooks/url_trie/trie_test.go b/server/internal/middleware/hooks/url_trie/trie_test.go index ba86753..9be1444 100644 --- a/server/internal/middleware/hooks/url_trie/trie_test.go +++ b/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")