├── .github ├── ISSUE_TEMPLATE │ ├── bug_report.md │ ├── feature_request.md │ └── question.md ├── dependabot.yml └── workflows │ ├── codeql-analysis.yml │ ├── issue-translator.yml │ └── issues.yml ├── .gitignore ├── README.md ├── auth └── casbin │ ├── README.md │ ├── auth_model.conf │ ├── auth_model_domain.conf │ ├── auth_policy.csv │ ├── auth_policy_domain.csv │ ├── casbin.go │ ├── casbin_test.go │ ├── go.mod │ └── go.sum ├── configcenter └── apollo │ ├── README.md │ ├── apollo.go │ ├── apollo_integration_test.go │ ├── apollo_test.go │ └── examples │ └── main.go ├── go.mod ├── go.sum ├── handler ├── README.md ├── etag.go ├── etag_test.go ├── go.mod ├── go.sum ├── header_types.go └── header_types_test.go ├── logx ├── logrusx │ ├── go.mod │ ├── go.sum │ ├── logrus.go │ └── readme.md ├── zapx │ ├── go.mod │ ├── go.sum │ ├── readme.md │ └── zap.go └── zerologx │ ├── go.mod │ ├── go.sum │ ├── readme.md │ ├── zerolog.go │ └── zerolog_test.go ├── rest └── registry │ └── etcd │ ├── README.md │ ├── go.mod │ ├── go.sum │ └── register.go ├── router ├── chi │ ├── README.md │ ├── go.mod │ ├── go.sum │ ├── router.go │ └── router_test.go ├── gin │ ├── README.md │ ├── config.go │ ├── go.mod │ ├── go.sum │ ├── router.go │ └── router_test.go └── mux │ ├── README.md │ ├── go.mod │ ├── go.sum │ ├── router.go │ └── router_test.go ├── stores ├── clickhouse │ ├── clickhouse.go │ ├── clickhouse_test.go │ ├── go.mod │ └── go.sum └── mongo │ ├── go.mod │ ├── go.sum │ ├── mon │ ├── bulkinserter.go │ ├── bulkinserter_test.go │ ├── clientmanager.go │ ├── clientmanager_test.go │ ├── collection.go │ ├── collection_test.go │ ├── model.go │ ├── model_test.go │ ├── options.go │ ├── options_test.go │ ├── trace.go │ ├── util.go │ └── util_test.go │ └── monc │ ├── cachedmodel.go │ └── cachedmodel_test.go └── zrpc └── registry ├── consul ├── README.md ├── builder.go ├── config.go ├── go.mod ├── go.sum ├── register.go ├── resolver.go ├── target.go └── tests │ └── client_test.go ├── nacos ├── README.md ├── builder.go ├── go.mod ├── go.sum ├── options.go ├── register.go ├── resolver.go └── target.go └── polaris ├── README.md ├── builder.go ├── constant.go ├── go.mod ├── go.sum ├── options.go ├── register.go ├── resolver.go └── target.go /.github/ISSUE_TEMPLATE/bug_report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Bug report 3 | about: Create a report to help us improve 4 | title: '' 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | **Describe the bug** 11 | A clear and concise description of what the bug is. 12 | 13 | **To Reproduce** 14 | Steps to reproduce the behavior, if applicable: 15 | 16 | 1. The code is 17 | 18 | ```go 19 | 20 | ``` 21 | 22 | 2. The error is 23 | 24 | ``` 25 | 26 | ``` 27 | 28 | **Expected behavior** 29 | A clear and concise description of what you expected to happen. 30 | 31 | **Screenshots** 32 | If applicable, add screenshots to help explain your problem. 33 | 34 | **Environments (please complete the following information):** 35 | - OS: [e.g. Linux] 36 | - go-zero version [e.g. 1.2.1] 37 | - goctl version [e.g. 1.2.1, optional] 38 | 39 | **More description** 40 | Add any other context about the problem here. 41 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Feature request 3 | about: Suggest an idea for this project 4 | title: '' 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | **Is your feature request related to a problem? Please describe.** 11 | A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] 12 | 13 | **Describe the solution you'd like** 14 | A clear and concise description of what you want to happen. 15 | 16 | **Describe alternatives you've considered** 17 | A clear and concise description of any alternative solutions or features you've considered. 18 | 19 | **Additional context** 20 | Add any other context or screenshots about the feature request here. 21 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/question.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Question 3 | about: Ask a question on using go-zero or goctl 4 | title: '' 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | 11 | -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | # To get started with Dependabot version updates, you'll need to specify which 2 | # package ecosystems to update and where the package manifests are located. 3 | # Please see the documentation for all configuration options: 4 | # https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates 5 | 6 | version: 2 7 | updates: 8 | - package-ecosystem: "gomod" # See documentation for possible values 9 | directory: "/auth/casbin" # Location of package manifests 10 | schedule: 11 | interval: "daily" 12 | - package-ecosystem: "gomod" # See documentation for possible values 13 | directory: "/handler" # Location of package manifests 14 | schedule: 15 | interval: "daily" 16 | - package-ecosystem: "gomod" # See documentation for possible values 17 | directory: "/logx/logrusx" # Location of package manifests 18 | schedule: 19 | interval: "daily" 20 | - package-ecosystem: "gomod" # See documentation for possible values 21 | directory: "/logx/zapx" # Location of package manifests 22 | schedule: 23 | interval: "daily" 24 | - package-ecosystem: "gomod" # See documentation for possible values 25 | directory: "/logx/zerologx" # Location of package manifests 26 | schedule: 27 | interval: "daily" 28 | - package-ecosystem: "gomod" # See documentation for possible values 29 | directory: "/rest/registry/etcd" # Location of package manifests 30 | schedule: 31 | interval: "daily" 32 | - package-ecosystem: "gomod" # See documentation for possible values 33 | directory: "/router/chi" # Location of package manifests 34 | schedule: 35 | interval: "daily" 36 | - package-ecosystem: "gomod" # See documentation for possible values 37 | directory: "/router/gin" # Location of package manifests 38 | schedule: 39 | interval: "daily" 40 | - package-ecosystem: "gomod" # See documentation for possible values 41 | directory: "/router/mux" # Location of package manifests 42 | schedule: 43 | interval: "daily" 44 | - package-ecosystem: "gomod" # See documentation for possible values 45 | directory: "/stores/clickhouse" # Location of package manifests 46 | schedule: 47 | interval: "daily" 48 | - package-ecosystem: "gomod" # See documentation for possible values 49 | directory: "/zrpc/registry/consul" # Location of package manifests 50 | schedule: 51 | interval: "daily" 52 | - package-ecosystem: "gomod" # See documentation for possible values 53 | directory: "/zrpc/registry/nacos" # Location of package manifests 54 | schedule: 55 | interval: "daily" 56 | - package-ecosystem: "gomod" # See documentation for possible values 57 | directory: "/zrpc/registry/polaris" # Location of package manifests 58 | schedule: 59 | interval: "daily" 60 | -------------------------------------------------------------------------------- /.github/workflows/codeql-analysis.yml: -------------------------------------------------------------------------------- 1 | # For most projects, this workflow file will not need changing; you simply need 2 | # to commit it to your repository. 3 | # 4 | # You may wish to alter this file to override the set of languages analyzed, 5 | # or to provide custom queries or build logic. 6 | # 7 | # ******** NOTE ******** 8 | # We have attempted to detect the languages in your repository. Please check 9 | # the `language` matrix defined below to confirm you have the correct set of 10 | # supported CodeQL languages. 11 | # 12 | name: "CodeQL" 13 | 14 | on: 15 | push: 16 | branches: [ main ] 17 | pull_request: 18 | # The branches below must be a subset of the branches above 19 | branches: [ main ] 20 | schedule: 21 | - cron: '18 19 * * 6' 22 | 23 | jobs: 24 | analyze: 25 | name: Analyze 26 | runs-on: ubuntu-latest 27 | 28 | strategy: 29 | fail-fast: false 30 | matrix: 31 | language: [ 'go' ] 32 | # CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python' ] 33 | # Learn more: 34 | # https://docs.github.com/en/free-pro-team@latest/github/finding-security-vulnerabilities-and-errors-in-your-code/configuring-code-scanning#changing-the-languages-that-are-analyzed 35 | 36 | steps: 37 | - name: Checkout repository 38 | uses: actions/checkout@v3 39 | 40 | # Initializes the CodeQL tools for scanning. 41 | - name: Initialize CodeQL 42 | uses: github/codeql-action/init@v2 43 | with: 44 | languages: ${{ matrix.language }} 45 | # If you wish to specify custom queries, you can do so here or in a config file. 46 | # By default, queries listed here will override any specified in a config file. 47 | # Prefix the list here with "+" to use these queries and those in the config file. 48 | # queries: ./path/to/local/query, your-org/your-repo/queries@main 49 | 50 | # Autobuild attempts to build any compiled languages (C/C++, C#, or Java). 51 | # If this step fails, then you should remove it and run the build manually (see below) 52 | - name: Autobuild 53 | uses: github/codeql-action/autobuild@v2 54 | 55 | # ℹ️ Command-line programs to run using the OS shell. 56 | # 📚 https://git.io/JvXDl 57 | 58 | # ✏️ If the Autobuild fails above, remove it and uncomment the following three lines 59 | # and modify them (or add more) to build your code if your project 60 | # uses a compiled language 61 | 62 | #- run: | 63 | # make bootstrap 64 | # make release 65 | 66 | - name: Perform CodeQL Analysis 67 | uses: github/codeql-action/analyze@v2 68 | -------------------------------------------------------------------------------- /.github/workflows/issue-translator.yml: -------------------------------------------------------------------------------- 1 | name: 'issue-translator' 2 | on: 3 | issue_comment: 4 | types: [created] 5 | issues: 6 | types: [opened] 7 | 8 | jobs: 9 | build: 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: usthe/issues-translate-action@v2.7 13 | with: 14 | IS_MODIFY_TITLE: true 15 | # not require, default false, . Decide whether to modify the issue title 16 | # if true, the robot account @Issues-translate-bot must have modification permissions, invite @Issues-translate-bot to your project or use your custom bot. 17 | CUSTOM_BOT_NOTE: Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿 18 | # not require. Customize the translation robot prefix message. 19 | -------------------------------------------------------------------------------- /.github/workflows/issues.yml: -------------------------------------------------------------------------------- 1 | name: Close inactive issues 2 | on: 3 | schedule: 4 | - cron: "30 1 * * *" 5 | 6 | jobs: 7 | close-issues: 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/stale@v6 11 | with: 12 | days-before-issue-stale: 365 13 | days-before-issue-close: 90 14 | stale-issue-label: "stale" 15 | stale-issue-message: "This issue is stale because it has been open for 30 days with no activity." 16 | close-issue-message: "This issue was closed because it has been inactive for 14 days since being marked as stale." 17 | days-before-pr-stale: -1 18 | days-before-pr-close: -1 19 | repo-token: ${{ secrets.GITHUB_TOKEN }} 20 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Binaries for programs and plugins 2 | *.exe 3 | *.exe~ 4 | *.dll 5 | *.so 6 | *.dylib 7 | 8 | # IDE integration 9 | /.vscode/* 10 | /.idea/* 11 | 12 | # Mac DS_Store 13 | .DS_Store 14 | 15 | # Test binary, built with `go test -c` 16 | *.test 17 | 18 | # Output of the go coverage tool, specifically when used with LiteIDE 19 | *.out 20 | 21 | # Dependency directories (remove the comment below to include it) 22 | vendor/ 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # gozero-contrib 2 | A collection of extensions and tools for go-zero 3 | -------------------------------------------------------------------------------- /auth/casbin/README.md: -------------------------------------------------------------------------------- 1 | ### Quick Start 2 | 3 | Prerequisites: 4 | 5 | * Install `go-zero`: go get -u github.com/zeromicro/go-zero 6 | 7 | Download the module: 8 | 9 | ```shell 10 | go get -u github.com/zeromicro/zero-contrib/auth/casbin 11 | ``` 12 | 13 | For example: 14 | 15 | ```go 16 | package main 17 | 18 | import ( 19 | stdcasbin "github.com/casbin/casbin/v2" 20 | "github.com/zeromicro/go-zero/rest" 21 | "github.com/zeromicro/zero-contrib/auth/casbin" 22 | ) 23 | 24 | func main() { 25 | // load the casbin model and policy from files, database is also supported. 26 | e, _ := stdcasbin.NewEnforcer("auth_model.conf", "auth_policy.csv") 27 | 28 | // define your router, and use the Casbin auth middleware. 29 | // the access that is denied by auth will return HTTP 403 error. 30 | // set username as the user unique identity field. 31 | authorizer := casbin.NewAuthorizer(e, casbin.WithUidField("username")) 32 | conf := rest.RestConf{} 33 | server := rest.MustNewServer(conf) 34 | server.Use(rest.ToMiddleware(authorizer)) 35 | } 36 | 37 | ``` 38 | 39 | documentation: 40 | 41 | The authorization determines a request based on {subject, object, action}, which means what subject can perform what 42 | action on what object. In this plugin, the meanings are: 43 | 44 | - subject: It comes from user unique identity in JWT Claims. 45 | - object: the URL path for the web resource like "dataset1/item1". 46 | - action: HTTP method like GET, POST, PUT, DELETE, or the high-level actions you defined like "read-file", "write-blog". 47 | 48 | For how to write authorization policy and other details, please refer 49 | to [the Casbin's documentation](https://github.com/casbin/casbin). 50 | 51 | gratitude: 52 | 53 | - https://github.com/gin-contrib/authz -------------------------------------------------------------------------------- /auth/casbin/auth_model.conf: -------------------------------------------------------------------------------- 1 | [request_definition] 2 | r = sub, obj, act 3 | 4 | [policy_definition] 5 | p = sub, obj, act 6 | 7 | [role_definition] 8 | g = _, _ 9 | 10 | [policy_effect] 11 | e = some(where (p.eft == allow)) 12 | 13 | [matchers] 14 | m = g(r.sub, p.sub) && keyMatch(r.obj, p.obj) && (r.act == p.act || p.act == "*") -------------------------------------------------------------------------------- /auth/casbin/auth_model_domain.conf: -------------------------------------------------------------------------------- 1 | [request_definition] 2 | r = sub, dom, obj, act 3 | 4 | [policy_definition] 5 | p = sub, dom, obj, act 6 | 7 | [role_definition] 8 | g = _, _, _ 9 | 10 | [policy_effect] 11 | e = some(where (p.eft == allow)) 12 | 13 | [matchers] 14 | m = g(r.sub, p.sub, r.dom) && r.dom == p.dom && r.obj == p.obj && (r.act == p.act|| p.act =="*") -------------------------------------------------------------------------------- /auth/casbin/auth_policy.csv: -------------------------------------------------------------------------------- 1 | p, alice, /dataset1/*, GET 2 | p, alice, /dataset1/resource1, POST 3 | p, bob, /dataset2/resource1, * 4 | p, bob, /dataset2/resource2, GET 5 | p, bob, /dataset2/folder1/*, POST 6 | p, dataset1_admin, /dataset1/*, * 7 | g, cathy, dataset1_admin -------------------------------------------------------------------------------- /auth/casbin/auth_policy_domain.csv: -------------------------------------------------------------------------------- 1 | p, alice,go-zero, /dataset1/resource1, POST 2 | p, admin,go-zero, /dataset1/resource2, * 3 | p, bob,domain1, /dataset2/resource1, POST 4 | p, bob,domain2, /dataset2/resource2, GET 5 | p, bob,go-zero, /dataset2/folder1/*, POST 6 | p, dataset1_admin,domain2, /dataset1/*, * 7 | 8 | g, cathy,domain1, dataset1_admin 9 | g, alice, admin, go-zero 10 | g, bob, admin, domain2 -------------------------------------------------------------------------------- /auth/casbin/casbin.go: -------------------------------------------------------------------------------- 1 | package casbin 2 | 3 | import ( 4 | "log" 5 | "net/http" 6 | 7 | "github.com/casbin/casbin/v2" 8 | "github.com/zeromicro/go-zero/core/logx" 9 | ) 10 | 11 | type ( 12 | // Authorizer stores the casbin handler 13 | Authorizer struct { 14 | enforcer *casbin.Enforcer 15 | uidField string 16 | domain string 17 | } 18 | // AuthorizerOption represents an option. 19 | AuthorizerOption func(opt *Authorizer) 20 | ) 21 | 22 | // WithUidField returns a custom user unique identity option. 23 | func WithUidField(uidField string) AuthorizerOption { 24 | return func(opt *Authorizer) { 25 | opt.uidField = uidField 26 | } 27 | } 28 | 29 | // WithDomain returns a custom domain option. 30 | func WithDomain(domain string) AuthorizerOption { 31 | return func(opt *Authorizer) { 32 | opt.domain = domain 33 | } 34 | } 35 | 36 | // NewAuthorizer returns the authorizer, uses a Casbin enforcer as input 37 | func NewAuthorizer(e *casbin.Enforcer, opts ...AuthorizerOption) func(http.Handler) http.Handler { 38 | a := &Authorizer{enforcer: e} 39 | // init an Authorizer 40 | a.init(opts...) 41 | 42 | return func(next http.Handler) http.Handler { 43 | return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { 44 | if !a.CheckPermission(request) { 45 | a.RequirePermission(writer) 46 | return 47 | } 48 | next.ServeHTTP(writer, request) 49 | }) 50 | } 51 | } 52 | 53 | func (a *Authorizer) init(opts ...AuthorizerOption) { 54 | a.uidField = "username" 55 | a.domain = "domain" 56 | for _, opt := range opts { 57 | opt(a) 58 | } 59 | } 60 | 61 | // GetUid gets the uid from the JWT Claims. 62 | func (a *Authorizer) GetUid(r *http.Request) (string, bool) { 63 | uid, ok := r.Context().Value(a.uidField).(string) 64 | return uid, ok 65 | } 66 | 67 | // GetDomain returns the domain from the request. 68 | func (a *Authorizer) GetDomain(r *http.Request) (string, bool) { 69 | domain, ok := r.Context().Value(a.domain).(string) 70 | return domain, ok 71 | } 72 | 73 | // CheckPermission checks the user/method/path combination from the request. 74 | // Returns true (permission granted) or false (permission forbidden) 75 | func (a *Authorizer) CheckPermission(r *http.Request) bool { 76 | uid, ok := a.GetUid(r) 77 | if !ok { 78 | return false 79 | } 80 | method := r.Method 81 | path := r.URL.Path 82 | var ( 83 | allowed = false 84 | err error 85 | ) 86 | domain, withDomain := a.GetDomain(r) 87 | log.Println("domain:", domain) 88 | if withDomain { 89 | allowed, err = a.enforcer.Enforce(uid, domain, path, method) 90 | } else { 91 | allowed, err = a.enforcer.Enforce(uid, path, method) 92 | } 93 | 94 | if err != nil { 95 | logx.WithContext(r.Context()).Errorf("[CASBIN] enforce err %s", err.Error()) 96 | } 97 | return allowed 98 | } 99 | 100 | // RequirePermission returns the 403 Forbidden to the client. 101 | func (a *Authorizer) RequirePermission(writer http.ResponseWriter) { 102 | writer.WriteHeader(http.StatusForbidden) 103 | 104 | } 105 | -------------------------------------------------------------------------------- /auth/casbin/casbin_test.go: -------------------------------------------------------------------------------- 1 | package casbin 2 | 3 | import ( 4 | "context" 5 | "net/http" 6 | "net/http/httptest" 7 | "testing" 8 | 9 | "github.com/casbin/casbin/v2" 10 | "github.com/stretchr/testify/assert" 11 | ) 12 | 13 | func testAuthWithUsernameRequest(t *testing.T, router http.Handler, user string, path string, method string, code int) { 14 | r, _ := http.NewRequestWithContext(context.Background(), method, path, nil) 15 | request := r.WithContext(context.WithValue(r.Context(), "username", user)) 16 | w := httptest.NewRecorder() 17 | router.ServeHTTP(w, request) 18 | 19 | if w.Code != code { 20 | t.Errorf("%s, %s, %s: %d, supposed to be %d", user, path, method, w.Code, code) 21 | } 22 | } 23 | func testDomainAuthWithUsernameRequest(t *testing.T, router http.Handler, user string, domain string, path string, method string, code int) { 24 | r, _ := http.NewRequestWithContext(context.Background(), method, path, nil) 25 | ctx := context.WithValue(r.Context(), "username", user) 26 | request := r.WithContext(context.WithValue(ctx, "domain", domain)) 27 | w := httptest.NewRecorder() 28 | router.ServeHTTP(w, request) 29 | 30 | if w.Code != code { 31 | t.Errorf("%s, %s,%s, %s: %d, supposed to be %d", user, domain, path, method, w.Code, code) 32 | } 33 | } 34 | 35 | func TestBasic(t *testing.T) { 36 | e, _ := casbin.NewEnforcer("auth_model.conf", "auth_policy.csv") 37 | router := NewAuthorizer(e, WithUidField("username"))( 38 | http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 39 | w.Header().Set("X-Test", "test") 40 | _, err := w.Write([]byte("content")) 41 | assert.Nil(t, err) 42 | 43 | flusher, ok := w.(http.Flusher) 44 | assert.True(t, ok) 45 | flusher.Flush() 46 | })) 47 | testAuthWithUsernameRequest(t, router, "alice", "/dataset1/resource1", "GET", 200) 48 | testAuthWithUsernameRequest(t, router, "alice", "/dataset1/resource1", "POST", 200) 49 | testAuthWithUsernameRequest(t, router, "alice", "/dataset1/resource2", "GET", 200) 50 | testAuthWithUsernameRequest(t, router, "alice", "/dataset1/resource2", "POST", 403) 51 | } 52 | 53 | func TestBasicDomain(t *testing.T) { 54 | e, err := casbin.NewEnforcer("auth_model_domain.conf", "auth_policy_domain.csv") 55 | if err != nil { 56 | t.Fatal(err) 57 | } 58 | router := NewAuthorizer(e, WithUidField("username"), WithDomain("domain"))( 59 | http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 60 | w.Header().Set("X-Test", "test") 61 | _, err := w.Write([]byte("content")) 62 | assert.Nil(t, err) 63 | 64 | flusher, ok := w.(http.Flusher) 65 | assert.True(t, ok) 66 | flusher.Flush() 67 | })) 68 | testDomainAuthWithUsernameRequest(t, router, "alice", "go-zero", "/dataset1/resource1", "POST", 200) 69 | testDomainAuthWithUsernameRequest(t, router, "bob", "domain1", "/dataset2/resource1", "POST", 200) 70 | testDomainAuthWithUsernameRequest(t, router, "alice", "go-zero", "/dataset1/resource2", "POST", 200) 71 | } 72 | 73 | func TestPathWildcard(t *testing.T) { 74 | e, _ := casbin.NewEnforcer("auth_model.conf", "auth_policy.csv") 75 | router := NewAuthorizer(e, WithUidField("username"))( 76 | http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 77 | w.Header().Set("X-Test", "test") 78 | _, err := w.Write([]byte("content")) 79 | assert.Nil(t, err) 80 | 81 | flusher, ok := w.(http.Flusher) 82 | assert.True(t, ok) 83 | flusher.Flush() 84 | })) 85 | 86 | testAuthWithUsernameRequest(t, router, "bob", "/dataset2/resource1", "GET", 200) 87 | testAuthWithUsernameRequest(t, router, "bob", "/dataset2/resource1", "POST", 200) 88 | testAuthWithUsernameRequest(t, router, "bob", "/dataset2/resource1", "DELETE", 200) 89 | testAuthWithUsernameRequest(t, router, "bob", "/dataset2/resource2", "GET", 200) 90 | testAuthWithUsernameRequest(t, router, "bob", "/dataset2/resource2", "POST", 403) 91 | testAuthWithUsernameRequest(t, router, "bob", "/dataset2/resource2", "DELETE", 403) 92 | 93 | testAuthWithUsernameRequest(t, router, "bob", "/dataset2/folder1/item1", "GET", 403) 94 | testAuthWithUsernameRequest(t, router, "bob", "/dataset2/folder1/item1", "POST", 200) 95 | testAuthWithUsernameRequest(t, router, "bob", "/dataset2/folder1/item1", "DELETE", 403) 96 | testAuthWithUsernameRequest(t, router, "bob", "/dataset2/folder1/item2", "GET", 403) 97 | testAuthWithUsernameRequest(t, router, "bob", "/dataset2/folder1/item2", "POST", 200) 98 | testAuthWithUsernameRequest(t, router, "bob", "/dataset2/folder1/item2", "DELETE", 403) 99 | } 100 | 101 | func TestRBAC(t *testing.T) { 102 | e, _ := casbin.NewEnforcer("auth_model.conf", "auth_policy.csv") 103 | router := NewAuthorizer(e, WithUidField("username"))( 104 | http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 105 | w.Header().Set("X-Test", "test") 106 | _, err := w.Write([]byte("content")) 107 | assert.Nil(t, err) 108 | 109 | flusher, ok := w.(http.Flusher) 110 | assert.True(t, ok) 111 | flusher.Flush() 112 | })) 113 | 114 | // cathy can access all /dataset1/* resources via all methods because it has the dataset1_admin role. 115 | testAuthWithUsernameRequest(t, router, "cathy", "/dataset1/item", "GET", 200) 116 | testAuthWithUsernameRequest(t, router, "cathy", "/dataset1/item", "POST", 200) 117 | testAuthWithUsernameRequest(t, router, "cathy", "/dataset1/item", "DELETE", 200) 118 | testAuthWithUsernameRequest(t, router, "cathy", "/dataset2/item", "GET", 403) 119 | testAuthWithUsernameRequest(t, router, "cathy", "/dataset2/item", "POST", 403) 120 | testAuthWithUsernameRequest(t, router, "cathy", "/dataset2/item", "DELETE", 403) 121 | 122 | // delete all roles on user cathy, so cathy cannot access any resources now. 123 | _, err := e.DeleteRolesForUser("cathy") 124 | if err != nil { 125 | t.Errorf("got error %v", err) 126 | } 127 | 128 | testAuthWithUsernameRequest(t, router, "cathy", "/dataset1/item", "GET", 403) 129 | testAuthWithUsernameRequest(t, router, "cathy", "/dataset1/item", "POST", 403) 130 | testAuthWithUsernameRequest(t, router, "cathy", "/dataset1/item", "DELETE", 403) 131 | testAuthWithUsernameRequest(t, router, "cathy", "/dataset2/item", "GET", 403) 132 | testAuthWithUsernameRequest(t, router, "cathy", "/dataset2/item", "POST", 403) 133 | testAuthWithUsernameRequest(t, router, "cathy", "/dataset2/item", "DELETE", 403) 134 | } 135 | 136 | func TestUsernameNotFounded(t *testing.T) { 137 | e, _ := casbin.NewEnforcer("auth_model.conf", "auth_policy.csv") 138 | router := NewAuthorizer(e, WithUidField("username"))( 139 | http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 140 | w.Header().Set("X-Test", "test") 141 | _, err := w.Write([]byte("content")) 142 | assert.Nil(t, err) 143 | 144 | flusher, ok := w.(http.Flusher) 145 | assert.True(t, ok) 146 | flusher.Flush() 147 | })) 148 | 149 | r, _ := http.NewRequestWithContext(context.Background(), "GET", "/dataset1/resource1", nil) 150 | w := httptest.NewRecorder() 151 | 152 | router.ServeHTTP(w, r) 153 | assert.EqualValues(t, 403, w.Code) 154 | } 155 | -------------------------------------------------------------------------------- /auth/casbin/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/auth/casbin 2 | 3 | go 1.16 4 | 5 | require ( 6 | github.com/casbin/casbin/v2 v2.73.0 7 | github.com/stretchr/testify v1.8.4 8 | github.com/zeromicro/go-zero v1.5.6 9 | ) 10 | -------------------------------------------------------------------------------- /configcenter/apollo/README.md: -------------------------------------------------------------------------------- 1 | # Apollo Config Center Integration 2 | 3 | Apollo configuration center integration for go-zero framework. 4 | 5 | ## Installation 6 | 7 | ```bash 8 | go get github.com/zeromicro/zero-contrib/configcenter/apollo 9 | ``` 10 | 11 | ## Quick Start 12 | 13 | ```go 14 | package main 15 | 16 | import ( 17 | "fmt" 18 | 19 | configurator "github.com/zeromicro/go-zero/core/configcenter" 20 | "github.com/zeromicro/zero-contrib/configcenter/apollo" 21 | ) 22 | 23 | type AppConfig struct { 24 | Name string `json:"name"` 25 | Timeout int64 `json:"timeout"` 26 | } 27 | 28 | func main() { 29 | // Create Apollo subscriber 30 | sub := apollo.MustNewApolloSubscriber(apollo.ApolloConf{ 31 | AppID: "your-app-id", 32 | Cluster: "default", 33 | NamespaceName: "application.json", 34 | MetaAddr: "http://localhost:8080", 35 | Format: "json", 36 | }) 37 | 38 | // Create config center 39 | cc := configurator.MustNewConfigCenter[AppConfig]( 40 | configurator.Config{Type: "json"}, 41 | sub, 42 | ) 43 | 44 | // Get config 45 | config, _ := cc.GetConfig() 46 | fmt.Printf("Config: %+v\n", config) 47 | 48 | // Listen for changes 49 | cc.AddListener(func() { 50 | newConfig, _ := cc.GetConfig() 51 | fmt.Printf("Config updated: %+v\n", newConfig) 52 | }) 53 | 54 | select {} // Keep running 55 | } 56 | ``` 57 | 58 | ## Configuration 59 | 60 | ### ApolloConf 61 | 62 | | Field | Type | Required | Default | Description | 63 | |-------|------|----------|---------|-------------| 64 | | AppID | string | Yes | - | Apollo application ID | 65 | | MetaAddr | string | Yes | - | Apollo meta server address | 66 | | Cluster | string | No | "default" | Cluster name | 67 | | NamespaceName | string | No | "application" | Namespace name | 68 | | Format | string | No | "json" | Config format: json/yaml/properties | 69 | | Key | string | No | - | Specific key to watch (empty = entire namespace) | 70 | | Secret | string | No | - | Secret key for authentication | 71 | | IsBackupConfig | bool | No | false | Enable local backup | 72 | | BackupPath | string | No | - | Backup directory path | 73 | 74 | ## Examples 75 | 76 | See [examples](./examples) directory for more examples. 77 | 78 | ## References 79 | 80 | - [Apollo Documentation](https://www.apolloconfig.com/) 81 | - [go-zero Config Center](https://go-zero.dev/docs/tutorials/configcenter/overview) 82 | -------------------------------------------------------------------------------- /configcenter/apollo/apollo.go: -------------------------------------------------------------------------------- 1 | package apollo 2 | 3 | import ( 4 | "encoding/json" 5 | "errors" 6 | "sync" 7 | 8 | "github.com/apolloconfig/agollo/v4" 9 | "github.com/apolloconfig/agollo/v4/env/config" 10 | "github.com/apolloconfig/agollo/v4/storage" 11 | "github.com/zeromicro/go-zero/core/configcenter/subscriber" 12 | "github.com/zeromicro/go-zero/core/logx" 13 | "gopkg.in/yaml.v2" 14 | ) 15 | 16 | type ( 17 | // ApolloConf is the configuration for Apollo. 18 | ApolloConf struct { 19 | AppID string 20 | Cluster string `json:",default=default"` 21 | NamespaceName string `json:",default=application"` 22 | IP string `json:",optional"` 23 | MetaAddr string // Apollo meta server address (required) 24 | Secret string `json:",optional"` 25 | IsBackupConfig bool `json:",optional"` 26 | BackupPath string `json:",optional"` 27 | MustStart bool `json:",optional"` 28 | Format string `json:",default=json,options=json|yaml|properties"` 29 | Key string `json:",optional"` // Specific key in namespace, empty means use all content 30 | } 31 | 32 | // apolloSubscriber is a subscriber that subscribes to Apollo config center. 33 | apolloSubscriber struct { 34 | client agollo.Client 35 | conf ApolloConf 36 | listeners []func() 37 | lock sync.RWMutex 38 | cache string 39 | cacheLock sync.RWMutex 40 | } 41 | ) 42 | 43 | var ( 44 | // ErrEmptyMetaAddr indicates that Apollo meta server address is empty. 45 | ErrEmptyMetaAddr = errors.New("empty Apollo meta server address") 46 | // ErrEmptyAppID indicates that Apollo app ID is empty. 47 | ErrEmptyAppID = errors.New("empty Apollo app ID") 48 | ) 49 | 50 | // MustNewApolloSubscriber returns an Apollo Subscriber, exits on errors. 51 | func MustNewApolloSubscriber(conf ApolloConf) subscriber.Subscriber { 52 | s, err := NewApolloSubscriber(conf) 53 | logx.Must(err) 54 | return s 55 | } 56 | 57 | // NewApolloSubscriber returns an Apollo Subscriber. 58 | func NewApolloSubscriber(conf ApolloConf) (subscriber.Subscriber, error) { 59 | if err := conf.Validate(); err != nil { 60 | return nil, err 61 | } 62 | 63 | apolloConf := buildApolloConfig(conf) 64 | client, err := agollo.StartWithConfig(func() (*config.AppConfig, error) { 65 | return apolloConf, nil 66 | }) 67 | if err != nil { 68 | return nil, err 69 | } 70 | 71 | sub := &apolloSubscriber{ 72 | client: client, 73 | conf: conf, 74 | } 75 | 76 | // Load initial value 77 | if err := sub.loadValue(); err != nil { 78 | return nil, err 79 | } 80 | 81 | // Register change listener 82 | client.AddChangeListener(sub) 83 | 84 | return sub, nil 85 | } 86 | 87 | // Validate validates the ApolloConf. 88 | func (c ApolloConf) Validate() error { 89 | if len(c.MetaAddr) == 0 { 90 | return ErrEmptyMetaAddr 91 | } 92 | if len(c.AppID) == 0 { 93 | return ErrEmptyAppID 94 | } 95 | return nil 96 | } 97 | 98 | // AddListener adds a listener to the subscriber. 99 | func (s *apolloSubscriber) AddListener(listener func()) error { 100 | s.lock.Lock() 101 | defer s.lock.Unlock() 102 | s.listeners = append(s.listeners, listener) 103 | return nil 104 | } 105 | 106 | // Value returns the value of the subscriber. 107 | func (s *apolloSubscriber) Value() (string, error) { 108 | s.cacheLock.RLock() 109 | defer s.cacheLock.RUnlock() 110 | return s.cache, nil 111 | } 112 | 113 | // OnChange is called when Apollo config changes. 114 | // Implements agollo.ChangeListener interface. 115 | func (s *apolloSubscriber) OnChange(event *storage.ChangeEvent) { 116 | s.handleConfigChange() 117 | } 118 | 119 | // OnNewestChange is called when Apollo config changes to newest. 120 | // Implements agollo.ChangeListener interface. 121 | func (s *apolloSubscriber) OnNewestChange(event *storage.FullChangeEvent) { 122 | s.handleConfigChange() 123 | } 124 | 125 | // handleConfigChange handles config reload and notifies all listeners. 126 | func (s *apolloSubscriber) handleConfigChange() { 127 | if err := s.loadValue(); err != nil { 128 | logx.Errorf("Apollo config reload failed: %v", err) 129 | return 130 | } 131 | 132 | s.lock.RLock() 133 | listeners := make([]func(), len(s.listeners)) 134 | copy(listeners, s.listeners) 135 | s.lock.RUnlock() 136 | 137 | for _, listener := range listeners { 138 | listener() 139 | } 140 | } 141 | 142 | func (s *apolloSubscriber) loadValue() error { 143 | var value string 144 | var err error 145 | 146 | // If specific key is set, get that key's value 147 | if len(s.conf.Key) > 0 { 148 | // Check if key exists in namespace 149 | cache := s.client.GetConfigCache(s.conf.NamespaceName) 150 | val, err := cache.Get(s.conf.Key) 151 | if err != nil { 152 | return errors.New("key not found in Apollo namespace") 153 | } 154 | value = toString(val) 155 | } else { 156 | // Get all content from namespace 157 | value, err = s.getAllContent() 158 | if err != nil { 159 | return err 160 | } 161 | } 162 | 163 | s.cacheLock.Lock() 164 | s.cache = value 165 | s.cacheLock.Unlock() 166 | 167 | return nil 168 | } 169 | 170 | func (s *apolloSubscriber) getAllContent() (string, error) { 171 | allConfig := make(map[string]interface{}) 172 | 173 | // Get all keys from the namespace 174 | cache := s.client.GetConfigCache(s.conf.NamespaceName) 175 | cache.Range(func(key, value interface{}) bool { 176 | if k, ok := key.(string); ok { 177 | allConfig[k] = value 178 | } 179 | return true 180 | }) 181 | 182 | var result []byte 183 | var err error 184 | 185 | switch s.conf.Format { 186 | case "json": 187 | result, err = json.Marshal(allConfig) 188 | case "yaml": 189 | result, err = yaml.Marshal(allConfig) 190 | case "properties": 191 | // For properties format, convert to key=value format 192 | props := "" 193 | for k, v := range allConfig { 194 | props += k + "=" + toString(v) + "\n" 195 | } 196 | result = []byte(props) 197 | default: 198 | result, err = json.Marshal(allConfig) 199 | } 200 | 201 | if err != nil { 202 | return "", err 203 | } 204 | 205 | return string(result), nil 206 | } 207 | 208 | func buildApolloConfig(conf ApolloConf) *config.AppConfig { 209 | apolloConf := &config.AppConfig{ 210 | AppID: conf.AppID, 211 | Cluster: conf.Cluster, 212 | NamespaceName: conf.NamespaceName, 213 | IP: conf.MetaAddr, 214 | IsBackupConfig: conf.IsBackupConfig, 215 | MustStart: conf.MustStart, 216 | } 217 | 218 | if len(conf.IP) > 0 { 219 | apolloConf.IP = conf.IP 220 | } 221 | 222 | if len(conf.Secret) > 0 { 223 | apolloConf.Secret = conf.Secret 224 | } 225 | 226 | if len(conf.BackupPath) > 0 { 227 | apolloConf.BackupConfigPath = conf.BackupPath 228 | } 229 | 230 | return apolloConf 231 | } 232 | 233 | func toString(v interface{}) string { 234 | if v == nil { 235 | return "" 236 | } 237 | if s, ok := v.(string); ok { 238 | return s 239 | } 240 | // Try to convert to string via json 241 | b, err := json.Marshal(v) 242 | if err != nil { 243 | return "" 244 | } 245 | return string(b) 246 | } 247 | -------------------------------------------------------------------------------- /configcenter/apollo/apollo_test.go: -------------------------------------------------------------------------------- 1 | package apollo 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestApolloConf_Validate(t *testing.T) { 10 | tests := []struct { 11 | name string 12 | conf ApolloConf 13 | wantErr bool 14 | }{ 15 | { 16 | name: "valid config", 17 | conf: ApolloConf{ 18 | AppID: "test-app", 19 | MetaAddr: "http://localhost:8080", 20 | }, 21 | wantErr: false, 22 | }, 23 | { 24 | name: "empty meta addr", 25 | conf: ApolloConf{ 26 | AppID: "test-app", 27 | }, 28 | wantErr: true, 29 | }, 30 | { 31 | name: "empty app id", 32 | conf: ApolloConf{ 33 | MetaAddr: "http://localhost:8080", 34 | }, 35 | wantErr: true, 36 | }, 37 | { 38 | name: "empty config", 39 | conf: ApolloConf{}, 40 | wantErr: true, 41 | }, 42 | } 43 | 44 | for _, tt := range tests { 45 | t.Run(tt.name, func(t *testing.T) { 46 | err := tt.conf.Validate() 47 | if tt.wantErr { 48 | assert.Error(t, err) 49 | } else { 50 | assert.NoError(t, err) 51 | } 52 | }) 53 | } 54 | } 55 | 56 | func TestBuildApolloConfig(t *testing.T) { 57 | conf := ApolloConf{ 58 | AppID: "test-app", 59 | Cluster: "prod", 60 | NamespaceName: "application.json", 61 | IP: "127.0.0.1", 62 | MetaAddr: "http://localhost:8080", 63 | Secret: "test-secret", 64 | IsBackupConfig: true, 65 | BackupPath: "/tmp/apollo", 66 | MustStart: false, 67 | } 68 | 69 | apolloConf := buildApolloConfig(conf) 70 | 71 | assert.Equal(t, "test-app", apolloConf.AppID) 72 | assert.Equal(t, "prod", apolloConf.Cluster) 73 | assert.Equal(t, "application.json", apolloConf.NamespaceName) 74 | assert.Equal(t, "127.0.0.1", apolloConf.IP) 75 | assert.Equal(t, "test-secret", apolloConf.Secret) 76 | assert.Equal(t, true, apolloConf.IsBackupConfig) 77 | assert.Equal(t, "/tmp/apollo", apolloConf.BackupConfigPath) 78 | assert.Equal(t, false, apolloConf.MustStart) 79 | } 80 | 81 | func TestBuildApolloConfig_Defaults(t *testing.T) { 82 | conf := ApolloConf{ 83 | AppID: "test-app", 84 | MetaAddr: "http://localhost:8080", 85 | } 86 | 87 | apolloConf := buildApolloConfig(conf) 88 | 89 | assert.Equal(t, "test-app", apolloConf.AppID) 90 | assert.Equal(t, "", apolloConf.Cluster) // Will use default in actual usage 91 | assert.Equal(t, "", apolloConf.NamespaceName) // Will use default in actual usage 92 | assert.Equal(t, "http://localhost:8080", apolloConf.IP) 93 | } 94 | 95 | func TestToString(t *testing.T) { 96 | tests := []struct { 97 | name string 98 | input interface{} 99 | expected string 100 | }{ 101 | { 102 | name: "nil value", 103 | input: nil, 104 | expected: "", 105 | }, 106 | { 107 | name: "string value", 108 | input: "test", 109 | expected: "test", 110 | }, 111 | { 112 | name: "int value", 113 | input: 123, 114 | expected: "123", 115 | }, 116 | { 117 | name: "bool value", 118 | input: true, 119 | expected: "true", 120 | }, 121 | { 122 | name: "map value", 123 | input: map[string]string{"key": "value"}, 124 | expected: `{"key":"value"}`, 125 | }, 126 | { 127 | name: "empty string value", 128 | input: "", 129 | expected: "", 130 | }, 131 | { 132 | name: "float value", 133 | input: 3.14, 134 | expected: "3.14", 135 | }, 136 | } 137 | 138 | for _, tt := range tests { 139 | t.Run(tt.name, func(t *testing.T) { 140 | result := toString(tt.input) 141 | assert.Equal(t, tt.expected, result) 142 | }) 143 | } 144 | } 145 | 146 | func TestBuildApolloConfig_IPOverride(t *testing.T) { 147 | tests := []struct { 148 | name string 149 | conf ApolloConf 150 | expected string 151 | }{ 152 | { 153 | name: "IP overrides MetaAddr with bare IP", 154 | conf: ApolloConf{ 155 | AppID: "test-app", 156 | MetaAddr: "http://localhost:8080", 157 | IP: "192.168.1.100", 158 | }, 159 | expected: "192.168.1.100", 160 | }, 161 | { 162 | name: "IP overrides MetaAddr with HTTP URL", 163 | conf: ApolloConf{ 164 | AppID: "test-app", 165 | MetaAddr: "http://localhost:8080", 166 | IP: "http://config.example.com:8080", 167 | }, 168 | expected: "http://config.example.com:8080", 169 | }, 170 | { 171 | name: "IP overrides MetaAddr with HTTPS URL", 172 | conf: ApolloConf{ 173 | AppID: "test-app", 174 | MetaAddr: "http://localhost:8080", 175 | IP: "https://secure.example.com", 176 | }, 177 | expected: "https://secure.example.com", 178 | }, 179 | { 180 | name: "IP overrides MetaAddr with URL including path", 181 | conf: ApolloConf{ 182 | AppID: "test-app", 183 | MetaAddr: "http://localhost:8080", 184 | IP: "http://config.example.com:8080/config-service", 185 | }, 186 | expected: "http://config.example.com:8080/config-service", 187 | }, 188 | } 189 | 190 | for _, tt := range tests { 191 | t.Run(tt.name, func(t *testing.T) { 192 | apolloConf := buildApolloConfig(tt.conf) 193 | assert.Equal(t, tt.expected, apolloConf.IP) 194 | }) 195 | } 196 | } 197 | 198 | func TestBuildApolloConfig_MetaAddrAsDefault(t *testing.T) { 199 | // Test that MetaAddr is used when IP is not provided 200 | conf := ApolloConf{ 201 | AppID: "test-app", 202 | MetaAddr: "http://localhost:8080", 203 | } 204 | 205 | apolloConf := buildApolloConfig(conf) 206 | 207 | // MetaAddr should be used as default for IP 208 | assert.Equal(t, "http://localhost:8080", apolloConf.IP) 209 | } 210 | -------------------------------------------------------------------------------- /configcenter/apollo/examples/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "time" 6 | 7 | configurator "github.com/zeromicro/go-zero/core/configcenter" 8 | "github.com/zeromicro/go-zero/core/logx" 9 | "github.com/zeromicro/zero-contrib/configcenter/apollo" 10 | ) 11 | 12 | type AppConfig struct { 13 | Name string `json:"name"` 14 | Version string `json:"version"` 15 | Timeout int64 `json:"timeout"` 16 | MaxConns int `json:"maxConns"` 17 | Features struct { 18 | EnableCache bool `json:"enableCache"` 19 | EnableTrace bool `json:"enableTrace"` 20 | } `json:"features"` 21 | } 22 | 23 | func main() { 24 | // Create Apollo subscriber 25 | sub := apollo.MustNewApolloSubscriber(apollo.ApolloConf{ 26 | AppID: "go-zero-demo", // Your Apollo AppID 27 | Cluster: "default", // Cluster name 28 | NamespaceName: "application.json", // Namespace name with format suffix 29 | MetaAddr: "http://localhost:8080", // Apollo meta server address 30 | Secret: "", // Optional: Apollo secret key 31 | IsBackupConfig: true, // Enable backup to local file 32 | BackupPath: "/tmp/apollo-backup", // Backup directory 33 | Format: "json", // Config format: json, yaml, or properties 34 | }) 35 | 36 | // Create config center with type-safe config 37 | cc := configurator.MustNewConfigCenter[AppConfig](configurator.Config{ 38 | Type: "json", 39 | Log: true, 40 | }, sub) 41 | 42 | // Get initial config 43 | config, err := cc.GetConfig() 44 | if err != nil { 45 | logx.Errorf("Failed to get config: %v", err) 46 | return 47 | } 48 | 49 | fmt.Printf("Initial config: %+v\n", config) 50 | 51 | // Add listener for config changes (hot reload) 52 | cc.AddListener(func() { 53 | newConfig, err := cc.GetConfig() 54 | if err != nil { 55 | logx.Errorf("Failed to get updated config: %v", err) 56 | return 57 | } 58 | fmt.Printf("Config updated: %+v\n", newConfig) 59 | 60 | // Apply new configuration 61 | // For example, update connection pool size, feature flags, etc. 62 | logx.Infof("Applying new config - MaxConns: %d, Timeout: %d", 63 | newConfig.MaxConns, newConfig.Timeout) 64 | }) 65 | 66 | // Keep running to receive config updates 67 | fmt.Println("Listening for config changes... Press Ctrl+C to exit") 68 | select {} 69 | } 70 | 71 | // Example: Using specific key instead of entire namespace 72 | func exampleSpecificKey() { 73 | // Get specific key from Apollo 74 | sub := apollo.MustNewApolloSubscriber(apollo.ApolloConf{ 75 | AppID: "go-zero-demo", 76 | Cluster: "default", 77 | NamespaceName: "application", 78 | MetaAddr: "http://localhost:8080", 79 | Key: "database.url", // Specific key to watch 80 | }) 81 | 82 | cc := configurator.MustNewConfigCenter[string](configurator.Config{ 83 | Type: "json", 84 | }, sub) 85 | 86 | dbUrl, err := cc.GetConfig() 87 | if err != nil { 88 | logx.Errorf("Failed to get database URL: %v", err) 89 | return 90 | } 91 | fmt.Printf("Database URL: %s\n", dbUrl) 92 | } 93 | 94 | // Example: Using properties format 95 | func examplePropertiesFormat() { 96 | sub := apollo.MustNewApolloSubscriber(apollo.ApolloConf{ 97 | AppID: "go-zero-demo", 98 | Cluster: "default", 99 | NamespaceName: "application.properties", 100 | MetaAddr: "http://localhost:8080", 101 | Format: "properties", // Properties format 102 | }) 103 | 104 | cc := configurator.MustNewConfigCenter[string](configurator.Config{ 105 | Type: "json", 106 | }, sub) 107 | 108 | properties, err := cc.GetConfig() 109 | if err != nil { 110 | logx.Errorf("Failed to get properties: %v", err) 111 | return 112 | } 113 | fmt.Printf("Properties: %s\n", properties) 114 | } 115 | 116 | // Example: Integration with existing go-zero service 117 | type ServiceConfig struct { 118 | configurator.Configurator[AppConfig] 119 | } 120 | 121 | func (sc *ServiceConfig) UpdateConfig(newConfig AppConfig) { 122 | // Update service behavior based on new config 123 | logx.Infof("Service config updated: %+v", newConfig) 124 | // Update connection pools, timeouts, feature flags, etc. 125 | } 126 | 127 | func exampleServiceIntegration() { 128 | sub := apollo.MustNewApolloSubscriber(apollo.ApolloConf{ 129 | AppID: "my-service", 130 | MetaAddr: "http://localhost:8080", 131 | }) 132 | 133 | cc := configurator.MustNewConfigCenter[AppConfig]( 134 | configurator.Config{Type: "json"}, 135 | sub, 136 | ) 137 | 138 | svc := &ServiceConfig{ 139 | Configurator: cc, 140 | } 141 | 142 | // Listen for changes 143 | cc.AddListener(func() { 144 | config, err := cc.GetConfig() 145 | if err != nil { 146 | logx.Errorf("Failed to get updated config: %v", err) 147 | return 148 | } 149 | svc.UpdateConfig(config) 150 | }) 151 | 152 | // Service logic continues... 153 | time.Sleep(time.Hour) 154 | } 155 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib 2 | 3 | go 1.23 4 | 5 | require ( 6 | github.com/apolloconfig/agollo/v4 v4.4.0 7 | github.com/stretchr/testify v1.11.1 8 | github.com/zeromicro/go-zero v1.9.2 9 | gopkg.in/yaml.v2 v2.4.0 10 | ) 11 | 12 | require ( 13 | github.com/coreos/go-semver v0.3.1 // indirect 14 | github.com/coreos/go-systemd/v22 v22.5.0 // indirect 15 | github.com/davecgh/go-spew v1.1.1 // indirect 16 | github.com/fatih/color v1.18.0 // indirect 17 | github.com/fsnotify/fsnotify v1.4.9 // indirect 18 | github.com/gogo/protobuf v1.3.2 // indirect 19 | github.com/golang/protobuf v1.5.4 // indirect 20 | github.com/hashicorp/hcl v1.0.0 // indirect 21 | github.com/magiconair/properties v1.8.5 // indirect 22 | github.com/mattn/go-colorable v0.1.13 // indirect 23 | github.com/mattn/go-isatty v0.0.20 // indirect 24 | github.com/mitchellh/mapstructure v1.4.1 // indirect 25 | github.com/pelletier/go-toml v1.9.3 // indirect 26 | github.com/pelletier/go-toml/v2 v2.2.2 // indirect 27 | github.com/pmezard/go-difflib v1.0.0 // indirect 28 | github.com/spaolacci/murmur3 v1.1.0 // indirect 29 | github.com/spf13/afero v1.6.0 // indirect 30 | github.com/spf13/cast v1.3.1 // indirect 31 | github.com/spf13/jwalterweatherman v1.1.0 // indirect 32 | github.com/spf13/pflag v1.0.5 // indirect 33 | github.com/spf13/viper v1.8.1 // indirect 34 | github.com/subosito/gotenv v1.2.0 // indirect 35 | go.etcd.io/etcd/api/v3 v3.5.15 // indirect 36 | go.etcd.io/etcd/client/pkg/v3 v3.5.15 // indirect 37 | go.etcd.io/etcd/client/v3 v3.5.15 // indirect 38 | go.opentelemetry.io/otel v1.24.0 // indirect 39 | go.opentelemetry.io/otel/trace v1.24.0 // indirect 40 | go.uber.org/atomic v1.10.0 // indirect 41 | go.uber.org/automaxprocs v1.6.0 // indirect 42 | go.uber.org/mock v0.4.0 // indirect 43 | go.uber.org/multierr v1.9.0 // indirect 44 | go.uber.org/zap v1.24.0 // indirect 45 | golang.org/x/net v0.35.0 // indirect 46 | golang.org/x/sys v0.30.0 // indirect 47 | golang.org/x/text v0.22.0 // indirect 48 | google.golang.org/genproto/googleapis/api v0.0.0-20240711142825-46eb208f015d // indirect 49 | google.golang.org/genproto/googleapis/rpc v0.0.0-20240701130421-f6361c86f094 // indirect 50 | google.golang.org/grpc v1.65.0 // indirect 51 | google.golang.org/protobuf v1.36.5 // indirect 52 | gopkg.in/ini.v1 v1.62.0 // indirect 53 | gopkg.in/yaml.v3 v3.0.1 // indirect 54 | ) 55 | -------------------------------------------------------------------------------- /handler/README.md: -------------------------------------------------------------------------------- 1 | ## Zero Handler Contrib List 2 | 3 | ETag: ETag Handler to support [ETag](https://en.wikipedia.org/wiki/HTTP_ETag) both for weak and strong validation 4 | 5 | ## ETag 6 | 7 | ### Prerequisites: 8 | 9 | * Install `go-zero`: go get -u github.com/zeromicro/go-zero 10 | 11 | ### Download the module: 12 | 13 | ```shell 14 | go get -u github.com/zeromicro/zero-contrib/handler 15 | ``` 16 | 17 | ### Example: 18 | 19 | ```go 20 | package api 21 | 22 | import ( 23 | ... 24 | 25 | "github.com/zeromicro/go-zero/rest" 26 | "github.com/zeromicro/zero-contrib/handler" 27 | ) 28 | 29 | func main() { 30 | ... 31 | 32 | server := rest.MustNewServer(c.RestConf) 33 | server.Use(handler.NewETagMiddleware(true).Handle) 34 | } 35 | 36 | ``` 37 | 38 | ### Gratitude: 39 | 40 | - https://github.com/go-http-utils/etag -------------------------------------------------------------------------------- /handler/etag.go: -------------------------------------------------------------------------------- 1 | package handler 2 | 3 | import ( 4 | "bytes" 5 | "crypto/sha1" 6 | "encoding/hex" 7 | "fmt" 8 | "hash" 9 | "net/http" 10 | "strconv" 11 | "strings" 12 | "time" 13 | ) 14 | 15 | type ETagMiddleware struct { 16 | weak bool 17 | } 18 | 19 | func NewETagMiddleware(weak bool) *ETagMiddleware { 20 | return &ETagMiddleware{weak: weak} 21 | } 22 | 23 | func (m *ETagMiddleware) Handle(h http.HandlerFunc) http.HandlerFunc { 24 | return func(res http.ResponseWriter, req *http.Request) { 25 | hw := hashWriter{rw: res, hash: sha1.New(), buf: bytes.NewBuffer(nil)} 26 | h.ServeHTTP(&hw, req) 27 | 28 | resHeader := res.Header() 29 | 30 | if hw.hash == nil || 31 | resHeader.Get(HeaderETag) != "" || 32 | strconv.Itoa(hw.status)[0] != '2' || 33 | hw.status == http.StatusNoContent || 34 | hw.buf.Len() == 0 { 35 | writeRaw(res, hw) 36 | return 37 | } 38 | 39 | etag := fmt.Sprintf("%v-%v", strconv.Itoa(hw.len), 40 | hex.EncodeToString(hw.hash.Sum(nil))) 41 | 42 | if m.weak { 43 | etag = "W/" + etag 44 | } 45 | 46 | resHeader.Set(HeaderETag, etag) 47 | 48 | if IsFresh(req.Header, resHeader) { 49 | res.WriteHeader(http.StatusNotModified) 50 | res.Write(nil) 51 | } else { 52 | writeRaw(res, hw) 53 | } 54 | } 55 | } 56 | 57 | type hashWriter struct { 58 | rw http.ResponseWriter 59 | hash hash.Hash 60 | buf *bytes.Buffer 61 | len int 62 | status int 63 | } 64 | 65 | func (hw hashWriter) Header() http.Header { 66 | return hw.rw.Header() 67 | } 68 | 69 | func (hw *hashWriter) WriteHeader(status int) { 70 | hw.status = status 71 | } 72 | 73 | func (hw *hashWriter) Write(b []byte) (int, error) { 74 | if hw.status == 0 { 75 | hw.status = http.StatusOK 76 | } 77 | // bytes.Buffer.Write(b) always return (len(b), nil), so just 78 | // ignore the return values. 79 | hw.buf.Write(b) 80 | 81 | l, err := hw.hash.Write(b) 82 | hw.len += l 83 | return l, err 84 | } 85 | 86 | func writeRaw(res http.ResponseWriter, hw hashWriter) { 87 | res.WriteHeader(hw.status) 88 | res.Write(hw.buf.Bytes()) 89 | } 90 | 91 | // IsFresh check whether cache can be used in this HTTP request 92 | func IsFresh(reqHeader http.Header, resHeader http.Header) bool { 93 | isEtagMatched, isModifiedMatched := false, false 94 | 95 | ifModifiedSince := reqHeader.Get(HeaderIfModifiedSince) 96 | ifUnmodifiedSince := reqHeader.Get(HeaderIfUnmodifiedSince) 97 | ifNoneMatch := reqHeader.Get(HeaderIfNoneMatch) 98 | ifMatch := reqHeader.Get(HeaderIfMatch) 99 | cacheControl := reqHeader.Get(HeaderCacheControl) 100 | 101 | etag := resHeader.Get(HeaderETag) 102 | lastModified := resHeader.Get(HeaderLastModified) 103 | 104 | if ifModifiedSince == "" && 105 | ifUnmodifiedSince == "" && 106 | ifNoneMatch == "" && 107 | ifMatch == "" { 108 | return false 109 | } 110 | 111 | if strings.Contains(cacheControl, "no-cache") { 112 | return false 113 | } 114 | 115 | if etag != "" && ifNoneMatch != "" { 116 | isEtagMatched = checkEtagNoneMatch(trimTags(strings.Split(ifNoneMatch, ",")), etag) 117 | } 118 | 119 | if etag != "" && ifMatch != "" && !isEtagMatched { 120 | isEtagMatched = checkEtagMatch(trimTags(strings.Split(ifMatch, ",")), etag) 121 | } 122 | 123 | if lastModified != "" && ifModifiedSince != "" { 124 | isModifiedMatched = checkModifedMatch(lastModified, ifModifiedSince) 125 | } 126 | 127 | if lastModified != "" && ifUnmodifiedSince != "" && !isModifiedMatched { 128 | isModifiedMatched = checkUnmodifedMatch(lastModified, ifUnmodifiedSince) 129 | } 130 | 131 | return isEtagMatched || isModifiedMatched 132 | } 133 | 134 | func trimTags(tags []string) []string { 135 | trimedTags := make([]string, len(tags)) 136 | 137 | for i, tag := range tags { 138 | trimedTags[i] = strings.TrimSpace(tag) 139 | } 140 | 141 | return trimedTags 142 | } 143 | 144 | func checkEtagNoneMatch(etagsToNoneMatch []string, etag string) bool { 145 | for _, etagToNoneMatch := range etagsToNoneMatch { 146 | if etagToNoneMatch == "*" || etagToNoneMatch == etag || etagToNoneMatch == "W/"+etag { 147 | return true 148 | } 149 | } 150 | 151 | return false 152 | } 153 | 154 | func checkEtagMatch(etagsToMatch []string, etag string) bool { 155 | for _, etagToMatch := range etagsToMatch { 156 | if etagToMatch == "*" { 157 | return false 158 | } 159 | 160 | if strings.HasPrefix(etagToMatch, "W/") { 161 | if etagToMatch == "W/"+etag { 162 | return false 163 | } 164 | } else { 165 | if etagToMatch == etag { 166 | return false 167 | } 168 | } 169 | } 170 | 171 | return true 172 | } 173 | 174 | func checkModifedMatch(lastModified, ifModifiedSince string) bool { 175 | if lm, ims, ok := parseTimePairs(lastModified, ifModifiedSince); ok { 176 | return lm.Before(ims) 177 | } 178 | 179 | return false 180 | } 181 | 182 | func checkUnmodifedMatch(lastModified, ifUnmodifiedSince string) bool { 183 | if lm, ius, ok := parseTimePairs(lastModified, ifUnmodifiedSince); ok { 184 | return lm.After(ius) 185 | } 186 | 187 | return false 188 | } 189 | 190 | func parseTimePairs(s1, s2 string) (t1 time.Time, t2 time.Time, ok bool) { 191 | if t1, err := time.Parse(http.TimeFormat, s1); err == nil { 192 | if t2, err := time.Parse(http.TimeFormat, s2); err == nil { 193 | return t1, t2, true 194 | } 195 | } 196 | 197 | return t1, t2, false 198 | } 199 | -------------------------------------------------------------------------------- /handler/etag_test.go: -------------------------------------------------------------------------------- 1 | package handler 2 | 3 | import ( 4 | "crypto/sha1" 5 | "encoding/hex" 6 | "fmt" 7 | "net/http" 8 | "net/http/httptest" 9 | "testing" 10 | "time" 11 | 12 | "github.com/stretchr/testify/suite" 13 | ) 14 | 15 | var testStrBytes = []byte("Hello World") 16 | var testStrEtag = "11-0a4d55a8d778e5022fab701977c5d840bbc486d0" 17 | 18 | type EmptyEtagSuite struct { 19 | suite.Suite 20 | 21 | server *httptest.Server 22 | } 23 | 24 | func (s *EmptyEtagSuite) SetupTest() { 25 | mux := http.NewServeMux() 26 | mux.Handle("/", NewETagMiddleware(true).Handle(emptyHandlerFunc)) 27 | 28 | s.server = httptest.NewServer(mux) 29 | } 30 | 31 | func (s *EmptyEtagSuite) TestNoEtag() { 32 | res, err := http.Get(s.server.URL + "/") 33 | 34 | s.Nil(err) 35 | s.Equal(http.StatusNoContent, res.StatusCode) 36 | s.Empty(res.Header.Get(HeaderETag)) 37 | } 38 | 39 | func TestEmptyEtag(t *testing.T) { 40 | suite.Run(t, new(EmptyEtagSuite)) 41 | } 42 | 43 | type EtagSuite struct { 44 | suite.Suite 45 | 46 | server *httptest.Server 47 | weakServer *httptest.Server 48 | } 49 | 50 | func (s *EtagSuite) SetupTest() { 51 | mux := http.NewServeMux() 52 | mux.Handle("/", NewETagMiddleware(false).Handle(handlerFunc)) 53 | 54 | s.server = httptest.NewServer(mux) 55 | 56 | wmux := http.NewServeMux() 57 | wmux.Handle("/", NewETagMiddleware(true).Handle(handlerFunc)) 58 | 59 | s.weakServer = httptest.NewServer(wmux) 60 | } 61 | 62 | func (s EtagSuite) TestEtagExists() { 63 | res, err := http.Get(s.server.URL + "/") 64 | 65 | s.Nil(err) 66 | s.Equal(http.StatusOK, res.StatusCode) 67 | 68 | h := sha1.New() 69 | h.Write(testStrBytes) 70 | 71 | s.Equal(fmt.Sprintf("%v-%v", len(testStrBytes), hex.EncodeToString(h.Sum(nil))), res.Header.Get(HeaderETag)) 72 | } 73 | 74 | func (s EtagSuite) TestWeakEtagExists() { 75 | res, err := http.Get(s.weakServer.URL + "/") 76 | 77 | s.Nil(err) 78 | s.Equal(http.StatusOK, res.StatusCode) 79 | 80 | h := sha1.New() 81 | h.Write(testStrBytes) 82 | 83 | s.Equal(fmt.Sprintf("W/%v-%v", len(testStrBytes), hex.EncodeToString(h.Sum(nil))), res.Header.Get(HeaderETag)) 84 | } 85 | 86 | func (s EtagSuite) TestMatch() { 87 | req, err := http.NewRequest(http.MethodGet, s.server.URL+"/", nil) 88 | s.Nil(err) 89 | 90 | req.Header.Set(HeaderIfNoneMatch, testStrEtag) 91 | 92 | cli := &http.Client{} 93 | res, err := cli.Do(req) 94 | 95 | s.Nil(err) 96 | s.Equal(http.StatusNotModified, res.StatusCode) 97 | } 98 | 99 | func TestEtag(t *testing.T) { 100 | suite.Run(t, new(EtagSuite)) 101 | } 102 | 103 | func emptyHandlerFunc(res http.ResponseWriter, req *http.Request) { 104 | res.WriteHeader(http.StatusNoContent) 105 | 106 | res.Write(nil) 107 | } 108 | 109 | func handlerFunc(res http.ResponseWriter, req *http.Request) { 110 | res.WriteHeader(http.StatusOK) 111 | 112 | res.Write(testStrBytes) 113 | } 114 | 115 | type FreshSuite struct { 116 | suite.Suite 117 | 118 | reqHeader http.Header 119 | resHeader http.Header 120 | } 121 | 122 | func (s *FreshSuite) SetupTest() { 123 | s.reqHeader = make(http.Header) 124 | s.resHeader = make(http.Header) 125 | } 126 | 127 | func (s FreshSuite) TestNoCache() { 128 | s.reqHeader.Set(HeaderCacheControl, "no-cache") 129 | s.reqHeader.Set(HeaderIfNoneMatch, "foo") 130 | 131 | s.False(IsFresh(s.reqHeader, s.resHeader)) 132 | } 133 | 134 | func (s FreshSuite) TestEtagEmpty() { 135 | s.False(IsFresh(s.reqHeader, s.resHeader)) 136 | } 137 | 138 | func (s FreshSuite) TestEtagMatch() { 139 | s.reqHeader.Set(HeaderIfNoneMatch, "foo") 140 | s.resHeader.Set(HeaderETag, "foo") 141 | 142 | s.True(IsFresh(s.reqHeader, s.resHeader)) 143 | } 144 | 145 | func (s FreshSuite) TestEtagMismatch() { 146 | s.reqHeader.Set(HeaderIfNoneMatch, "foo") 147 | s.resHeader.Set(HeaderETag, "bar") 148 | 149 | s.False(IsFresh(s.reqHeader, s.resHeader)) 150 | } 151 | 152 | func (s FreshSuite) TestEtagMissing() { 153 | s.reqHeader.Set(HeaderIfNoneMatch, "foo") 154 | 155 | s.False(IsFresh(s.reqHeader, s.resHeader)) 156 | } 157 | 158 | func (s FreshSuite) TestWeakEtagMatch() { 159 | s.reqHeader.Set(HeaderIfNoneMatch, `W/"foo"`) 160 | s.resHeader.Set(HeaderETag, `W/"foo"`) 161 | 162 | s.True(IsFresh(s.reqHeader, s.resHeader)) 163 | } 164 | 165 | func (s FreshSuite) TestEtagStrongMatch() { 166 | s.reqHeader.Set(HeaderIfNoneMatch, `W/"foo"`) 167 | s.resHeader.Set(HeaderETag, `"foo"`) 168 | 169 | s.True(IsFresh(s.reqHeader, s.resHeader)) 170 | } 171 | 172 | func (s FreshSuite) TestEtagIfMatch() { 173 | s.reqHeader.Set(HeaderIfMatch, "foo") 174 | s.resHeader.Set(HeaderETag, "bar") 175 | 176 | s.True(IsFresh(s.reqHeader, s.resHeader)) 177 | } 178 | 179 | func (s FreshSuite) TestWeakEtagIfMatch() { 180 | s.reqHeader.Set(HeaderIfMatch, "W/foo") 181 | s.resHeader.Set(HeaderETag, "W/bar") 182 | 183 | s.True(IsFresh(s.reqHeader, s.resHeader)) 184 | } 185 | 186 | func (s FreshSuite) TestStarEtagIfMatch() { 187 | s.reqHeader.Set(HeaderIfMatch, "*") 188 | s.resHeader.Set(HeaderETag, "W/bar") 189 | 190 | s.False(IsFresh(s.reqHeader, s.resHeader)) 191 | } 192 | 193 | func (s FreshSuite) TestWeakEtagIfMatchMatched() { 194 | s.reqHeader.Set(HeaderIfMatch, "W/bar") 195 | s.resHeader.Set(HeaderETag, "bar") 196 | 197 | s.False(IsFresh(s.reqHeader, s.resHeader)) 198 | } 199 | 200 | func (s FreshSuite) TestEtagIfMatchMatched() { 201 | s.reqHeader.Set(HeaderIfMatch, "bar") 202 | s.resHeader.Set(HeaderETag, "bar") 203 | 204 | s.False(IsFresh(s.reqHeader, s.resHeader)) 205 | } 206 | 207 | func (s FreshSuite) TestStaleOnEtagWeakMatch() { 208 | s.reqHeader.Set(HeaderIfNoneMatch, `"foo"`) 209 | s.resHeader.Set(HeaderETag, `W/"foo"`) 210 | 211 | s.False(IsFresh(s.reqHeader, s.resHeader)) 212 | } 213 | 214 | func (s FreshSuite) TestEtagAsterisk() { 215 | s.reqHeader.Set(HeaderIfNoneMatch, "*") 216 | s.resHeader.Set(HeaderETag, `"foo"`) 217 | 218 | s.True(IsFresh(s.reqHeader, s.resHeader)) 219 | } 220 | 221 | func (s FreshSuite) TestModifiedFresh() { 222 | s.reqHeader.Set(HeaderIfModifiedSince, getFormattedTime(4*time.Second)) 223 | s.resHeader.Set(HeaderLastModified, getFormattedTime(2*time.Second)) 224 | 225 | s.True(IsFresh(s.reqHeader, s.resHeader)) 226 | } 227 | 228 | func (s FreshSuite) TestModifiedStale() { 229 | s.reqHeader.Set(HeaderIfModifiedSince, getFormattedTime(2*time.Second)) 230 | s.resHeader.Set(HeaderLastModified, getFormattedTime(4*time.Second)) 231 | 232 | s.False(IsFresh(s.reqHeader, s.resHeader)) 233 | } 234 | 235 | func (s FreshSuite) TestUnmodifiedFresh() { 236 | s.reqHeader.Set(HeaderIfUnmodifiedSince, getFormattedTime(2*time.Second)) 237 | s.resHeader.Set(HeaderLastModified, getFormattedTime(4*time.Second)) 238 | 239 | s.True(IsFresh(s.reqHeader, s.resHeader)) 240 | } 241 | 242 | func (s FreshSuite) TestUnmodifiedStale() { 243 | s.reqHeader.Set(HeaderIfUnmodifiedSince, getFormattedTime(4*time.Second)) 244 | s.resHeader.Set(HeaderLastModified, getFormattedTime(2*time.Second)) 245 | 246 | s.False(IsFresh(s.reqHeader, s.resHeader)) 247 | } 248 | 249 | func (s FreshSuite) TestEmptyLastModified() { 250 | s.reqHeader.Set(HeaderIfModifiedSince, getFormattedTime(4*time.Second)) 251 | 252 | s.False(IsFresh(s.reqHeader, s.resHeader)) 253 | } 254 | 255 | func (s FreshSuite) TestBoshAndModifiedFresh() { 256 | s.reqHeader.Set(HeaderIfNoneMatch, "foo") 257 | s.reqHeader.Set(HeaderIfModifiedSince, getFormattedTime(4*time.Second)) 258 | 259 | s.resHeader.Set(HeaderETag, "bar") 260 | s.resHeader.Set(HeaderLastModified, getFormattedTime(2*time.Second)) 261 | 262 | s.True(IsFresh(s.reqHeader, s.resHeader)) 263 | } 264 | 265 | func (s FreshSuite) TestBoshAndETagFresh() { 266 | s.reqHeader.Set(HeaderIfNoneMatch, "foo") 267 | s.reqHeader.Set(HeaderIfModifiedSince, getFormattedTime(2*time.Second)) 268 | 269 | s.resHeader.Set(HeaderETag, "foo") 270 | s.resHeader.Set(HeaderLastModified, getFormattedTime(4*time.Second)) 271 | 272 | s.True(IsFresh(s.reqHeader, s.resHeader)) 273 | } 274 | 275 | func (s FreshSuite) TestBoshFresh() { 276 | s.reqHeader.Set(HeaderIfNoneMatch, "foo") 277 | s.reqHeader.Set(HeaderIfModifiedSince, getFormattedTime(4*time.Second)) 278 | 279 | s.resHeader.Set(HeaderETag, "foo") 280 | s.resHeader.Set(HeaderLastModified, getFormattedTime(2*time.Second)) 281 | 282 | s.True(IsFresh(s.reqHeader, s.resHeader)) 283 | } 284 | 285 | func (s FreshSuite) TestBoshStale() { 286 | s.reqHeader.Set(HeaderIfNoneMatch, "foo") 287 | s.reqHeader.Set(HeaderIfModifiedSince, getFormattedTime(2*time.Second)) 288 | 289 | s.resHeader.Set(HeaderETag, "bar") 290 | s.resHeader.Set(HeaderLastModified, getFormattedTime(4*time.Second)) 291 | 292 | s.False(IsFresh(s.reqHeader, s.resHeader)) 293 | } 294 | 295 | func TestFresh(t *testing.T) { 296 | suite.Run(t, new(FreshSuite)) 297 | } 298 | 299 | func getFormattedTime(d time.Duration) string { 300 | return time.Now().Add(d).Format(http.TimeFormat) 301 | } 302 | -------------------------------------------------------------------------------- /handler/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/handler 2 | 3 | go 1.16 4 | 5 | require github.com/stretchr/testify v1.8.4 6 | -------------------------------------------------------------------------------- /handler/go.sum: -------------------------------------------------------------------------------- 1 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 2 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 3 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 4 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 5 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 6 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= 7 | github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= 8 | github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= 9 | github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= 10 | github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= 11 | github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= 12 | github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= 13 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= 14 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 15 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 16 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= 17 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 18 | -------------------------------------------------------------------------------- /handler/header_types.go: -------------------------------------------------------------------------------- 1 | package handler 2 | 3 | import "net/textproto" 4 | 5 | // HTTP headers 6 | const ( 7 | HeaderAccept = "Accept" 8 | HeaderAcceptCharset = "Accept-Charset" 9 | HeaderAcceptEncoding = "Accept-Encoding" 10 | HeaderAcceptLanguage = "Accept-Language" 11 | HeaderAuthorization = "Authorization" 12 | HeaderCacheControl = "Cache-Control" 13 | HeaderContentLength = "Content-Length" 14 | HeaderContentMD5 = "Content-MD5" 15 | HeaderContentType = "Content-Type" 16 | HeaderDoNotTrack = "DNT" 17 | HeaderIfMatch = "If-Match" 18 | HeaderIfModifiedSince = "If-Modified-Since" 19 | HeaderIfNoneMatch = "If-None-Match" 20 | HeaderIfRange = "If-Range" 21 | HeaderIfUnmodifiedSince = "If-Unmodified-Since" 22 | HeaderMaxForwards = "Max-Forwards" 23 | HeaderProxyAuthorization = "Proxy-Authorization" 24 | HeaderPragma = "Pragma" 25 | HeaderRange = "Range" 26 | HeaderReferer = "Referer" 27 | HeaderUserAgent = "User-Agent" 28 | HeaderTE = "TE" 29 | HeaderVia = "Via" 30 | HeaderWarning = "Warning" 31 | HeaderCookie = "Cookie" 32 | HeaderOrigin = "Origin" 33 | HeaderAcceptDatetime = "Accept-Datetime" 34 | HeaderXRequestedWith = "X-Requested-With" 35 | HeaderAccessControlAllowOrigin = "Access-Control-Allow-Origin" 36 | HeaderAccessControlAllowMethods = "Access-Control-Allow-Methods" 37 | HeaderAccessControlAllowHeaders = "Access-Control-Allow-Headers" 38 | HeaderAccessControlAllowCredentials = "Access-Control-Allow-Credentials" 39 | HeaderAccessControlExposeHeaders = "Access-Control-Expose-Headers" 40 | HeaderAccessControlMaxAge = "Access-Control-Max-Age" 41 | HeaderAccessControlRequestMethod = "Access-Control-Request-Method" 42 | HeaderAccessControlRequestHeaders = "Access-Control-Request-Headers" 43 | HeaderAcceptPatch = "Accept-Patch" 44 | HeaderAcceptRanges = "Accept-Ranges" 45 | HeaderAllow = "Allow" 46 | HeaderContentEncoding = "Content-Encoding" 47 | HeaderContentLanguage = "Content-Language" 48 | HeaderContentLocation = "Content-Location" 49 | HeaderContentDisposition = "Content-Disposition" 50 | HeaderContentRange = "Content-Range" 51 | HeaderETag = "ETag" 52 | HeaderExpires = "Expires" 53 | HeaderLastModified = "Last-Modified" 54 | HeaderLink = "Link" 55 | HeaderLocation = "Location" 56 | HeaderP3P = "P3P" 57 | HeaderProxyAuthenticate = "Proxy-Authenticate" 58 | HeaderRefresh = "Refresh" 59 | HeaderRetryAfter = "Retry-After" 60 | HeaderServer = "Server" 61 | HeaderSetCookie = "Set-Cookie" 62 | HeaderStrictTransportSecurity = "Strict-Transport-Security" 63 | HeaderTransferEncoding = "Transfer-Encoding" 64 | HeaderUpgrade = "Upgrade" 65 | HeaderVary = "Vary" 66 | HeaderWWWAuthenticate = "WWW-Authenticate" 67 | 68 | // Non-Standard 69 | HeaderXFrameOptions = "X-Frame-Options" 70 | HeaderXXSSProtection = "X-XSS-Protection" 71 | HeaderContentSecurityPolicy = "Content-Security-Policy" 72 | HeaderXContentSecurityPolicy = "X-Content-Security-Policy" 73 | HeaderXWebKitCSP = "X-WebKit-CSP" 74 | HeaderXContentTypeOptions = "X-Content-Type-Options" 75 | HeaderXPoweredBy = "X-Powered-By" 76 | HeaderXUACompatible = "X-UA-Compatible" 77 | HeaderXForwardedProto = "X-Forwarded-Proto" 78 | HeaderXHTTPMethodOverride = "X-HTTP-Method-Override" 79 | HeaderXForwardedFor = "X-Forwarded-For" 80 | HeaderXRealIP = "X-Real-IP" 81 | HeaderXCSRFToken = "X-CSRF-Token" 82 | HeaderXRatelimitLimit = "X-Ratelimit-Limit" 83 | HeaderXRatelimitRemaining = "X-Ratelimit-Remaining" 84 | HeaderXRatelimitReset = "X-Ratelimit-Reset" 85 | HeaderXForwardedPrefix = "X-Forwarded-Prefix" 86 | ) 87 | 88 | // Normalize formats the input header to the formation of "Xxx-Xxx". 89 | func Normalize(header string) string { 90 | return textproto.CanonicalMIMEHeaderKey(header) 91 | } 92 | -------------------------------------------------------------------------------- /handler/header_types_test.go: -------------------------------------------------------------------------------- 1 | package handler 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestNormalize(t *testing.T) { 10 | assert := assert.New(t) 11 | 12 | assert.Equal("Content-Type", Normalize("content-type")) 13 | assert.Equal("Content-Type", Normalize("CONTENT-TYPE")) 14 | assert.Equal("Content-Type", Normalize("cONtENT-tYpE")) 15 | } 16 | -------------------------------------------------------------------------------- /logx/logrusx/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/logx/logrusx 2 | 3 | go 1.16 4 | 5 | require ( 6 | github.com/sirupsen/logrus v1.9.3 7 | github.com/zeromicro/go-zero v1.5.4 8 | ) 9 | -------------------------------------------------------------------------------- /logx/logrusx/logrus.go: -------------------------------------------------------------------------------- 1 | package logrusx 2 | 3 | import ( 4 | "github.com/sirupsen/logrus" 5 | "github.com/zeromicro/go-zero/core/logx" 6 | ) 7 | 8 | type LogrusWriter struct { 9 | logger *logrus.Logger 10 | } 11 | 12 | func NewLogrusWriter(opts ...func(logger *logrus.Logger)) logx.Writer { 13 | logger := logrus.New() 14 | for _, opt := range opts { 15 | opt(logger) 16 | } 17 | 18 | return &LogrusWriter{ 19 | logger: logger, 20 | } 21 | } 22 | 23 | func (w *LogrusWriter) Alert(v interface{}) { 24 | w.logger.Error(v) 25 | } 26 | 27 | func (w *LogrusWriter) Close() error { 28 | w.logger.Exit(0) 29 | return nil 30 | } 31 | 32 | func (w *LogrusWriter) Debug(v interface{}, fields ...logx.LogField) { 33 | w.logger.WithFields(toLogrusFields(fields...)).Debug(v) 34 | } 35 | 36 | func (w *LogrusWriter) Error(v interface{}, fields ...logx.LogField) { 37 | w.logger.WithFields(toLogrusFields(fields...)).Error(v) 38 | } 39 | 40 | func (w *LogrusWriter) Info(v interface{}, fields ...logx.LogField) { 41 | w.logger.WithFields(toLogrusFields(fields...)).Info(v) 42 | } 43 | 44 | func (w *LogrusWriter) Severe(v interface{}) { 45 | w.logger.Fatal(v) 46 | } 47 | 48 | func (w *LogrusWriter) Slow(v interface{}, fields ...logx.LogField) { 49 | w.logger.WithFields(toLogrusFields(fields...)).Warn(v) 50 | } 51 | 52 | func (w *LogrusWriter) Stack(v interface{}) { 53 | w.logger.Error(v) 54 | } 55 | 56 | func (w *LogrusWriter) Stat(v interface{}, fields ...logx.LogField) { 57 | w.logger.WithFields(toLogrusFields(fields...)).Info(v) 58 | } 59 | 60 | func toLogrusFields(fields ...logx.LogField) logrus.Fields { 61 | logrusFields := make(logrus.Fields) 62 | for _, field := range fields { 63 | logrusFields[field.Key] = field.Value 64 | } 65 | return logrusFields 66 | } 67 | -------------------------------------------------------------------------------- /logx/logrusx/readme.md: -------------------------------------------------------------------------------- 1 | ### Quick Start 2 | 3 | Prerequesites: 4 | 5 | - Install `go-zero`: 6 | 7 | ```console 8 | go get -u github.com/zeromicro/go-zero 9 | ``` 10 | 11 | Download the module: 12 | 13 | ```console 14 | go get -u github.com/zeromicro/zero-contrib/logx/logrusx 15 | ``` 16 | 17 | For example: 18 | 19 | ```go 20 | package main 21 | 22 | import ( 23 | "context" 24 | "time" 25 | 26 | "github.com/zeromicro/go-zero/core/logx" 27 | "github.com/zeromicro/zero-contrib/logx/logrusx" 28 | ) 29 | 30 | func main() { 31 | writer := logrusx.NewLogrusWriter(func(logger *logrus.Logger) { 32 | logger.SetFormatter(&logrus.JSONFormatter{}) 33 | }) 34 | logx.SetWriter(writer) 35 | 36 | logx.Infow("infow foo", 37 | logx.Field("url", "http://localhost:8080/hello"), 38 | logx.Field("attempt", 3), 39 | logx.Field("backoff", time.Second), 40 | ) 41 | logx.Errorw("errorw foo", 42 | logx.Field("url", "http://localhost:8080/hello"), 43 | logx.Field("attempt", 3), 44 | logx.Field("backoff", time.Second), 45 | ) 46 | logx.Sloww("sloww foo", 47 | logx.Field("url", "http://localhost:8080/hello"), 48 | logx.Field("attempt", 3), 49 | logx.Field("backoff", time.Second), 50 | ) 51 | logx.Error("error") 52 | logx.Infov(map[string]interface{}{ 53 | "url": "localhost:8080/hello", 54 | "attempt": 3, 55 | "backoff": time.Second, 56 | "value": "foo", 57 | }) 58 | logx.WithDuration(1100*time.Microsecond).Infow("infow withduration", 59 | logx.Field("url", "localhost:8080/hello"), 60 | logx.Field("attempt", 3), 61 | logx.Field("backoff", time.Second), 62 | ) 63 | logx.WithContext(context.Background()).WithDuration(1100*time.Microsecond).Errorw( 64 | "errorw withcontext withduration", 65 | logx.Field("url", "localhost:8080/hello"), 66 | logx.Field("attempt", 3), 67 | logx.Field("backoff", time.Second), 68 | ) 69 | logx.WithDuration(1100*time.Microsecond).WithContext(context.Background()).Errorw( 70 | "errorw withduration withcontext", 71 | logx.Field("url", "localhost:8080/hello"), 72 | logx.Field("attempt", 3), 73 | logx.Field("backoff", time.Second), 74 | ) 75 | } 76 | ``` 77 | -------------------------------------------------------------------------------- /logx/zapx/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/logx/zapx 2 | 3 | go 1.16 4 | 5 | require ( 6 | github.com/zeromicro/go-zero v1.5.4 7 | go.uber.org/zap v1.26.0 8 | ) 9 | -------------------------------------------------------------------------------- /logx/zapx/readme.md: -------------------------------------------------------------------------------- 1 | ### Quick Start 2 | 3 | Prerequesites: 4 | 5 | - Install `go-zero`: 6 | 7 | ```console 8 | go get -u github.com/zeromicro/go-zero 9 | ``` 10 | 11 | Download the module: 12 | 13 | ```console 14 | go get -u github.com/zeromicro/zero-contrib/logx/zapx 15 | ``` 16 | 17 | For example: 18 | 19 | ```go 20 | package main 21 | 22 | import ( 23 | "context" 24 | "time" 25 | 26 | "github.com/zeromicro/go-zero/core/logx" 27 | "github.com/zeromicro/zero-contrib/logx/zapx" 28 | ) 29 | 30 | func main() { 31 | writer, err := zapx.NewZapWriter() 32 | logx.Must(err) 33 | logx.SetWriter(writer) 34 | 35 | logx.Infow("infow foo", 36 | logx.Field("url", "http://localhost:8080/hello"), 37 | logx.Field("attempt", 3), 38 | logx.Field("backoff", time.Second), 39 | ) 40 | logx.Errorw("errorw foo", 41 | logx.Field("url", "http://localhost:8080/hello"), 42 | logx.Field("attempt", 3), 43 | logx.Field("backoff", time.Second), 44 | ) 45 | logx.Sloww("sloww foo", 46 | logx.Field("url", "http://localhost:8080/hello"), 47 | logx.Field("attempt", 3), 48 | logx.Field("backoff", time.Second), 49 | ) 50 | logx.Error("error") 51 | logx.Infov(map[string]interface{}{ 52 | "url": "localhost:8080/hello", 53 | "attempt": 3, 54 | "backoff": time.Second, 55 | "value": "foo", 56 | }) 57 | logx.WithDuration(1100*time.Microsecond).Infow("infow withduration", 58 | logx.Field("url", "localhost:8080/hello"), 59 | logx.Field("attempt", 3), 60 | logx.Field("backoff", time.Second), 61 | ) 62 | logx.WithContext(context.Background()).WithDuration(1100*time.Microsecond).Errorw( 63 | "errorw withcontext withduration", 64 | logx.Field("url", "localhost:8080/hello"), 65 | logx.Field("attempt", 3), 66 | logx.Field("backoff", time.Second), 67 | ) 68 | logx.WithDuration(1100*time.Microsecond).WithContext(context.Background()).Errorw( 69 | "errorw withduration withcontext", 70 | logx.Field("url", "localhost:8080/hello"), 71 | logx.Field("attempt", 3), 72 | logx.Field("backoff", time.Second), 73 | ) 74 | } 75 | ``` 76 | -------------------------------------------------------------------------------- /logx/zapx/zap.go: -------------------------------------------------------------------------------- 1 | package zapx 2 | 3 | import ( 4 | "fmt" 5 | 6 | "github.com/zeromicro/go-zero/core/logx" 7 | "go.uber.org/zap" 8 | ) 9 | 10 | const callerSkipOffset = 3 11 | 12 | type ZapWriter struct { 13 | logger *zap.Logger 14 | } 15 | 16 | func NewZapWriter(opts ...zap.Option) (logx.Writer, error) { 17 | opts = append(opts, zap.AddCallerSkip(callerSkipOffset)) 18 | logger, err := zap.NewProduction(opts...) 19 | if err != nil { 20 | return nil, err 21 | } 22 | 23 | return &ZapWriter{ 24 | logger: logger, 25 | }, nil 26 | } 27 | 28 | func (w *ZapWriter) Alert(v interface{}) { 29 | w.logger.Error(fmt.Sprint(v)) 30 | } 31 | 32 | func (w *ZapWriter) Close() error { 33 | return w.logger.Sync() 34 | } 35 | 36 | func (w *ZapWriter) Debug(v interface{}, fields ...logx.LogField) { 37 | w.logger.Debug(fmt.Sprint(v), toZapFields(fields...)...) 38 | } 39 | 40 | func (w *ZapWriter) Error(v interface{}, fields ...logx.LogField) { 41 | w.logger.Error(fmt.Sprint(v), toZapFields(fields...)...) 42 | } 43 | 44 | func (w *ZapWriter) Info(v interface{}, fields ...logx.LogField) { 45 | w.logger.Info(fmt.Sprint(v), toZapFields(fields...)...) 46 | } 47 | 48 | func (w *ZapWriter) Severe(v interface{}) { 49 | w.logger.Fatal(fmt.Sprint(v)) 50 | } 51 | 52 | func (w *ZapWriter) Slow(v interface{}, fields ...logx.LogField) { 53 | w.logger.Warn(fmt.Sprint(v), toZapFields(fields...)...) 54 | } 55 | 56 | func (w *ZapWriter) Stack(v interface{}) { 57 | w.logger.Error(fmt.Sprint(v), zap.Stack("stack")) 58 | } 59 | 60 | func (w *ZapWriter) Stat(v interface{}, fields ...logx.LogField) { 61 | w.logger.Info(fmt.Sprint(v), toZapFields(fields...)...) 62 | } 63 | 64 | func toZapFields(fields ...logx.LogField) []zap.Field { 65 | zapFields := make([]zap.Field, 0, len(fields)) 66 | for _, f := range fields { 67 | zapFields = append(zapFields, zap.Any(f.Key, f.Value)) 68 | } 69 | return zapFields 70 | } 71 | -------------------------------------------------------------------------------- /logx/zerologx/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/logx/zerologx 2 | 3 | go 1.16 4 | 5 | require ( 6 | github.com/rs/zerolog v1.31.0 7 | github.com/zeromicro/go-zero v1.5.4 8 | ) 9 | -------------------------------------------------------------------------------- /logx/zerologx/readme.md: -------------------------------------------------------------------------------- 1 | ### Quick Start 2 | 3 | Prerequesites: 4 | 5 | - Install `go-zero`: 6 | 7 | ```console 8 | go get -u github.com/zeromicro/go-zero 9 | ``` 10 | 11 | Download the module: 12 | 13 | ```console 14 | go get -u github.com/zeromicro/zero-contrib/logx/zerologx 15 | ``` 16 | 17 | For example: 18 | 19 | ```go 20 | package main 21 | 22 | import ( 23 | "context" 24 | "github.com/rs/zerolog" 25 | "github.com/zeromicro/go-zero/core/logx" 26 | "github.com/zeromicro/zero-contrib/logx/zerologx" 27 | "os" 28 | "time" 29 | ) 30 | 31 | func main() { 32 | logger := zerolog.New(os.Stderr).With().Timestamp().Logger() 33 | writer := zerologx.NewZeroLogWriter(logger) 34 | logx.SetWriter(writer) 35 | 36 | logx.Infow("infow foo", 37 | logx.Field("url", "http://localhost:8080/hello"), 38 | logx.Field("attempt", 3), 39 | logx.Field("backoff", time.Second), 40 | ) 41 | logx.Errorw("errorw foo", 42 | logx.Field("url", "http://localhost:8080/hello"), 43 | logx.Field("attempt", 3), 44 | logx.Field("backoff", time.Second), 45 | ) 46 | logx.Sloww("sloww foo", 47 | logx.Field("url", "http://localhost:8080/hello"), 48 | logx.Field("attempt", 3), 49 | logx.Field("backoff", time.Second), 50 | ) 51 | logx.Error("error") 52 | logx.Infov(map[string]interface{}{ 53 | "url": "localhost:8080/hello", 54 | "attempt": 3, 55 | "backoff": time.Second, 56 | "value": "foo", 57 | }) 58 | logx.WithDuration(1100*time.Microsecond).Infow("infow withduration", 59 | logx.Field("url", "localhost:8080/hello"), 60 | logx.Field("attempt", 3), 61 | logx.Field("backoff", time.Second), 62 | ) 63 | logx.WithContext(context.Background()).WithDuration(1100*time.Microsecond).Errorw( 64 | "errorw withcontext withduration", 65 | logx.Field("url", "localhost:8080/hello"), 66 | logx.Field("attempt", 3), 67 | logx.Field("backoff", time.Second), 68 | ) 69 | logx.WithDuration(1100*time.Microsecond).WithContext(context.Background()).Errorw( 70 | "errorw withduration withcontext", 71 | logx.Field("url", "localhost:8080/hello"), 72 | logx.Field("attempt", 3), 73 | logx.Field("backoff", time.Second), 74 | ) 75 | } 76 | ``` -------------------------------------------------------------------------------- /logx/zerologx/zerolog.go: -------------------------------------------------------------------------------- 1 | package zerologx 2 | 3 | import ( 4 | "fmt" 5 | 6 | "github.com/rs/zerolog" 7 | "github.com/zeromicro/go-zero/core/logx" 8 | ) 9 | 10 | type ( 11 | ZeroLogWriter struct { 12 | logger zerolog.Logger 13 | } 14 | ) 15 | 16 | func NewZeroLogWriter(logger zerolog.Logger) logx.Writer { 17 | return &ZeroLogWriter{logger: logger} 18 | } 19 | 20 | func (w *ZeroLogWriter) Alert(v interface{}) { 21 | w.logger.Error().Msg(fmt.Sprint(v)) 22 | } 23 | 24 | func (w *ZeroLogWriter) Close() error { 25 | w.logger.Fatal().Msg("") 26 | return nil 27 | } 28 | 29 | func (w *ZeroLogWriter) Debug(v interface{}, fields ...logx.LogField) { 30 | toZeroLogInterface(w.logger.Debug(), fields...).Msgf(fmt.Sprint(v)) 31 | } 32 | 33 | func (w *ZeroLogWriter) Error(v interface{}, fields ...logx.LogField) { 34 | toZeroLogInterface(w.logger.Error(), fields...).Msgf(fmt.Sprint(v)) 35 | } 36 | 37 | func (w *ZeroLogWriter) Info(v interface{}, fields ...logx.LogField) { 38 | toZeroLogInterface(w.logger.Info(), fields...).Msgf(fmt.Sprint(v)) 39 | } 40 | 41 | func (w *ZeroLogWriter) Severe(v interface{}) { 42 | w.logger.Fatal().Msg(fmt.Sprint(v)) 43 | } 44 | 45 | func (w *ZeroLogWriter) Slow(v interface{}, fields ...logx.LogField) { 46 | toZeroLogInterface(w.logger.Warn(), fields...).Msgf(fmt.Sprint(v)) 47 | } 48 | 49 | func (w *ZeroLogWriter) Stack(v interface{}) { 50 | w.logger.Error().Msgf(fmt.Sprint(v)) 51 | } 52 | 53 | func (w *ZeroLogWriter) Stat(v interface{}, fields ...logx.LogField) { 54 | toZeroLogInterface(w.logger.Info(), fields...).Msgf(fmt.Sprint(v)) 55 | } 56 | 57 | func toZeroLogInterface(event *zerolog.Event, fields ...logx.LogField) *zerolog.Event { 58 | for _, field := range fields { 59 | event = event.Interface(field.Key, field.Value) 60 | } 61 | return event 62 | } 63 | -------------------------------------------------------------------------------- /logx/zerologx/zerolog_test.go: -------------------------------------------------------------------------------- 1 | package zerologx 2 | 3 | import ( 4 | "context" 5 | "github.com/rs/zerolog" 6 | "github.com/zeromicro/go-zero/core/logx" 7 | "os" 8 | "testing" 9 | "time" 10 | ) 11 | 12 | func Test_Log(t *testing.T) { 13 | logger := zerolog.New(os.Stderr).With().Timestamp().Logger() 14 | writer := NewZeroLogWriter(logger) 15 | logx.SetWriter(writer) 16 | 17 | logx.Infow("infow foo", 18 | logx.Field("url", "http://localhost:8080/hello"), 19 | logx.Field("attempt", 3), 20 | logx.Field("backoff", time.Second), 21 | ) 22 | logx.Errorw("errorw foo", 23 | logx.Field("url", "http://localhost:8080/hello"), 24 | logx.Field("attempt", 3), 25 | logx.Field("backoff", time.Second), 26 | ) 27 | logx.Sloww("sloww foo", 28 | logx.Field("url", "http://localhost:8080/hello"), 29 | logx.Field("attempt", 3), 30 | logx.Field("backoff", time.Second), 31 | ) 32 | logx.Error("error") 33 | logx.Infov(map[string]interface{}{ 34 | "url": "localhost:8080/hello", 35 | "attempt": 3, 36 | "backoff": time.Second, 37 | "value": "foo", 38 | }) 39 | logx.WithDuration(1100*time.Microsecond).Infow("infow withduration", 40 | logx.Field("url", "localhost:8080/hello"), 41 | logx.Field("attempt", 3), 42 | logx.Field("backoff", time.Second), 43 | ) 44 | logx.WithContext(context.Background()).WithDuration(1100*time.Microsecond).Errorw( 45 | "errorw withcontext withduration", 46 | logx.Field("url", "localhost:8080/hello"), 47 | logx.Field("attempt", 3), 48 | logx.Field("backoff", time.Second), 49 | ) 50 | logx.WithDuration(1100*time.Microsecond).WithContext(context.Background()).Errorw( 51 | "errorw withduration withcontext", 52 | logx.Field("url", "localhost:8080/hello"), 53 | logx.Field("attempt", 3), 54 | logx.Field("backoff", time.Second), 55 | ) 56 | } 57 | -------------------------------------------------------------------------------- /rest/registry/etcd/README.md: -------------------------------------------------------------------------------- 1 | ### Quick Start 2 | 3 | Prerequisites: 4 | 5 | Download the module: 6 | 7 | ```console 8 | go get -u github.com/zeromicro/zero-contrib/rest/registry/etcd 9 | ``` 10 | 11 | For example: 12 | 13 | ## 修改REST服务的代码 14 | 15 | - etc/*.yaml 16 | 17 | ```yaml 18 | Name: user.api 19 | Host: 0.0.0.0 20 | Port: 8888 21 | Etcd: 22 | Hosts: 23 | - etcd:2379 24 | Key: user.api 25 | 26 | ``` 27 | 28 | - internal/config/config.go 29 | 30 | ```go 31 | type Config struct { 32 | rest.RestConf 33 | Etcd discov.EtcdConf // etcd register center config 34 | } 35 | ``` 36 | 37 | - main.go 38 | 39 | ```go 40 | package main 41 | 42 | import ( 43 | "flag" 44 | 45 | "github.com/zeromicro/go-zero/conf" 46 | "github.com/zeromicro/go-zero/rest" 47 | "github.com/zeromicro/zero-contrib/rest/registry/etcd" 48 | ) 49 | 50 | var configFile = flag.String("f", "etc/user-api.yaml", "the config file") 51 | 52 | func main() { 53 | flag.Parse() 54 | 55 | var c Config 56 | conf.MustLoad(*configFile, &c) 57 | 58 | server := rest.MustNewServer(c.RestConf) 59 | // register rest to etcd 60 | logx.Must(etcd.RegisterRest(c.Etcd, c.RestConf)) 61 | 62 | server.Start() 63 | } 64 | ``` -------------------------------------------------------------------------------- /rest/registry/etcd/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/rest/registry/etcd 2 | 3 | go 1.16 4 | 5 | require github.com/zeromicro/go-zero v1.5.4 6 | -------------------------------------------------------------------------------- /rest/registry/etcd/register.go: -------------------------------------------------------------------------------- 1 | package etcd 2 | 3 | import ( 4 | "fmt" 5 | "os" 6 | "strings" 7 | 8 | "github.com/zeromicro/go-zero/core/discov" 9 | "github.com/zeromicro/go-zero/core/netx" 10 | "github.com/zeromicro/go-zero/core/proc" 11 | "github.com/zeromicro/go-zero/rest" 12 | ) 13 | 14 | const ( 15 | allEths = "0.0.0.0" 16 | envPodIP = "POD_IP" 17 | ) 18 | 19 | // RegisterRest register reset to etcd. 20 | func RegisterRest(etcd discov.EtcdConf, svrConf rest.RestConf) error { 21 | if err := etcd.Validate(); err != nil { 22 | return err 23 | } 24 | 25 | listenOn := fmt.Sprintf("%s:%d", svrConf.Host, svrConf.Port) 26 | pubListenOn := figureOutListenOn(listenOn) 27 | var pubOpts []discov.PubOption 28 | if etcd.HasAccount() { 29 | pubOpts = append(pubOpts, discov.WithPubEtcdAccount(etcd.User, etcd.Pass)) 30 | } 31 | if etcd.HasTLS() { 32 | pubOpts = append(pubOpts, discov.WithPubEtcdTLS(etcd.CertFile, etcd.CertKeyFile, 33 | etcd.CACertFile, etcd.InsecureSkipVerify)) 34 | } 35 | 36 | pubClient := discov.NewPublisher(etcd.Hosts, etcd.Key, pubListenOn, pubOpts...) 37 | proc.AddShutdownListener(func() { 38 | pubClient.Stop() 39 | }) 40 | 41 | return pubClient.KeepAlive() 42 | } 43 | 44 | func figureOutListenOn(listenOn string) string { 45 | fields := strings.Split(listenOn, ":") 46 | if len(fields) == 0 { 47 | return listenOn 48 | } 49 | 50 | host := fields[0] 51 | if len(host) > 0 && host != allEths { 52 | return listenOn 53 | } 54 | 55 | ip := os.Getenv(envPodIP) 56 | if len(ip) == 0 { 57 | ip = netx.InternalIp() 58 | } 59 | if len(ip) == 0 { 60 | return listenOn 61 | } 62 | 63 | return strings.Join(append([]string{ip}, fields[1:]...), ":") 64 | } 65 | -------------------------------------------------------------------------------- /router/chi/README.md: -------------------------------------------------------------------------------- 1 | ### Quick Start 2 | Prerequesites: 3 | * Install `go-zero`: go get -u github.com/zeromicro/go-zero@master 4 | 5 | 6 | Download the module: 7 | ```console 8 | go get -u github.com/zeromicro/zero-contrib/router/chi 9 | ``` 10 | 11 | For example: 12 | 13 | ```go 14 | package main 15 | 16 | import ( 17 | "github.com/zeromicro/go-zero/core/logx" 18 | "github.com/zeromicro/go-zero/core/service" 19 | "github.com/zeromicro/go-zero/rest" 20 | "github.com/zeromicro/go-zero/rest/httpx" 21 | "github.com/zeromicro/zero-contrib/router/chi" 22 | "net/http" 23 | "strings" 24 | ) 25 | 26 | type CommonPath struct { 27 | Year int `path:"year"` 28 | Month int `path:"month"` 29 | Day int `path:"day"` 30 | } 31 | 32 | func (c *CommonPath) String() string { 33 | var builder strings.Builder 34 | builder.WriteString("CommonPath(") 35 | builder.WriteString(fmt.Sprintf("Year=%v", c.Year)) 36 | builder.WriteString(fmt.Sprintf(", Month=%v", c.Month)) 37 | builder.WriteString(fmt.Sprintf(", Day=%v", c.Day)) 38 | builder.WriteByte(')') 39 | return builder.String() 40 | } 41 | 42 | func main() { 43 | r := chi.NewRouter() 44 | engine := rest.MustNewServer(rest.RestConf{ 45 | ServiceConf: service.ServiceConf{ 46 | Log: logx.LogConf{ 47 | Mode: "console", 48 | }, 49 | }, 50 | Port: 3345, 51 | Timeout: 20000, 52 | MaxConns: 500, 53 | }, rest.WithRouter(r)) 54 | 55 | // NotFound defines a handler to respond whenever a route could 56 | // not be found. 57 | r.SetNotFoundHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 58 | w.WriteHeader(404) 59 | w.Write([]byte("nothing here")) 60 | })) 61 | // MethodNotAllowed defines a handler to respond whenever a method is 62 | // not allowed. 63 | r.SetNotAllowedHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 64 | w.WriteHeader(405) 65 | w.Write([]byte("405")) 66 | })) 67 | defer engine.Stop() 68 | 69 | engine.AddRoute(rest.Route{ 70 | Method: http.MethodGet, 71 | Path: "/api/{month}-{day}-{year}", // GET /articles/01-16-2017 72 | Handler: func(w http.ResponseWriter, r *http.Request) { 73 | var commonPath CommonPath 74 | err := httpx.Parse(r, &commonPath) 75 | if err != nil { 76 | return 77 | } 78 | w.Write([]byte(commonPath.String())) 79 | 80 | }, 81 | }) 82 | engine.Start() 83 | } 84 | 85 | ``` 86 | -------------------------------------------------------------------------------- /router/chi/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/router/chi 2 | 3 | go 1.15 4 | 5 | require ( 6 | github.com/go-chi/chi v1.5.4 7 | github.com/stretchr/testify v1.8.4 8 | github.com/zeromicro/go-zero v1.5.4 9 | ) 10 | -------------------------------------------------------------------------------- /router/chi/router.go: -------------------------------------------------------------------------------- 1 | package chi 2 | 3 | import ( 4 | "errors" 5 | "net/http" 6 | 7 | "github.com/go-chi/chi" 8 | "github.com/zeromicro/go-zero/rest/httpx" 9 | "github.com/zeromicro/go-zero/rest/pathvar" 10 | ) 11 | 12 | var ( 13 | // ErrInvalidMethod is an error that indicates not a valid http method. 14 | ErrInvalidMethod = errors.New("not a valid http method") 15 | // ErrInvalidPath is an error that indicates path is not start with /. 16 | ErrInvalidPath = errors.New("path must begin with '/'") 17 | ) 18 | 19 | type chiRouter struct { 20 | mux *chi.Mux 21 | } 22 | 23 | // NewRouter returns a chi.Router. 24 | func NewRouter() httpx.Router { 25 | return &chiRouter{ 26 | mux: chi.NewRouter(), 27 | } 28 | } 29 | 30 | func (pr *chiRouter) Handle(method, reqPath string, handler http.Handler) error { 31 | if !validMethod(method) { 32 | return ErrInvalidMethod 33 | } 34 | 35 | if len(reqPath) == 0 || reqPath[0] != '/' { 36 | return ErrInvalidPath 37 | } 38 | 39 | pr.getHandleFunc(method)(reqPath, func(w http.ResponseWriter, r *http.Request) { 40 | handler.ServeHTTP(w, pr.withUrlParamsToContext(r)) 41 | }) 42 | return nil 43 | } 44 | 45 | func (pr *chiRouter) withUrlParamsToContext(r *http.Request) *http.Request { 46 | if ctx := chi.RouteContext(r.Context()); ctx != nil { 47 | urlParams := ctx.URLParams 48 | params := make(map[string]string) 49 | for i := 0; i < len(urlParams.Values); i++ { 50 | params[urlParams.Keys[i]] = urlParams.Values[i] 51 | } 52 | if len(params) > 0 { 53 | r = pathvar.WithVars(r, params) 54 | } 55 | } 56 | return r 57 | } 58 | 59 | func (pr *chiRouter) getHandleFunc(method string) func(pattern string, handlerFn http.HandlerFunc) { 60 | switch method { 61 | case http.MethodGet: 62 | return pr.mux.Get 63 | case http.MethodPost: 64 | return pr.mux.Post 65 | case http.MethodPut: 66 | return pr.mux.Put 67 | case http.MethodDelete: 68 | return pr.mux.Delete 69 | case http.MethodHead: 70 | return pr.mux.Head 71 | case http.MethodOptions: 72 | return pr.mux.Options 73 | case http.MethodPatch: 74 | return pr.mux.Patch 75 | default: 76 | panic(ErrInvalidMethod) 77 | } 78 | } 79 | 80 | func (pr *chiRouter) ServeHTTP(w http.ResponseWriter, r *http.Request) { 81 | pr.mux.ServeHTTP(w, r) 82 | } 83 | 84 | func (pr *chiRouter) SetNotFoundHandler(handler http.Handler) { 85 | pr.mux.NotFound(func(w http.ResponseWriter, r *http.Request) { 86 | handler.ServeHTTP(w, r) 87 | }) 88 | } 89 | 90 | func (pr *chiRouter) SetNotAllowedHandler(handler http.Handler) { 91 | pr.mux.MethodNotAllowed(func(w http.ResponseWriter, r *http.Request) { 92 | handler.ServeHTTP(w, r) 93 | }) 94 | } 95 | 96 | func validMethod(method string) bool { 97 | return method == http.MethodDelete || method == http.MethodGet || 98 | method == http.MethodHead || method == http.MethodOptions || 99 | method == http.MethodPatch || method == http.MethodPost || 100 | method == http.MethodPut 101 | } 102 | -------------------------------------------------------------------------------- /router/chi/router_test.go: -------------------------------------------------------------------------------- 1 | package chi 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "io" 7 | "net/http" 8 | "net/http/httptest" 9 | "testing" 10 | 11 | "github.com/stretchr/testify/assert" 12 | "github.com/zeromicro/go-zero/rest/httpx" 13 | "github.com/zeromicro/go-zero/rest/pathvar" 14 | ) 15 | 16 | type mockedResponseWriter struct { 17 | code int 18 | } 19 | 20 | func (m *mockedResponseWriter) Header() http.Header { 21 | return http.Header{} 22 | } 23 | 24 | func (m *mockedResponseWriter) Write(p []byte) (int, error) { 25 | return len(p), nil 26 | } 27 | 28 | func (m *mockedResponseWriter) WriteHeader(code int) { 29 | m.code = code 30 | } 31 | 32 | func TestChiRouterNotFound(t *testing.T) { 33 | var notFound bool 34 | router := NewRouter() 35 | router.SetNotFoundHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 36 | notFound = true 37 | })) 38 | err := router.Handle(http.MethodGet, "/a/b", 39 | http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) 40 | assert.Nil(t, err) 41 | r, _ := http.NewRequest(http.MethodGet, "/b/c", nil) 42 | w := new(mockedResponseWriter) 43 | router.ServeHTTP(w, r) 44 | assert.True(t, notFound) 45 | } 46 | 47 | func TestChiRouterNotAllowed(t *testing.T) { 48 | var notAllowed bool 49 | router := NewRouter() 50 | router.SetNotAllowedHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 51 | notAllowed = true 52 | })) 53 | err := router.Handle(http.MethodGet, "/a/b", 54 | http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) 55 | assert.Nil(t, err) 56 | r, _ := http.NewRequest(http.MethodPost, "/a/b", nil) 57 | w := new(mockedResponseWriter) 58 | router.ServeHTTP(w, r) 59 | assert.True(t, notAllowed) 60 | } 61 | 62 | func TestChiRouter(t *testing.T) { 63 | tests := []struct { 64 | method string 65 | path string 66 | expect bool 67 | code int 68 | err error 69 | }{ 70 | // we don't explicitly set status code, framework will do it. 71 | {http.MethodGet, "/a/b", true, 0, nil}, 72 | {http.MethodGet, "/a/b?a=b", true, 0, nil}, 73 | {http.MethodGet, "/a/b/c?a=b", true, 0, nil}, 74 | {http.MethodGet, "/b/d", false, http.StatusNotFound, nil}, 75 | } 76 | 77 | for _, test := range tests { 78 | t.Run(test.method+":"+test.path, func(t *testing.T) { 79 | routed := false 80 | router := NewRouter() 81 | err := router.Handle(test.method, "/a/{b}", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 82 | routed = true 83 | assert.Equal(t, 1, len(pathvar.Vars(r))) 84 | })) 85 | assert.Nil(t, err) 86 | err = router.Handle(test.method, "/a/b/c", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 87 | routed = true 88 | assert.Nil(t, pathvar.Vars(r)) 89 | })) 90 | assert.Nil(t, err) 91 | err = router.Handle(test.method, "/b/c", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 92 | routed = true 93 | })) 94 | assert.Nil(t, err) 95 | 96 | w := new(mockedResponseWriter) 97 | r, _ := http.NewRequest(test.method, test.path, nil) 98 | router.ServeHTTP(w, r) 99 | assert.Equal(t, test.expect, routed) 100 | assert.Equal(t, test.code, w.code) 101 | 102 | if test.code == 0 { 103 | r, _ = http.NewRequest(http.MethodPut, test.path, nil) 104 | router.ServeHTTP(w, r) 105 | assert.Equal(t, http.StatusMethodNotAllowed, w.code) 106 | } 107 | }) 108 | } 109 | } 110 | 111 | func TestParseJsonPost(t *testing.T) { 112 | r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000", 113 | bytes.NewBufferString(`{"location": "shanghai", "time": 20170912}`)) 114 | assert.Nil(t, err) 115 | r.Header.Set(httpx.ContentType, httpx.JsonContentType) 116 | 117 | router := NewRouter() 118 | err = router.Handle(http.MethodPost, "/{name}/{year}", http.HandlerFunc(func( 119 | w http.ResponseWriter, r *http.Request) { 120 | v := struct { 121 | Name string `path:"name"` 122 | Year int `path:"year"` 123 | Nickname string `form:"nickname"` 124 | Zipcode int64 `form:"zipcode"` 125 | Location string `json:"location"` 126 | Time int64 `json:"time"` 127 | }{} 128 | 129 | err = httpx.Parse(r, &v) 130 | assert.Nil(t, err) 131 | _, err = io.WriteString(w, fmt.Sprintf("%s:%d:%s:%d:%s:%d", v.Name, v.Year, 132 | v.Nickname, v.Zipcode, v.Location, v.Time)) 133 | assert.Nil(t, err) 134 | })) 135 | assert.Nil(t, err) 136 | 137 | rr := httptest.NewRecorder() 138 | router.ServeHTTP(rr, r) 139 | 140 | assert.Equal(t, "kevin:2017:whatever:200000:shanghai:20170912", rr.Body.String()) 141 | } 142 | 143 | func TestParseJsonPostWithIntSlice(t *testing.T) { 144 | r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017", 145 | bytes.NewBufferString(`{"ages": [1, 2], "years": [3, 4]}`)) 146 | assert.Nil(t, err) 147 | r.Header.Set(httpx.ContentType, httpx.JsonContentType) 148 | 149 | router := NewRouter() 150 | err = router.Handle(http.MethodPost, "/{name}/{year}", http.HandlerFunc(func( 151 | w http.ResponseWriter, r *http.Request) { 152 | v := struct { 153 | Name string `path:"name"` 154 | Year int `path:"year"` 155 | Ages []int `json:"ages"` 156 | Years []int64 `json:"years"` 157 | }{} 158 | 159 | err = httpx.Parse(r, &v) 160 | assert.Nil(t, err) 161 | assert.ElementsMatch(t, []int{1, 2}, v.Ages) 162 | assert.ElementsMatch(t, []int64{3, 4}, v.Years) 163 | })) 164 | assert.Nil(t, err) 165 | 166 | rr := httptest.NewRecorder() 167 | router.ServeHTTP(rr, r) 168 | } 169 | 170 | func TestParseJsonPostError(t *testing.T) { 171 | payload := `[{"abcd": "cdef"}]` 172 | r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000", 173 | bytes.NewBufferString(payload)) 174 | assert.Nil(t, err) 175 | r.Header.Set(httpx.ContentType, httpx.JsonContentType) 176 | 177 | router := NewRouter() 178 | err = router.Handle(http.MethodPost, "/{name}/{year}", http.HandlerFunc( 179 | func(w http.ResponseWriter, r *http.Request) { 180 | v := struct { 181 | Name string `path:"name"` 182 | Year int `path:"year"` 183 | Nickname string `form:"nickname"` 184 | Zipcode int64 `form:"zipcode"` 185 | Location string `json:"location"` 186 | Time int64 `json:"time"` 187 | }{} 188 | 189 | err = httpx.Parse(r, &v) 190 | assert.NotNil(t, err) 191 | })) 192 | assert.Nil(t, err) 193 | 194 | rr := httptest.NewRecorder() 195 | router.ServeHTTP(rr, r) 196 | } 197 | 198 | func TestParsePath(t *testing.T) { 199 | r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin/2017", nil) 200 | assert.Nil(t, err) 201 | 202 | router := NewRouter() 203 | err = router.Handle(http.MethodGet, "/{name}/{year}", http.HandlerFunc( 204 | func(w http.ResponseWriter, r *http.Request) { 205 | v := struct { 206 | Name string `path:"name"` 207 | Year int `path:"year"` 208 | }{} 209 | 210 | err = httpx.Parse(r, &v) 211 | assert.Nil(t, err) 212 | _, err = io.WriteString(w, fmt.Sprintf("%s in %d", v.Name, v.Year)) 213 | assert.Nil(t, err) 214 | })) 215 | assert.Nil(t, err) 216 | 217 | rr := httptest.NewRecorder() 218 | router.ServeHTTP(rr, r) 219 | 220 | assert.Equal(t, "kevin in 2017", rr.Body.String()) 221 | } 222 | 223 | func TestParsePathRequired(t *testing.T) { 224 | r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin", nil) 225 | assert.Nil(t, err) 226 | 227 | router := NewRouter() 228 | err = router.Handle(http.MethodGet, "/{name}/", http.HandlerFunc( 229 | func(w http.ResponseWriter, r *http.Request) { 230 | v := struct { 231 | Name string `path:"name"` 232 | Year int `path:"year"` 233 | }{} 234 | 235 | err = httpx.Parse(r, &v) 236 | assert.NotNil(t, err) 237 | })) 238 | assert.Nil(t, err) 239 | 240 | rr := httptest.NewRecorder() 241 | router.ServeHTTP(rr, r) 242 | } 243 | 244 | func TestParseQuery(t *testing.T) { 245 | r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000", nil) 246 | assert.Nil(t, err) 247 | 248 | router := NewRouter() 249 | err = router.Handle(http.MethodGet, "/{name}/{year}", http.HandlerFunc( 250 | func(w http.ResponseWriter, r *http.Request) { 251 | v := struct { 252 | Nickname string `form:"nickname"` 253 | Zipcode int64 `form:"zipcode"` 254 | }{} 255 | 256 | err = httpx.Parse(r, &v) 257 | assert.Nil(t, err) 258 | _, err = io.WriteString(w, fmt.Sprintf("%s:%d", v.Nickname, v.Zipcode)) 259 | assert.Nil(t, err) 260 | })) 261 | assert.Nil(t, err) 262 | 263 | rr := httptest.NewRecorder() 264 | router.ServeHTTP(rr, r) 265 | 266 | assert.Equal(t, "whatever:200000", rr.Body.String()) 267 | } 268 | 269 | func TestParseOptional(t *testing.T) { 270 | r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin/2017?nickname=whatever&zipcode=", nil) 271 | assert.Nil(t, err) 272 | 273 | router := NewRouter() 274 | err = router.Handle(http.MethodGet, "/{name}/{year}", http.HandlerFunc( 275 | func(w http.ResponseWriter, r *http.Request) { 276 | v := struct { 277 | Nickname string `form:"nickname"` 278 | Zipcode int64 `form:"zipcode,optional"` 279 | }{} 280 | 281 | err = httpx.Parse(r, &v) 282 | assert.Nil(t, err) 283 | _, err = io.WriteString(w, fmt.Sprintf("%s:%d", v.Nickname, v.Zipcode)) 284 | assert.Nil(t, err) 285 | })) 286 | assert.Nil(t, err) 287 | 288 | rr := httptest.NewRecorder() 289 | router.ServeHTTP(rr, r) 290 | 291 | assert.Equal(t, "whatever:0", rr.Body.String()) 292 | } 293 | 294 | func BenchmarkChiRouter(b *testing.B) { 295 | b.ReportAllocs() 296 | 297 | router := NewRouter() 298 | router.Handle(http.MethodGet, "/api/{articleSlug:[a-z-]+}", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 299 | })) 300 | w := &mockedResponseWriter{} 301 | r, _ := http.NewRequest(http.MethodGet, "/api/home-is-toronto", nil) 302 | for i := 0; i < b.N; i++ { 303 | router.ServeHTTP(w, r) 304 | } 305 | } 306 | -------------------------------------------------------------------------------- /router/gin/README.md: -------------------------------------------------------------------------------- 1 | ### Quick Start 2 | 3 | Prerequesites: 4 | 5 | - Install `go-zero`: 6 | 7 | ```console 8 | go get -u github.com/zeromicro/go-zero 9 | ``` 10 | 11 | Download the module: 12 | 13 | ```console 14 | go get -u github.com/zeromicro/zero-contrib/router/gin 15 | ``` 16 | 17 | For example: 18 | 19 | ```go 20 | package main 21 | 22 | import ( 23 | "net/http" 24 | "strings" 25 | 26 | "github.com/zeromicro/go-zero/core/logx" 27 | "github.com/zeromicro/go-zero/core/service" 28 | "github.com/zeromicro/go-zero/rest" 29 | "github.com/zeromicro/go-zero/rest/httpx" 30 | "github.com/zeromicro/zero-contrib/router/gin" 31 | stdgin "github.com/gin-gonic/gin" 32 | ) 33 | 34 | type CommonPathID struct { 35 | ID int `path:"id"` 36 | Name string `path:"name"` 37 | } 38 | 39 | func (c *CommonPathID) String() string { 40 | var builder strings.Builder 41 | builder.WriteString("CommonPathID(") 42 | builder.WriteString(fmt.Sprintf("ID=%v", c.ID)) 43 | builder.WriteString(fmt.Sprintf(", Name=%s", c.Name)) 44 | builder.WriteByte(')') 45 | return builder.String() 46 | } 47 | 48 | func init() { 49 | stdgin.SetMode(stdgin.ReleaseMode) 50 | } 51 | 52 | func main() { 53 | r := gin.NewRouter() 54 | engine := rest.MustNewServer(rest.RestConf{ 55 | ServiceConf: service.ServiceConf{ 56 | Log: logx.LogConf{ 57 | Mode: "console", 58 | }, 59 | }, 60 | Port: 3345, 61 | Timeout: 20000, 62 | MaxConns: 500, 63 | }, rest.WithRouter(r)) 64 | 65 | // NotFound defines a handler to respond whenever a route could 66 | // not be found. 67 | r.SetNotFoundHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 68 | w.WriteHeader(404) 69 | w.Write([]byte("nothing here")) 70 | })) 71 | // MethodNotAllowed defines a handler to respond whenever a method is 72 | // not allowed. 73 | r.SetNotAllowedHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 74 | w.WriteHeader(405) 75 | w.Write([]byte("405")) 76 | })) 77 | defer engine.Stop() 78 | 79 | engine.AddRoute(rest.Route{ 80 | Method: http.MethodGet, 81 | Path: "/api/:name/:id", // GET /api/joh/123 82 | Handler: func(w http.ResponseWriter, r *http.Request) { 83 | var commonPath CommonPathID 84 | err := httpx.Parse(r, &commonPath) 85 | if err != nil { 86 | return 87 | } 88 | w.Write([]byte(commonPath.String())) 89 | }, 90 | }) 91 | engine.Start() 92 | } 93 | ``` 94 | -------------------------------------------------------------------------------- /router/gin/config.go: -------------------------------------------------------------------------------- 1 | package gin 2 | 3 | import "errors" 4 | 5 | var ( 6 | // ErrInvalidMethod is an error that indicates not a valid http method. 7 | ErrInvalidMethod = errors.New("not a valid http method") 8 | // ErrInvalidPath is an error that indicates path is not start with /. 9 | ErrInvalidPath = errors.New("path must begin with '/'") 10 | ) 11 | 12 | type config struct { 13 | redirectTrailingSlash bool 14 | redirectFixedPath bool 15 | } 16 | 17 | func (c *config) options(opts ...Option) { 18 | for _, opt := range opts { 19 | opt(c) 20 | } 21 | } 22 | 23 | type Option func(o *config) 24 | 25 | func WithRedirectTrailingSlash(redirectTrailingSlash bool) Option { 26 | return func(c *config) { 27 | c.redirectTrailingSlash = redirectTrailingSlash 28 | } 29 | } 30 | 31 | func WithRedirectFixedPath(redirectFixedPath bool) Option { 32 | return func(c *config) { 33 | c.redirectFixedPath = redirectFixedPath 34 | } 35 | } 36 | -------------------------------------------------------------------------------- /router/gin/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/router/gin 2 | 3 | go 1.15 4 | 5 | require ( 6 | github.com/gin-gonic/gin v1.9.1 7 | github.com/stretchr/testify v1.8.4 8 | github.com/zeromicro/go-zero v1.5.4 9 | ) 10 | -------------------------------------------------------------------------------- /router/gin/router.go: -------------------------------------------------------------------------------- 1 | package gin 2 | 3 | import ( 4 | "net/http" 5 | "strings" 6 | 7 | "github.com/gin-gonic/gin" 8 | "github.com/zeromicro/go-zero/rest/httpx" 9 | "github.com/zeromicro/go-zero/rest/pathvar" 10 | ) 11 | 12 | type ginRouter struct { 13 | g *gin.Engine 14 | } 15 | 16 | // NewRouter returns a gin.Router. 17 | func NewRouter(opts ...Option) httpx.Router { 18 | g := gin.New() 19 | cfg := config{ 20 | redirectTrailingSlash: true, 21 | redirectFixedPath: false, 22 | } 23 | cfg.options(opts...) 24 | 25 | g.RedirectTrailingSlash = cfg.redirectTrailingSlash 26 | g.RedirectFixedPath = cfg.redirectFixedPath 27 | return &ginRouter{g: g} 28 | } 29 | 30 | func (pr *ginRouter) Handle(method, reqPath string, handler http.Handler) error { 31 | if !validMethod(method) { 32 | return ErrInvalidMethod 33 | } 34 | 35 | if len(reqPath) == 0 || reqPath[0] != '/' { 36 | return ErrInvalidPath 37 | } 38 | 39 | pr.g.Handle(strings.ToUpper(method), reqPath, func(ctx *gin.Context) { 40 | params := make(map[string]string) 41 | for i := 0; i < len(ctx.Params); i++ { 42 | params[ctx.Params[i].Key] = ctx.Params[i].Value 43 | } 44 | if len(params) > 0 { 45 | ctx.Request = pathvar.WithVars(ctx.Request, params) 46 | } 47 | handler.ServeHTTP(ctx.Writer, ctx.Request) 48 | }) 49 | return nil 50 | } 51 | 52 | func (pr *ginRouter) ServeHTTP(w http.ResponseWriter, r *http.Request) { 53 | pr.g.ServeHTTP(w, r) 54 | } 55 | 56 | func (pr *ginRouter) SetNotFoundHandler(handler http.Handler) { 57 | pr.g.NoRoute(gin.WrapH(handler)) 58 | } 59 | 60 | func (pr *ginRouter) SetNotAllowedHandler(handler http.Handler) { 61 | pr.g.NoMethod(gin.WrapH(handler)) 62 | } 63 | 64 | func validMethod(method string) bool { 65 | return method == http.MethodDelete || method == http.MethodGet || 66 | method == http.MethodHead || method == http.MethodOptions || 67 | method == http.MethodPatch || method == http.MethodPost || 68 | method == http.MethodPut 69 | } 70 | -------------------------------------------------------------------------------- /router/gin/router_test.go: -------------------------------------------------------------------------------- 1 | package gin 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "io" 7 | "net/http" 8 | "net/http/httptest" 9 | "testing" 10 | 11 | "github.com/stretchr/testify/assert" 12 | "github.com/zeromicro/go-zero/rest/httpx" 13 | "github.com/zeromicro/go-zero/rest/pathvar" 14 | ) 15 | 16 | type mockedResponseWriter struct { 17 | code int 18 | } 19 | 20 | func (m *mockedResponseWriter) Header() http.Header { 21 | return http.Header{} 22 | } 23 | 24 | func (m *mockedResponseWriter) Write(p []byte) (int, error) { 25 | return len(p), nil 26 | } 27 | 28 | func (m *mockedResponseWriter) WriteHeader(code int) { 29 | m.code = code 30 | } 31 | 32 | func TestChiRouterNotFound(t *testing.T) { 33 | var notFound bool 34 | router := NewRouter() 35 | router.SetNotFoundHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 36 | notFound = true 37 | })) 38 | err := router.Handle(http.MethodGet, "/a/b", 39 | http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) 40 | assert.Nil(t, err) 41 | r, _ := http.NewRequest(http.MethodGet, "/b/c", nil) 42 | w := new(mockedResponseWriter) 43 | router.ServeHTTP(w, r) 44 | assert.True(t, notFound) 45 | } 46 | 47 | func TestChiRouterNotAllowed(t *testing.T) { 48 | var notAllowed bool 49 | router := NewRouter() 50 | router.SetNotAllowedHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 51 | notAllowed = true 52 | })) 53 | err := router.Handle(http.MethodGet, "/a/b", 54 | http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) 55 | assert.Nil(t, err) 56 | r, _ := http.NewRequest(http.MethodPost, "/a/b", nil) 57 | w := new(mockedResponseWriter) 58 | router.ServeHTTP(w, r) 59 | assert.True(t, notAllowed) 60 | } 61 | 62 | func TestGinRouter(t *testing.T) { 63 | tests := []struct { 64 | method string 65 | path string 66 | expect bool 67 | code int 68 | err error 69 | }{ 70 | // we don't explicitly set status code, framework will do it. 71 | {http.MethodGet, "/test/john/smith/is/super/great", true, 200, nil}, 72 | {http.MethodGet, "/a/b/c?a=b", true, 200, nil}, 73 | {http.MethodGet, "/b/d", false, http.StatusNotFound, nil}, 74 | } 75 | 76 | for _, test := range tests { 77 | t.Run(test.method+":"+test.path, func(t *testing.T) { 78 | routed := false 79 | router := NewRouter() 80 | 81 | err := router.Handle(test.method, "/test/:name/:last_name/*wild", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 82 | routed = true 83 | assert.Equal(t, 3, len(pathvar.Vars(r))) 84 | })) 85 | assert.Nil(t, err) 86 | err = router.Handle(test.method, "/a/b/c", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 87 | routed = true 88 | assert.Nil(t, pathvar.Vars(r)) 89 | })) 90 | assert.Nil(t, err) 91 | err = router.Handle(test.method, "/b/c", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 92 | routed = true 93 | })) 94 | assert.Nil(t, err) 95 | 96 | w := new(mockedResponseWriter) 97 | r, _ := http.NewRequest(test.method, test.path, nil) 98 | router.ServeHTTP(w, r) 99 | 100 | assert.Equal(t, test.expect, routed) 101 | assert.Equal(t, test.code, w.code) 102 | 103 | if test.code == 0 { 104 | r, _ = http.NewRequest(http.MethodPut, test.path, nil) 105 | router.ServeHTTP(w, r) 106 | assert.Equal(t, http.StatusNotFound, w.code) 107 | } 108 | }) 109 | } 110 | } 111 | 112 | func TestParseJsonPost(t *testing.T) { 113 | r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000", 114 | bytes.NewBufferString(`{"location": "shanghai", "time": 20170912}`)) 115 | assert.Nil(t, err) 116 | r.Header.Set(httpx.ContentType, httpx.JsonContentType) 117 | 118 | router := NewRouter() 119 | err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc(func( 120 | w http.ResponseWriter, r *http.Request) { 121 | v := struct { 122 | Name string `path:"name"` 123 | Year int `path:"year"` 124 | Nickname string `form:"nickname"` 125 | Zipcode int64 `form:"zipcode"` 126 | Location string `json:"location"` 127 | Time int64 `json:"time"` 128 | }{} 129 | 130 | err = httpx.Parse(r, &v) 131 | assert.Nil(t, err) 132 | _, err = io.WriteString(w, fmt.Sprintf("%s:%d:%s:%d:%s:%d", v.Name, v.Year, 133 | v.Nickname, v.Zipcode, v.Location, v.Time)) 134 | assert.Nil(t, err) 135 | })) 136 | assert.Nil(t, err) 137 | 138 | rr := httptest.NewRecorder() 139 | router.ServeHTTP(rr, r) 140 | 141 | assert.Equal(t, "kevin:2017:whatever:200000:shanghai:20170912", rr.Body.String()) 142 | } 143 | 144 | func TestParseJsonPostWithIntSlice(t *testing.T) { 145 | r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017", 146 | bytes.NewBufferString(`{"ages": [1, 2], "years": [3, 4]}`)) 147 | assert.Nil(t, err) 148 | r.Header.Set(httpx.ContentType, httpx.JsonContentType) 149 | 150 | router := NewRouter() 151 | err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc(func( 152 | w http.ResponseWriter, r *http.Request) { 153 | v := struct { 154 | Name string `path:"name"` 155 | Year int `path:"year"` 156 | Ages []int `json:"ages"` 157 | Years []int64 `json:"years"` 158 | }{} 159 | 160 | err = httpx.Parse(r, &v) 161 | assert.Nil(t, err) 162 | assert.ElementsMatch(t, []int{1, 2}, v.Ages) 163 | assert.ElementsMatch(t, []int64{3, 4}, v.Years) 164 | })) 165 | assert.Nil(t, err) 166 | 167 | rr := httptest.NewRecorder() 168 | router.ServeHTTP(rr, r) 169 | } 170 | 171 | func TestParseJsonPostError(t *testing.T) { 172 | payload := `[{"abcd": "cdef"}]` 173 | r, err := http.NewRequest(http.MethodPost, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000", 174 | bytes.NewBufferString(payload)) 175 | assert.Nil(t, err) 176 | r.Header.Set(httpx.ContentType, httpx.JsonContentType) 177 | 178 | router := NewRouter() 179 | err = router.Handle(http.MethodPost, "/:name/:year", http.HandlerFunc( 180 | func(w http.ResponseWriter, r *http.Request) { 181 | v := struct { 182 | Name string `path:"name"` 183 | Year int `path:"year"` 184 | Nickname string `form:"nickname"` 185 | Zipcode int64 `form:"zipcode"` 186 | Location string `json:"location"` 187 | Time int64 `json:"time"` 188 | }{} 189 | 190 | err = httpx.Parse(r, &v) 191 | assert.NotNil(t, err) 192 | })) 193 | assert.Nil(t, err) 194 | 195 | rr := httptest.NewRecorder() 196 | router.ServeHTTP(rr, r) 197 | } 198 | 199 | func TestParsePath(t *testing.T) { 200 | r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin/2017", nil) 201 | assert.Nil(t, err) 202 | 203 | router := NewRouter() 204 | err = router.Handle(http.MethodGet, "/:name/:year", http.HandlerFunc( 205 | func(w http.ResponseWriter, r *http.Request) { 206 | v := struct { 207 | Name string `path:"name"` 208 | Year int `path:"year"` 209 | }{} 210 | 211 | err = httpx.Parse(r, &v) 212 | assert.Nil(t, err) 213 | _, err = io.WriteString(w, fmt.Sprintf("%s in %d", v.Name, v.Year)) 214 | assert.Nil(t, err) 215 | })) 216 | assert.Nil(t, err) 217 | 218 | rr := httptest.NewRecorder() 219 | router.ServeHTTP(rr, r) 220 | 221 | assert.Equal(t, "kevin in 2017", rr.Body.String()) 222 | } 223 | 224 | func TestParsePathRequired(t *testing.T) { 225 | r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin", nil) 226 | assert.Nil(t, err) 227 | 228 | router := NewRouter() 229 | err = router.Handle(http.MethodGet, "/:name/", http.HandlerFunc( 230 | func(w http.ResponseWriter, r *http.Request) { 231 | v := struct { 232 | Name string `path:"name"` 233 | Year int `path:"year"` 234 | }{} 235 | 236 | err = httpx.Parse(r, &v) 237 | assert.NotNil(t, err) 238 | })) 239 | assert.Nil(t, err) 240 | 241 | rr := httptest.NewRecorder() 242 | router.ServeHTTP(rr, r) 243 | } 244 | 245 | func TestParseQuery(t *testing.T) { 246 | r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin/2017?nickname=whatever&zipcode=200000", nil) 247 | assert.Nil(t, err) 248 | 249 | router := NewRouter() 250 | err = router.Handle(http.MethodGet, "/:name/:year", http.HandlerFunc( 251 | func(w http.ResponseWriter, r *http.Request) { 252 | v := struct { 253 | Nickname string `form:"nickname"` 254 | Zipcode int64 `form:"zipcode"` 255 | }{} 256 | 257 | err = httpx.Parse(r, &v) 258 | assert.Nil(t, err) 259 | _, err = io.WriteString(w, fmt.Sprintf("%s:%d", v.Nickname, v.Zipcode)) 260 | assert.Nil(t, err) 261 | })) 262 | assert.Nil(t, err) 263 | 264 | rr := httptest.NewRecorder() 265 | router.ServeHTTP(rr, r) 266 | 267 | assert.Equal(t, "whatever:200000", rr.Body.String()) 268 | } 269 | 270 | func TestParseOptional(t *testing.T) { 271 | r, err := http.NewRequest(http.MethodGet, "http://hello.com/kevin/2017?nickname=whatever&zipcode=", nil) 272 | assert.Nil(t, err) 273 | 274 | router := NewRouter() 275 | err = router.Handle(http.MethodGet, "/:name/:year", http.HandlerFunc( 276 | func(w http.ResponseWriter, r *http.Request) { 277 | v := struct { 278 | Nickname string `form:"nickname"` 279 | Zipcode int64 `form:"zipcode,optional"` 280 | }{} 281 | 282 | err = httpx.Parse(r, &v) 283 | assert.Nil(t, err) 284 | _, err = io.WriteString(w, fmt.Sprintf("%s:%d", v.Nickname, v.Zipcode)) 285 | assert.Nil(t, err) 286 | })) 287 | assert.Nil(t, err) 288 | 289 | rr := httptest.NewRecorder() 290 | router.ServeHTTP(rr, r) 291 | 292 | assert.Equal(t, "whatever:0", rr.Body.String()) 293 | } 294 | 295 | func BenchmarkGinRouter(b *testing.B) { 296 | b.ReportAllocs() 297 | 298 | router := NewRouter() 299 | router.Handle(http.MethodGet, "/api/param/:param1/:params2/:param3/:param4/:param5", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 300 | })) 301 | w := &mockedResponseWriter{} 302 | r, _ := http.NewRequest(http.MethodGet, "/api/param/path/to/parameter/john/12345", nil) 303 | for i := 0; i < b.N; i++ { 304 | router.ServeHTTP(w, r) 305 | } 306 | } 307 | -------------------------------------------------------------------------------- /router/mux/README.md: -------------------------------------------------------------------------------- 1 | ### Quick Start 2 | 3 | Prerequesites: 4 | 5 | - Install `go-zero`: 6 | 7 | ```console 8 | go get -u github.com/zeromicro/go-zero 9 | ``` 10 | 11 | Download the module: 12 | 13 | ```console 14 | go get -u github.com/zeromicro/zero-contrib/router/mux 15 | ``` 16 | 17 | For example: 18 | 19 | ```go 20 | package main 21 | 22 | import ( 23 | "encoding/json" 24 | "fmt" 25 | "github.com/zeromicro/go-zero/core/logx" 26 | "github.com/zeromicro/go-zero/core/service" 27 | "github.com/zeromicro/go-zero/rest" 28 | "github.com/zeromicro/go-zero/rest/httpx" 29 | "github.com/zeromicro/zero-contrib/router/mux" 30 | "net/http" 31 | ) 32 | 33 | type User struct { 34 | Id int64 `path:"id"` //`form:"id"` | `json:"id"` 35 | Name string `path:"name"` //`form:"name"` | `json:"name"` 36 | } 37 | 38 | func main() { 39 | r := mux.NewRouter() 40 | engine := rest.MustNewServer(rest.RestConf{ 41 | ServiceConf: service.ServiceConf{ 42 | Log: logx.LogConf{ 43 | Mode: "console", 44 | }, 45 | }, 46 | Port: 3345, 47 | Timeout: 20000, 48 | MaxConns: 500, 49 | }, rest.WithRouter(r)) 50 | 51 | // NotFound defines a handler to respond whenever a route could 52 | // not be found. 53 | r.SetNotFoundHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 54 | w.WriteHeader(404) 55 | w.Write([]byte("404 callback here")) 56 | })) 57 | // MethodNotAllowed defines a handler to respond whenever a method is 58 | // not allowed. 59 | r.SetNotAllowedHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 60 | w.WriteHeader(405) 61 | w.Write([]byte("405 callback here")) 62 | })) 63 | defer engine.Stop() 64 | 65 | engine.AddRoute(rest.Route{ 66 | Method: http.MethodGet, 67 | Path: "/api/{name}/{id}", // GET /api/joh/123 68 | Handler: func(w http.ResponseWriter, r *http.Request) { 69 | var user User 70 | err := httpx.Parse(r, &user) 71 | if err != nil { 72 | return 73 | } 74 | userStr, _ := json.Marshal(user) 75 | fmt.Println(userStr) 76 | w.Write([]byte(userStr)) 77 | }, 78 | }) 79 | engine.Start() 80 | } 81 | ``` 82 | -------------------------------------------------------------------------------- /router/mux/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/router/mux 2 | 3 | go 1.15 4 | 5 | require ( 6 | github.com/gorilla/mux v1.8.0 7 | github.com/stretchr/testify v1.8.4 8 | github.com/zeromicro/go-zero v1.5.6 9 | ) 10 | -------------------------------------------------------------------------------- /router/mux/router.go: -------------------------------------------------------------------------------- 1 | package mux 2 | 3 | import ( 4 | "errors" 5 | "net/http" 6 | "strings" 7 | 8 | "github.com/gorilla/mux" 9 | "github.com/zeromicro/go-zero/rest/httpx" 10 | "github.com/zeromicro/go-zero/rest/pathvar" 11 | ) 12 | 13 | var ( 14 | // ErrInvalidMethod is an error that indicates not a valid http method. 15 | ErrInvalidMethod = errors.New("not a valid http method") 16 | // ErrInvalidPath is an error that indicates path is not start with /. 17 | ErrInvalidPath = errors.New("path must begin with '/'") 18 | ) 19 | 20 | type muxRouter struct { 21 | g *mux.Router 22 | } 23 | 24 | // NewRouter returns a mux.Router. 25 | func NewRouter() httpx.Router { 26 | g := mux.NewRouter() 27 | return &muxRouter{g: g} 28 | } 29 | 30 | func (pr *muxRouter) Handle(method, reqPath string, handler http.Handler) error { 31 | if !validMethod(method) { 32 | return ErrInvalidMethod 33 | } 34 | 35 | if len(reqPath) == 0 || reqPath[0] != '/' { 36 | return ErrInvalidPath 37 | } 38 | 39 | pr.g.HandleFunc(reqPath, func(w http.ResponseWriter, r *http.Request) { 40 | params := make(map[string]string) 41 | vars := mux.Vars(r) 42 | for key, val := range vars { 43 | params[key] = val 44 | } 45 | if len(params) > 0 { 46 | r = pathvar.WithVars(r, params) 47 | } 48 | handler.ServeHTTP(w, r) 49 | }).Methods(strings.ToUpper(method)) 50 | return nil 51 | } 52 | 53 | func (pr *muxRouter) ServeHTTP(w http.ResponseWriter, r *http.Request) { 54 | pr.g.ServeHTTP(w, r) 55 | } 56 | 57 | func (pr *muxRouter) SetNotFoundHandler(handler http.Handler) { 58 | pr.g.NotFoundHandler = handler 59 | } 60 | 61 | func (pr *muxRouter) SetNotAllowedHandler(handler http.Handler) { 62 | pr.g.MethodNotAllowedHandler = handler 63 | } 64 | 65 | func validMethod(method string) bool { 66 | return method == http.MethodDelete || method == http.MethodGet || 67 | method == http.MethodHead || method == http.MethodOptions || 68 | method == http.MethodPatch || method == http.MethodPost || 69 | method == http.MethodPut 70 | } 71 | -------------------------------------------------------------------------------- /router/mux/router_test.go: -------------------------------------------------------------------------------- 1 | package mux 2 | 3 | import ( 4 | "bytes" 5 | "fmt" 6 | "github.com/stretchr/testify/assert" 7 | "github.com/zeromicro/go-zero/rest/httpx" 8 | "github.com/zeromicro/go-zero/rest/pathvar" 9 | "io" 10 | "net/http" 11 | "net/http/httptest" 12 | "testing" 13 | ) 14 | 15 | type mockedResponseWriter struct { 16 | code int 17 | } 18 | 19 | func (m *mockedResponseWriter) Header() http.Header { 20 | return http.Header{} 21 | } 22 | 23 | func (m *mockedResponseWriter) Write(p []byte) (int, error) { 24 | return len(p), nil 25 | } 26 | 27 | func (m *mockedResponseWriter) WriteHeader(code int) { 28 | m.code = code 29 | } 30 | 31 | func TestMuxRouter(t *testing.T) { 32 | tests := []struct { 33 | method string 34 | path string 35 | expect bool 36 | code int 37 | err error 38 | }{ 39 | // we don't explicitly set status code, framework will do it. 40 | {http.MethodGet, "/test/{john}/{smith}", true, 200, nil}, 41 | {http.MethodGet, "/a/b/c?a=b", true, 200, nil}, 42 | {http.MethodGet, "/b/d", false, http.StatusNotFound, nil}, 43 | } 44 | 45 | for _, test := range tests { 46 | t.Run(test.method+":"+test.path, func(t *testing.T) { 47 | routed := false 48 | router := NewRouter() 49 | 50 | err := router.Handle(test.method, "/test/{name}/{last_name}", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 51 | routed = true 52 | assert.Equal(t, 2, len(pathvar.Vars(r))) 53 | w.WriteHeader(200) 54 | })) 55 | assert.Nil(t, err) 56 | err = router.Handle(test.method, "/a/b/c", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 57 | routed = true 58 | assert.Nil(t, pathvar.Vars(r)) 59 | w.WriteHeader(200) 60 | })) 61 | assert.Nil(t, err) 62 | err = router.Handle(test.method, "/b/c", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 63 | routed = true 64 | w.WriteHeader(200) 65 | })) 66 | assert.Nil(t, err) 67 | 68 | w := new(mockedResponseWriter) 69 | r, _ := http.NewRequest(test.method, test.path, nil) 70 | router.ServeHTTP(w, r) 71 | 72 | assert.Equal(t, test.expect, routed) 73 | assert.Equal(t, test.code, w.code) 74 | }) 75 | } 76 | } 77 | 78 | 79 | func TestParseJsonPost(t *testing.T) { 80 | r, err := http.NewRequest(http.MethodPost, "http://hello.com/mikael/2022?nickname=whatever&zipcode=200000", 81 | bytes.NewBufferString(`{"location": "shenzhen", "time": 20220225}`)) 82 | assert.Nil(t, err) 83 | r.Header.Set(httpx.ContentType, httpx.ApplicationJson) 84 | 85 | router := NewRouter() 86 | err = router.Handle(http.MethodPost, "/{name}/{year}", http.HandlerFunc(func( 87 | w http.ResponseWriter, r *http.Request) { 88 | v := struct { 89 | Name string `path:"name"` 90 | Year int `path:"year"` 91 | Nickname string `form:"nickname"` 92 | Zipcode int64 `form:"zipcode"` 93 | Location string `json:"location"` 94 | Time int64 `json:"time"` 95 | }{} 96 | 97 | err = httpx.Parse(r, &v) 98 | assert.Nil(t, err) 99 | _, err = io.WriteString(w, fmt.Sprintf("%s:%d:%s:%d:%s:%d", v.Name, v.Year, 100 | v.Nickname, v.Zipcode, v.Location, v.Time)) 101 | assert.Nil(t, err) 102 | })) 103 | assert.Nil(t, err) 104 | 105 | rr := httptest.NewRecorder() 106 | router.ServeHTTP(rr, r) 107 | 108 | assert.Equal(t, "mikael:2022:whatever:200000:shenzhen:20220225", rr.Body.String()) 109 | } 110 | 111 | func TestParseJsonPostWithIntSlice(t *testing.T) { 112 | r, err := http.NewRequest(http.MethodPost, "http://hello.com/mikael/2022", 113 | bytes.NewBufferString(`{"ages": [1, 2], "years": [3, 4]}`)) 114 | assert.Nil(t, err) 115 | r.Header.Set(httpx.ContentType, httpx.ApplicationJson) 116 | 117 | router := NewRouter() 118 | err = router.Handle(http.MethodPost, "/{name}/{year}", http.HandlerFunc(func( 119 | w http.ResponseWriter, r *http.Request) { 120 | v := struct { 121 | Name string `path:"name"` 122 | Year int `path:"year"` 123 | Ages []int `json:"ages"` 124 | Years []int64 `json:"years"` 125 | }{} 126 | 127 | err = httpx.Parse(r, &v) 128 | 129 | 130 | assert.Nil(t, err) 131 | assert.ElementsMatch(t, []int{1, 2}, v.Ages) 132 | assert.ElementsMatch(t, []int64{3, 4}, v.Years) 133 | })) 134 | assert.Nil(t, err) 135 | 136 | rr := httptest.NewRecorder() 137 | router.ServeHTTP(rr, r) 138 | } 139 | 140 | func TestParseJsonPostError(t *testing.T) { 141 | payload := `[{"abcd": "cdef"}]` 142 | r, err := http.NewRequest(http.MethodPost, "http://hello.com/mikael/2022?nickname=whatever&zipcode=200000", 143 | bytes.NewBufferString(payload)) 144 | assert.Nil(t, err) 145 | r.Header.Set(httpx.ContentType, httpx.ApplicationJson) 146 | 147 | router := NewRouter() 148 | err = router.Handle(http.MethodPost, "/{name}/{year}", http.HandlerFunc( 149 | func(w http.ResponseWriter, r *http.Request) { 150 | v := struct { 151 | Name string `path:"name"` 152 | Year int `path:"year"` 153 | Nickname string `form:"nickname"` 154 | Zipcode int64 `form:"zipcode"` 155 | Location string `json:"location"` 156 | Time int64 `json:"time"` 157 | }{} 158 | 159 | err = httpx.Parse(r, &v) 160 | assert.NotNil(t, err) 161 | })) 162 | assert.Nil(t, err) 163 | 164 | rr := httptest.NewRecorder() 165 | router.ServeHTTP(rr, r) 166 | } 167 | 168 | func TestParsePath(t *testing.T) { 169 | r, err := http.NewRequest(http.MethodGet, "http://hello.com/mikael/2022", nil) 170 | assert.Nil(t, err) 171 | 172 | router := NewRouter() 173 | err = router.Handle(http.MethodGet, "/{name}/{year}", http.HandlerFunc( 174 | func(w http.ResponseWriter, r *http.Request) { 175 | v := struct { 176 | Name string `path:"name"` 177 | Year int `path:"year"` 178 | }{} 179 | 180 | err = httpx.Parse(r, &v) 181 | assert.Nil(t, err) 182 | _, err = io.WriteString(w, fmt.Sprintf("%s in %d", v.Name, v.Year)) 183 | assert.Nil(t, err) 184 | })) 185 | assert.Nil(t, err) 186 | 187 | rr := httptest.NewRecorder() 188 | router.ServeHTTP(rr, r) 189 | 190 | assert.Equal(t, "mikael in 2022", rr.Body.String()) 191 | } 192 | 193 | func TestParsePathRequired(t *testing.T) { 194 | r, err := http.NewRequest(http.MethodGet, "http://hello.com/mikael", nil) 195 | assert.Nil(t, err) 196 | 197 | router := NewRouter() 198 | err = router.Handle(http.MethodGet, "/{name}/", http.HandlerFunc( 199 | func(w http.ResponseWriter, r *http.Request) { 200 | v := struct { 201 | Name string `path:"name"` 202 | Year int `path:"year"` 203 | }{} 204 | 205 | err = httpx.Parse(r, &v) 206 | assert.NotNil(t, err) 207 | })) 208 | assert.Nil(t, err) 209 | 210 | rr := httptest.NewRecorder() 211 | router.ServeHTTP(rr, r) 212 | } 213 | 214 | func TestParseQuery(t *testing.T) { 215 | r, err := http.NewRequest(http.MethodGet, "http://hello.com/mikael/2022?nickname=whatever&zipcode=200000", nil) 216 | assert.Nil(t, err) 217 | 218 | router := NewRouter() 219 | err = router.Handle(http.MethodGet, "/{name}/{year}", http.HandlerFunc( 220 | func(w http.ResponseWriter, r *http.Request) { 221 | v := struct { 222 | Nickname string `form:"nickname"` 223 | Zipcode int64 `form:"zipcode"` 224 | }{} 225 | 226 | err = httpx.Parse(r, &v) 227 | assert.Nil(t, err) 228 | _, err = io.WriteString(w, fmt.Sprintf("%s:%d", v.Nickname, v.Zipcode)) 229 | assert.Nil(t, err) 230 | })) 231 | assert.Nil(t, err) 232 | 233 | rr := httptest.NewRecorder() 234 | router.ServeHTTP(rr, r) 235 | 236 | assert.Equal(t, "whatever:200000", rr.Body.String()) 237 | } 238 | 239 | func TestParseOptional(t *testing.T) { 240 | r, err := http.NewRequest(http.MethodGet, "http://hello.com/mikael/2022?nickname=whatever&zipcode=", nil) 241 | assert.Nil(t, err) 242 | 243 | router := NewRouter() 244 | err = router.Handle(http.MethodGet, "/{name}/{year}", http.HandlerFunc( 245 | func(w http.ResponseWriter, r *http.Request) { 246 | v := struct { 247 | Nickname string `form:"nickname"` 248 | Zipcode int64 `form:"zipcode,optional"` 249 | }{} 250 | 251 | err = httpx.Parse(r, &v) 252 | assert.Nil(t, err) 253 | _, err = io.WriteString(w, fmt.Sprintf("%s:%d", v.Nickname, v.Zipcode)) 254 | assert.Nil(t, err) 255 | })) 256 | assert.Nil(t, err) 257 | 258 | rr := httptest.NewRecorder() 259 | router.ServeHTTP(rr, r) 260 | 261 | assert.Equal(t, "whatever:0", rr.Body.String()) 262 | } 263 | 264 | 265 | func BenchmarkMuxRouter(b *testing.B) { 266 | b.ReportAllocs() 267 | 268 | router := NewRouter() 269 | router.Handle(http.MethodGet, "/api/param/{param1}/{params2}/{param3}/{param4}/{param5}", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 270 | })) 271 | w := &mockedResponseWriter{} 272 | r, _ := http.NewRequest(http.MethodGet, "/api/param/path/to/parameter/john/12345", nil) 273 | for i := 0; i < b.N; i++ { 274 | router.ServeHTTP(w, r) 275 | } 276 | } 277 | -------------------------------------------------------------------------------- /stores/clickhouse/clickhouse.go: -------------------------------------------------------------------------------- 1 | package clickhouse 2 | 3 | import ( 4 | // imports the driver, don't remove this comment, golint requires. 5 | _ "github.com/ClickHouse/clickhouse-go/v2" 6 | "github.com/zeromicro/go-zero/core/stores/sqlx" 7 | ) 8 | 9 | const clickHouseDriverName = "clickhouse" 10 | 11 | // New returns a clickhouse connection. 12 | func New(datasource string, opts ...sqlx.SqlOption) sqlx.SqlConn { 13 | return sqlx.NewSqlConn(clickHouseDriverName, datasource, opts...) 14 | } 15 | -------------------------------------------------------------------------------- /stores/clickhouse/clickhouse_test.go: -------------------------------------------------------------------------------- 1 | package clickhouse 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestClickHouse(t *testing.T) { 10 | assert.NotNil(t, New("clickhouse")) 11 | } 12 | -------------------------------------------------------------------------------- /stores/clickhouse/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/stores/clickhouse 2 | 3 | go 1.16 4 | 5 | require ( 6 | github.com/ClickHouse/clickhouse-go/v2 v2.14.1 7 | github.com/pierrec/lz4 v2.6.1+incompatible // indirect 8 | github.com/stretchr/testify v1.8.4 9 | github.com/zeromicro/go-zero v1.5.4 10 | ) 11 | -------------------------------------------------------------------------------- /stores/mongo/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/stores/mongo 2 | 3 | go 1.21 4 | 5 | require ( 6 | github.com/alicebob/miniredis/v2 v2.35.0 7 | github.com/stretchr/testify v1.10.0 8 | github.com/zeromicro/go-zero v1.8.5 9 | go.mongodb.org/mongo-driver v1.17.4 10 | go.opentelemetry.io/otel v1.24.0 11 | go.opentelemetry.io/otel/trace v1.24.0 12 | ) 13 | 14 | require ( 15 | github.com/beorn7/perks v1.0.1 // indirect 16 | github.com/cenkalti/backoff/v4 v4.3.0 // indirect 17 | github.com/cespare/xxhash/v2 v2.3.0 // indirect 18 | github.com/davecgh/go-spew v1.1.1 // indirect 19 | github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect 20 | github.com/fatih/color v1.18.0 // indirect 21 | github.com/go-logr/logr v1.4.2 // indirect 22 | github.com/go-logr/stdr v1.2.2 // indirect 23 | github.com/golang/snappy v0.0.4 // indirect 24 | github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0 // indirect 25 | github.com/klauspost/compress v1.17.11 // indirect 26 | github.com/mattn/go-colorable v0.1.13 // indirect 27 | github.com/mattn/go-isatty v0.0.20 // indirect 28 | github.com/montanaflynn/stats v0.7.1 // indirect 29 | github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect 30 | github.com/openzipkin/zipkin-go v0.4.3 // indirect 31 | github.com/pelletier/go-toml/v2 v2.2.2 // indirect 32 | github.com/pmezard/go-difflib v1.0.0 // indirect 33 | github.com/prometheus/client_golang v1.21.1 // indirect 34 | github.com/prometheus/client_model v0.6.1 // indirect 35 | github.com/prometheus/common v0.62.0 // indirect 36 | github.com/prometheus/procfs v0.15.1 // indirect 37 | github.com/redis/go-redis/v9 v9.11.0 // indirect 38 | github.com/spaolacci/murmur3 v1.1.0 // indirect 39 | github.com/xdg-go/pbkdf2 v1.0.0 // indirect 40 | github.com/xdg-go/scram v1.1.2 // indirect 41 | github.com/xdg-go/stringprep v1.0.4 // indirect 42 | github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 // indirect 43 | github.com/yuin/gopher-lua v1.1.1 // indirect 44 | go.opentelemetry.io/otel/exporters/jaeger v1.17.0 // indirect 45 | go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.24.0 // indirect 46 | go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.24.0 // indirect 47 | go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.24.0 // indirect 48 | go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.24.0 // indirect 49 | go.opentelemetry.io/otel/exporters/zipkin v1.24.0 // indirect 50 | go.opentelemetry.io/otel/metric v1.24.0 // indirect 51 | go.opentelemetry.io/otel/sdk v1.24.0 // indirect 52 | go.opentelemetry.io/proto/otlp v1.3.1 // indirect 53 | go.uber.org/automaxprocs v1.6.0 // indirect 54 | golang.org/x/crypto v0.33.0 // indirect 55 | golang.org/x/net v0.35.0 // indirect 56 | golang.org/x/sync v0.11.0 // indirect 57 | golang.org/x/sys v0.30.0 // indirect 58 | golang.org/x/text v0.22.0 // indirect 59 | google.golang.org/genproto/googleapis/api v0.0.0-20240711142825-46eb208f015d // indirect 60 | google.golang.org/genproto/googleapis/rpc v0.0.0-20240701130421-f6361c86f094 // indirect 61 | google.golang.org/grpc v1.65.0 // indirect 62 | google.golang.org/protobuf v1.36.5 // indirect 63 | gopkg.in/yaml.v2 v2.4.0 // indirect 64 | gopkg.in/yaml.v3 v3.0.1 // indirect 65 | ) 66 | -------------------------------------------------------------------------------- /stores/mongo/mon/bulkinserter.go: -------------------------------------------------------------------------------- 1 | package mon 2 | 3 | import ( 4 | "context" 5 | "time" 6 | 7 | "github.com/zeromicro/go-zero/core/executors" 8 | "github.com/zeromicro/go-zero/core/logx" 9 | "go.mongodb.org/mongo-driver/mongo" 10 | ) 11 | 12 | const ( 13 | flushInterval = time.Second 14 | maxBulkRows = 1000 15 | ) 16 | 17 | type ( 18 | // ResultHandler is a handler that used to handle results. 19 | ResultHandler func(*mongo.InsertManyResult, error) 20 | 21 | // A BulkInserter is used to insert bulk of mongo records. 22 | BulkInserter struct { 23 | executor *executors.PeriodicalExecutor 24 | inserter *dbInserter 25 | } 26 | ) 27 | 28 | // NewBulkInserter returns a BulkInserter. 29 | func NewBulkInserter(coll Collection, interval ...time.Duration) (*BulkInserter, error) { 30 | cloneColl, err := coll.Clone() 31 | if err != nil { 32 | return nil, err 33 | } 34 | 35 | inserter := &dbInserter{ 36 | collection: cloneColl, 37 | } 38 | 39 | duration := flushInterval 40 | if len(interval) > 0 { 41 | duration = interval[0] 42 | } 43 | 44 | return &BulkInserter{ 45 | executor: executors.NewPeriodicalExecutor(duration, inserter), 46 | inserter: inserter, 47 | }, nil 48 | } 49 | 50 | // Flush flushes the inserter, writes all pending records. 51 | func (bi *BulkInserter) Flush() { 52 | bi.executor.Flush() 53 | } 54 | 55 | // Insert inserts doc. 56 | func (bi *BulkInserter) Insert(doc any) { 57 | bi.executor.Add(doc) 58 | } 59 | 60 | // SetResultHandler sets the result handler. 61 | func (bi *BulkInserter) SetResultHandler(handler ResultHandler) { 62 | bi.executor.Sync(func() { 63 | bi.inserter.resultHandler = handler 64 | }) 65 | } 66 | 67 | type dbInserter struct { 68 | collection *mongo.Collection 69 | documents []any 70 | resultHandler ResultHandler 71 | } 72 | 73 | func (in *dbInserter) AddTask(doc any) bool { 74 | in.documents = append(in.documents, doc) 75 | return len(in.documents) >= maxBulkRows 76 | } 77 | 78 | func (in *dbInserter) Execute(objs any) { 79 | docs := objs.([]any) 80 | if len(docs) == 0 { 81 | return 82 | } 83 | 84 | result, err := in.collection.InsertMany(context.Background(), docs) 85 | if in.resultHandler != nil { 86 | in.resultHandler(result, err) 87 | } else if err != nil { 88 | logx.Error(err) 89 | } 90 | } 91 | 92 | func (in *dbInserter) RemoveAll() any { 93 | documents := in.documents 94 | in.documents = nil 95 | return documents 96 | } 97 | -------------------------------------------------------------------------------- /stores/mongo/mon/bulkinserter_test.go: -------------------------------------------------------------------------------- 1 | package mon 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | "go.mongodb.org/mongo-driver/bson" 8 | "go.mongodb.org/mongo-driver/mongo" 9 | "go.mongodb.org/mongo-driver/mongo/integration/mtest" 10 | ) 11 | 12 | func TestBulkInserter(t *testing.T) { 13 | mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock)) 14 | mt.Run("test", func(mt *mtest.T) { 15 | mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "ok", Value: 1}}...)) 16 | bulk, err := NewBulkInserter(createModel(mt).Collection) 17 | assert.Equal(t, err, nil) 18 | bulk.SetResultHandler(func(result *mongo.InsertManyResult, err error) { 19 | assert.Nil(t, err) 20 | assert.Equal(t, 2, len(result.InsertedIDs)) 21 | }) 22 | bulk.Insert(bson.D{{Key: "foo", Value: "bar"}}) 23 | bulk.Insert(bson.D{{Key: "foo", Value: "baz"}}) 24 | bulk.Flush() 25 | }) 26 | } 27 | -------------------------------------------------------------------------------- /stores/mongo/mon/clientmanager.go: -------------------------------------------------------------------------------- 1 | package mon 2 | 3 | import ( 4 | "context" 5 | "io" 6 | 7 | "github.com/zeromicro/go-zero/core/syncx" 8 | "go.mongodb.org/mongo-driver/mongo" 9 | mopt "go.mongodb.org/mongo-driver/mongo/options" 10 | ) 11 | 12 | var clientManager = syncx.NewResourceManager() 13 | 14 | // ClosableClient wraps *mongo.Client and provides a Close method. 15 | type ClosableClient struct { 16 | *mongo.Client 17 | } 18 | 19 | // Close disconnects the underlying *mongo.Client. 20 | func (cs *ClosableClient) Close() error { 21 | return cs.Client.Disconnect(context.Background()) 22 | } 23 | 24 | // Inject injects a *mongo.Client into the client manager. 25 | // Typically, this is used to inject a *mongo.Client for test purpose. 26 | func Inject(key string, client *mongo.Client) { 27 | clientManager.Inject(key, &ClosableClient{client}) 28 | } 29 | 30 | func getClient(url string, opts ...Option) (*mongo.Client, error) { 31 | val, err := clientManager.GetResource(url, func() (io.Closer, error) { 32 | o := mopt.Client().ApplyURI(url) 33 | opts = append([]Option{defaultTimeoutOption()}, opts...) 34 | for _, opt := range opts { 35 | opt(o) 36 | } 37 | 38 | cli, err := mongo.Connect(context.Background(), o) 39 | if err != nil { 40 | return nil, err 41 | } 42 | 43 | err = cli.Ping(context.Background(), nil) 44 | if err != nil { 45 | return nil, err 46 | } 47 | 48 | concurrentSess := &ClosableClient{ 49 | Client: cli, 50 | } 51 | 52 | return concurrentSess, nil 53 | }) 54 | if err != nil { 55 | return nil, err 56 | } 57 | 58 | return val.(*ClosableClient).Client, nil 59 | } 60 | -------------------------------------------------------------------------------- /stores/mongo/mon/clientmanager_test.go: -------------------------------------------------------------------------------- 1 | package mon 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | "go.mongodb.org/mongo-driver/mongo/integration/mtest" 8 | ) 9 | 10 | func init() { 11 | _ = mtest.Setup() 12 | } 13 | 14 | func TestClientManger_getClient(t *testing.T) { 15 | mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock)) 16 | mt.Run("test", func(mt *mtest.T) { 17 | Inject(mtest.ClusterURI(), mt.Client) 18 | cli, err := getClient(mtest.ClusterURI()) 19 | assert.Nil(t, err) 20 | assert.Equal(t, mt.Client, cli) 21 | }) 22 | } 23 | -------------------------------------------------------------------------------- /stores/mongo/mon/model.go: -------------------------------------------------------------------------------- 1 | package mon 2 | 3 | import ( 4 | "context" 5 | "strings" 6 | 7 | "github.com/zeromicro/go-zero/core/breaker" 8 | "github.com/zeromicro/go-zero/core/logx" 9 | "github.com/zeromicro/go-zero/core/timex" 10 | "go.mongodb.org/mongo-driver/mongo" 11 | mopt "go.mongodb.org/mongo-driver/mongo/options" 12 | ) 13 | 14 | const ( 15 | startSession = "StartSession" 16 | abortTransaction = "AbortTransaction" 17 | commitTransaction = "CommitTransaction" 18 | withTransaction = "WithTransaction" 19 | endSession = "EndSession" 20 | ) 21 | 22 | type ( 23 | // Model is a mongodb store model that represents a collection. 24 | Model struct { 25 | Collection 26 | name string 27 | cli *mongo.Client 28 | brk breaker.Breaker 29 | opts []Option 30 | } 31 | 32 | wrappedSession struct { 33 | mongo.Session 34 | name string 35 | brk breaker.Breaker 36 | } 37 | ) 38 | 39 | // MustNewModel returns a Model, exits on errors. 40 | func MustNewModel(uri, db, collection string, opts ...Option) *Model { 41 | model, err := NewModel(uri, db, collection, opts...) 42 | logx.Must(err) 43 | return model 44 | } 45 | 46 | // NewModel returns a Model. 47 | func NewModel(uri, db, collection string, opts ...Option) (*Model, error) { 48 | cli, err := getClient(uri, opts...) 49 | if err != nil { 50 | return nil, err 51 | } 52 | 53 | name := strings.Join([]string{uri, collection}, "/") 54 | brk := breaker.GetBreaker(uri) 55 | coll := newCollection(cli.Database(db).Collection(collection), brk) 56 | return newModel(name, cli, coll, brk, opts...), nil 57 | } 58 | 59 | func newModel(name string, cli *mongo.Client, coll Collection, brk breaker.Breaker, 60 | opts ...Option) *Model { 61 | return &Model{ 62 | name: name, 63 | Collection: coll, 64 | cli: cli, 65 | brk: brk, 66 | opts: opts, 67 | } 68 | } 69 | 70 | // StartSession starts a new session. 71 | func (m *Model) StartSession(opts ...*mopt.SessionOptions) (sess mongo.Session, err error) { 72 | starTime := timex.Now() 73 | defer func() { 74 | logDuration(context.Background(), m.name, startSession, starTime, err) 75 | }() 76 | 77 | session, sessionErr := m.cli.StartSession(opts...) 78 | if sessionErr != nil { 79 | return nil, sessionErr 80 | } 81 | 82 | return &wrappedSession{ 83 | Session: session, 84 | name: m.name, 85 | brk: m.brk, 86 | }, nil 87 | } 88 | 89 | // Aggregate executes an aggregation pipeline. 90 | func (m *Model) Aggregate(ctx context.Context, v, pipeline any, opts ...*mopt.AggregateOptions) error { 91 | cur, err := m.Collection.Aggregate(ctx, pipeline, opts...) 92 | if err != nil { 93 | return err 94 | } 95 | defer cur.Close(ctx) 96 | 97 | return cur.All(ctx, v) 98 | } 99 | 100 | // DeleteMany deletes documents that match the filter. 101 | func (m *Model) DeleteMany(ctx context.Context, filter any, opts ...*mopt.DeleteOptions) (int64, error) { 102 | res, err := m.Collection.DeleteMany(ctx, filter, opts...) 103 | if err != nil { 104 | return 0, err 105 | } 106 | 107 | return res.DeletedCount, nil 108 | } 109 | 110 | // DeleteOne deletes the first document that matches the filter. 111 | func (m *Model) DeleteOne(ctx context.Context, filter any, opts ...*mopt.DeleteOptions) (int64, error) { 112 | res, err := m.Collection.DeleteOne(ctx, filter, opts...) 113 | if err != nil { 114 | return 0, err 115 | } 116 | 117 | return res.DeletedCount, nil 118 | } 119 | 120 | // Find finds documents that match the filter. 121 | func (m *Model) Find(ctx context.Context, v, filter any, opts ...*mopt.FindOptions) error { 122 | cur, err := m.Collection.Find(ctx, filter, opts...) 123 | if err != nil { 124 | return err 125 | } 126 | defer cur.Close(ctx) 127 | 128 | return cur.All(ctx, v) 129 | } 130 | 131 | // FindOne finds the first document that matches the filter. 132 | func (m *Model) FindOne(ctx context.Context, v, filter any, opts ...*mopt.FindOneOptions) error { 133 | res, err := m.Collection.FindOne(ctx, filter, opts...) 134 | if err != nil { 135 | return err 136 | } 137 | 138 | return res.Decode(v) 139 | } 140 | 141 | // FindOneAndDelete finds a single document and deletes it. 142 | func (m *Model) FindOneAndDelete(ctx context.Context, v, filter any, 143 | opts ...*mopt.FindOneAndDeleteOptions) error { 144 | res, err := m.Collection.FindOneAndDelete(ctx, filter, opts...) 145 | if err != nil { 146 | return err 147 | } 148 | 149 | return res.Decode(v) 150 | } 151 | 152 | // FindOneAndReplace finds a single document and replaces it. 153 | func (m *Model) FindOneAndReplace(ctx context.Context, v, filter, replacement any, 154 | opts ...*mopt.FindOneAndReplaceOptions) error { 155 | res, err := m.Collection.FindOneAndReplace(ctx, filter, replacement, opts...) 156 | if err != nil { 157 | return err 158 | } 159 | 160 | return res.Decode(v) 161 | } 162 | 163 | // FindOneAndUpdate finds a single document and updates it. 164 | func (m *Model) FindOneAndUpdate(ctx context.Context, v, filter, update any, 165 | opts ...*mopt.FindOneAndUpdateOptions) error { 166 | res, err := m.Collection.FindOneAndUpdate(ctx, filter, update, opts...) 167 | if err != nil { 168 | return err 169 | } 170 | 171 | return res.Decode(v) 172 | } 173 | 174 | // AbortTransaction implements the mongo.Session interface. 175 | func (w *wrappedSession) AbortTransaction(ctx context.Context) (err error) { 176 | ctx, span := startSpan(ctx, abortTransaction) 177 | defer func() { 178 | endSpan(span, err) 179 | }() 180 | 181 | return w.brk.DoWithAcceptableCtx(ctx, func() error { 182 | starTime := timex.Now() 183 | defer func() { 184 | logDuration(ctx, w.name, abortTransaction, starTime, err) 185 | }() 186 | 187 | return w.Session.AbortTransaction(ctx) 188 | }, acceptable) 189 | } 190 | 191 | // CommitTransaction implements the mongo.Session interface. 192 | func (w *wrappedSession) CommitTransaction(ctx context.Context) (err error) { 193 | ctx, span := startSpan(ctx, commitTransaction) 194 | defer func() { 195 | endSpan(span, err) 196 | }() 197 | 198 | return w.brk.DoWithAcceptableCtx(ctx, func() error { 199 | starTime := timex.Now() 200 | defer func() { 201 | logDuration(ctx, w.name, commitTransaction, starTime, err) 202 | }() 203 | 204 | return w.Session.CommitTransaction(ctx) 205 | }, acceptable) 206 | } 207 | 208 | // WithTransaction implements the mongo.Session interface. 209 | func (w *wrappedSession) WithTransaction( 210 | ctx context.Context, 211 | fn func(sessCtx mongo.SessionContext) (any, error), 212 | opts ...*mopt.TransactionOptions, 213 | ) (res any, err error) { 214 | ctx, span := startSpan(ctx, withTransaction) 215 | defer func() { 216 | endSpan(span, err) 217 | }() 218 | 219 | err = w.brk.DoWithAcceptableCtx(ctx, func() error { 220 | starTime := timex.Now() 221 | defer func() { 222 | logDuration(ctx, w.name, withTransaction, starTime, err) 223 | }() 224 | 225 | res, err = w.Session.WithTransaction(ctx, fn, opts...) 226 | return err 227 | }, acceptable) 228 | 229 | return 230 | } 231 | 232 | // EndSession implements the mongo.Session interface. 233 | func (w *wrappedSession) EndSession(ctx context.Context) { 234 | var err error 235 | ctx, span := startSpan(ctx, endSession) 236 | defer func() { 237 | endSpan(span, err) 238 | }() 239 | 240 | err = w.brk.DoWithAcceptableCtx(ctx, func() error { 241 | starTime := timex.Now() 242 | defer func() { 243 | logDuration(ctx, w.name, endSession, starTime, err) 244 | }() 245 | 246 | w.Session.EndSession(ctx) 247 | return nil 248 | }, acceptable) 249 | } 250 | -------------------------------------------------------------------------------- /stores/mongo/mon/model_test.go: -------------------------------------------------------------------------------- 1 | package mon 2 | 3 | import ( 4 | "context" 5 | "testing" 6 | 7 | "github.com/stretchr/testify/assert" 8 | "go.mongodb.org/mongo-driver/bson" 9 | "go.mongodb.org/mongo-driver/mongo" 10 | "go.mongodb.org/mongo-driver/mongo/integration/mtest" 11 | ) 12 | 13 | func TestModel_StartSession(t *testing.T) { 14 | mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock)) 15 | mt.Run("test", func(mt *mtest.T) { 16 | m := createModel(mt) 17 | sess, err := m.StartSession() 18 | assert.Nil(t, err) 19 | defer sess.EndSession(context.Background()) 20 | 21 | _, err = sess.WithTransaction(context.Background(), func(sessCtx mongo.SessionContext) (any, error) { 22 | _ = sessCtx.StartTransaction() 23 | sessCtx.Client().Database("1") 24 | sessCtx.EndSession(context.Background()) 25 | return nil, nil 26 | }) 27 | assert.Nil(t, err) 28 | assert.NoError(t, sess.CommitTransaction(context.Background())) 29 | assert.Error(t, sess.AbortTransaction(context.Background())) 30 | }) 31 | } 32 | 33 | func TestModel_Aggregate(t *testing.T) { 34 | mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock)) 35 | mt.Run("test", func(mt *mtest.T) { 36 | m := createModel(mt) 37 | find := mtest.CreateCursorResponse( 38 | 1, 39 | "DBName.CollectionName", 40 | mtest.FirstBatch, 41 | bson.D{ 42 | {Key: "name", Value: "John"}, 43 | }) 44 | getMore := mtest.CreateCursorResponse( 45 | 1, 46 | "DBName.CollectionName", 47 | mtest.NextBatch, 48 | bson.D{ 49 | {Key: "name", Value: "Mary"}, 50 | }) 51 | killCursors := mtest.CreateCursorResponse( 52 | 0, 53 | "DBName.CollectionName", 54 | mtest.NextBatch) 55 | mt.AddMockResponses(find, getMore, killCursors) 56 | var result []any 57 | err := m.Aggregate(context.Background(), &result, mongo.Pipeline{}) 58 | assert.Nil(t, err) 59 | assert.Equal(t, 2, len(result)) 60 | assert.Equal(t, "John", result[0].(bson.D).Map()["name"]) 61 | assert.Equal(t, "Mary", result[1].(bson.D).Map()["name"]) 62 | 63 | triggerBreaker(m) 64 | assert.Equal(t, errDummy, m.Aggregate(context.Background(), &result, mongo.Pipeline{})) 65 | }) 66 | } 67 | 68 | func TestModel_DeleteMany(t *testing.T) { 69 | mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock)) 70 | mt.Run("test", func(mt *mtest.T) { 71 | m := createModel(mt) 72 | mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...)) 73 | val, err := m.DeleteMany(context.Background(), bson.D{}) 74 | assert.Nil(t, err) 75 | assert.Equal(t, int64(1), val) 76 | 77 | triggerBreaker(m) 78 | _, err = m.DeleteMany(context.Background(), bson.D{}) 79 | assert.Equal(t, errDummy, err) 80 | }) 81 | } 82 | 83 | func TestModel_DeleteOne(t *testing.T) { 84 | mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock)) 85 | mt.Run("test", func(mt *mtest.T) { 86 | m := createModel(mt) 87 | mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{{Key: "n", Value: 1}}...)) 88 | val, err := m.DeleteOne(context.Background(), bson.D{}) 89 | assert.Nil(t, err) 90 | assert.Equal(t, int64(1), val) 91 | 92 | triggerBreaker(m) 93 | _, err = m.DeleteOne(context.Background(), bson.D{}) 94 | assert.Equal(t, errDummy, err) 95 | }) 96 | } 97 | 98 | func TestModel_Find(t *testing.T) { 99 | mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock)) 100 | mt.Run("test", func(mt *mtest.T) { 101 | m := createModel(mt) 102 | find := mtest.CreateCursorResponse( 103 | 1, 104 | "DBName.CollectionName", 105 | mtest.FirstBatch, 106 | bson.D{ 107 | {Key: "name", Value: "John"}, 108 | }) 109 | getMore := mtest.CreateCursorResponse( 110 | 1, 111 | "DBName.CollectionName", 112 | mtest.NextBatch, 113 | bson.D{ 114 | {Key: "name", Value: "Mary"}, 115 | }) 116 | killCursors := mtest.CreateCursorResponse( 117 | 0, 118 | "DBName.CollectionName", 119 | mtest.NextBatch) 120 | mt.AddMockResponses(find, getMore, killCursors) 121 | var result []any 122 | err := m.Find(context.Background(), &result, bson.D{}) 123 | assert.Nil(t, err) 124 | assert.Equal(t, 2, len(result)) 125 | assert.Equal(t, "John", result[0].(bson.D).Map()["name"]) 126 | assert.Equal(t, "Mary", result[1].(bson.D).Map()["name"]) 127 | 128 | triggerBreaker(m) 129 | assert.Equal(t, errDummy, m.Find(context.Background(), &result, bson.D{})) 130 | }) 131 | } 132 | 133 | func TestModel_FindOne(t *testing.T) { 134 | mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock)) 135 | mt.Run("test", func(mt *mtest.T) { 136 | m := createModel(mt) 137 | find := mtest.CreateCursorResponse( 138 | 1, 139 | "DBName.CollectionName", 140 | mtest.FirstBatch, 141 | bson.D{ 142 | {Key: "name", Value: "John"}, 143 | }) 144 | killCursors := mtest.CreateCursorResponse( 145 | 0, 146 | "DBName.CollectionName", 147 | mtest.NextBatch) 148 | mt.AddMockResponses(find, killCursors) 149 | var result bson.D 150 | err := m.FindOne(context.Background(), &result, bson.D{}) 151 | assert.Nil(t, err) 152 | assert.Equal(t, "John", result.Map()["name"]) 153 | 154 | triggerBreaker(m) 155 | assert.Equal(t, errDummy, m.FindOne(context.Background(), &result, bson.D{})) 156 | }) 157 | } 158 | 159 | func TestModel_FindOneAndDelete(t *testing.T) { 160 | mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock)) 161 | mt.Run("test", func(mt *mtest.T) { 162 | m := createModel(mt) 163 | mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{ 164 | {Key: "value", Value: bson.D{{Key: "name", Value: "John"}}}, 165 | }...)) 166 | var result bson.D 167 | err := m.FindOneAndDelete(context.Background(), &result, bson.D{}) 168 | assert.Nil(t, err) 169 | assert.Equal(t, "John", result.Map()["name"]) 170 | 171 | triggerBreaker(m) 172 | assert.Equal(t, errDummy, m.FindOneAndDelete(context.Background(), &result, bson.D{})) 173 | }) 174 | } 175 | 176 | func TestModel_FindOneAndReplace(t *testing.T) { 177 | mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock)) 178 | mt.Run("test", func(mt *mtest.T) { 179 | m := createModel(mt) 180 | mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{ 181 | {Key: "value", Value: bson.D{{Key: "name", Value: "John"}}}, 182 | }...)) 183 | var result bson.D 184 | err := m.FindOneAndReplace(context.Background(), &result, bson.D{}, bson.D{ 185 | {Key: "name", Value: "Mary"}, 186 | }) 187 | assert.Nil(t, err) 188 | assert.Equal(t, "John", result.Map()["name"]) 189 | 190 | triggerBreaker(m) 191 | assert.Equal(t, errDummy, m.FindOneAndReplace(context.Background(), &result, bson.D{}, bson.D{ 192 | {Key: "name", Value: "Mary"}, 193 | })) 194 | }) 195 | } 196 | 197 | func TestModel_FindOneAndUpdate(t *testing.T) { 198 | mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock)) 199 | mt.Run("test", func(mt *mtest.T) { 200 | m := createModel(mt) 201 | mt.AddMockResponses(mtest.CreateSuccessResponse(bson.D{ 202 | {Key: "value", Value: bson.D{{Key: "name", Value: "John"}}}, 203 | }...)) 204 | var result bson.D 205 | err := m.FindOneAndUpdate(context.Background(), &result, bson.D{}, bson.D{ 206 | {Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}}, 207 | }) 208 | assert.Nil(t, err) 209 | assert.Equal(t, "John", result.Map()["name"]) 210 | 211 | triggerBreaker(m) 212 | assert.Equal(t, errDummy, m.FindOneAndUpdate(context.Background(), &result, bson.D{}, bson.D{ 213 | {Key: "$set", Value: bson.D{{Key: "name", Value: "Mary"}}}, 214 | })) 215 | }) 216 | } 217 | 218 | func createModel(mt *mtest.T) *Model { 219 | Inject(mt.Name(), mt.Client) 220 | return MustNewModel(mt.Name(), mt.DB.Name(), mt.Coll.Name()) 221 | } 222 | 223 | func triggerBreaker(m *Model) { 224 | m.Collection.(*decoratedCollection).brk = new(dropBreaker) 225 | } 226 | -------------------------------------------------------------------------------- /stores/mongo/mon/options.go: -------------------------------------------------------------------------------- 1 | package mon 2 | 3 | import ( 4 | "reflect" 5 | "time" 6 | 7 | "github.com/zeromicro/go-zero/core/syncx" 8 | "go.mongodb.org/mongo-driver/bson" 9 | "go.mongodb.org/mongo-driver/bson/bsoncodec" 10 | mopt "go.mongodb.org/mongo-driver/mongo/options" 11 | ) 12 | 13 | const defaultTimeout = time.Second * 3 14 | 15 | var ( 16 | slowThreshold = syncx.ForAtomicDuration(defaultSlowThreshold) 17 | logMon = syncx.ForAtomicBool(true) 18 | logSlowMon = syncx.ForAtomicBool(true) 19 | ) 20 | 21 | type ( 22 | // Option defines the method to customize a mongo model. 23 | Option func(opts *options) 24 | 25 | // TypeCodec is a struct that stores specific type Encoder/Decoder. 26 | TypeCodec struct { 27 | ValueType reflect.Type 28 | Encoder bsoncodec.ValueEncoder 29 | Decoder bsoncodec.ValueDecoder 30 | } 31 | 32 | options = mopt.ClientOptions 33 | ) 34 | 35 | // DisableLog disables logging of mongo commands, includes info and slow logs. 36 | func DisableLog() { 37 | logMon.Set(false) 38 | logSlowMon.Set(false) 39 | } 40 | 41 | // DisableInfoLog disables info logging of mongo commands, but keeps slow logs. 42 | func DisableInfoLog() { 43 | logMon.Set(false) 44 | } 45 | 46 | // SetSlowThreshold sets the slow threshold. 47 | func SetSlowThreshold(threshold time.Duration) { 48 | slowThreshold.Set(threshold) 49 | } 50 | 51 | // WithTimeout set the mon client operation timeout. 52 | func WithTimeout(timeout time.Duration) Option { 53 | return func(opts *options) { 54 | opts.SetTimeout(timeout) 55 | } 56 | } 57 | 58 | // WithTypeCodec registers TypeCodecs to convert custom types. 59 | func WithTypeCodec(typeCodecs ...TypeCodec) Option { 60 | return func(opts *options) { 61 | registry := bson.NewRegistry() 62 | for _, v := range typeCodecs { 63 | registry.RegisterTypeEncoder(v.ValueType, v.Encoder) 64 | registry.RegisterTypeDecoder(v.ValueType, v.Decoder) 65 | } 66 | opts.SetRegistry(registry) 67 | } 68 | } 69 | 70 | func defaultTimeoutOption() Option { 71 | return func(opts *options) { 72 | opts.SetTimeout(defaultTimeout) 73 | } 74 | } 75 | -------------------------------------------------------------------------------- /stores/mongo/mon/options_test.go: -------------------------------------------------------------------------------- 1 | package mon 2 | 3 | import ( 4 | "fmt" 5 | "reflect" 6 | "testing" 7 | "time" 8 | 9 | "github.com/stretchr/testify/assert" 10 | "go.mongodb.org/mongo-driver/bson/bsoncodec" 11 | "go.mongodb.org/mongo-driver/bson/bsonrw" 12 | mopt "go.mongodb.org/mongo-driver/mongo/options" 13 | ) 14 | 15 | func TestSetSlowThreshold(t *testing.T) { 16 | assert.Equal(t, defaultSlowThreshold, slowThreshold.Load()) 17 | SetSlowThreshold(time.Second) 18 | assert.Equal(t, time.Second, slowThreshold.Load()) 19 | } 20 | 21 | func Test_defaultTimeoutOption(t *testing.T) { 22 | opts := mopt.Client() 23 | defaultTimeoutOption()(opts) 24 | assert.Equal(t, defaultTimeout, *opts.Timeout) 25 | } 26 | 27 | func TestWithTimeout(t *testing.T) { 28 | opts := mopt.Client() 29 | WithTimeout(time.Second)(opts) 30 | assert.Equal(t, time.Second, *opts.Timeout) 31 | } 32 | 33 | func TestDisableLog(t *testing.T) { 34 | assert.True(t, logMon.True()) 35 | assert.True(t, logSlowMon.True()) 36 | defer func() { 37 | logMon.Set(true) 38 | logSlowMon.Set(true) 39 | }() 40 | 41 | DisableLog() 42 | assert.False(t, logMon.True()) 43 | assert.False(t, logSlowMon.True()) 44 | } 45 | 46 | func TestDisableInfoLog(t *testing.T) { 47 | assert.True(t, logMon.True()) 48 | assert.True(t, logSlowMon.True()) 49 | defer func() { 50 | logMon.Set(true) 51 | logSlowMon.Set(true) 52 | }() 53 | 54 | DisableInfoLog() 55 | assert.False(t, logMon.True()) 56 | assert.True(t, logSlowMon.True()) 57 | } 58 | 59 | func TestWithRegistryForTimestampRegisterType(t *testing.T) { 60 | opts := mopt.Client() 61 | 62 | // mongoDateTimeEncoder allow user convert time.Time to primitive.DateTime. 63 | var mongoDateTimeEncoder bsoncodec.ValueEncoderFunc = func(ect bsoncodec.EncodeContext, w bsonrw.ValueWriter, value reflect.Value) error { 64 | // Use reflect, determine if it can be converted to time.Time. 65 | dec, ok := value.Interface().(time.Time) 66 | if !ok { 67 | return fmt.Errorf("value %v to encode is not of type time.Time", value) 68 | } 69 | return w.WriteDateTime(dec.Unix()) 70 | } 71 | 72 | // mongoDateTimeEncoder allow user convert primitive.DateTime to time.Time. 73 | var mongoDateTimeDecoder bsoncodec.ValueDecoderFunc = func(ect bsoncodec.DecodeContext, r bsonrw.ValueReader, value reflect.Value) error { 74 | primTime, err := r.ReadDateTime() 75 | if err != nil { 76 | return fmt.Errorf("error reading primitive.DateTime from ValueReader: %v", err) 77 | } 78 | value.Set(reflect.ValueOf(time.Unix(primTime, 0))) 79 | return nil 80 | } 81 | 82 | codecs := []TypeCodec{ 83 | { 84 | ValueType: reflect.TypeOf(time.Time{}), 85 | Encoder: mongoDateTimeEncoder, 86 | Decoder: mongoDateTimeDecoder, 87 | }, 88 | } 89 | WithTypeCodec(codecs...)(opts) 90 | 91 | for _, v := range codecs { 92 | // Validate Encoder 93 | enc, err := opts.Registry.LookupEncoder(v.ValueType) 94 | if err != nil { 95 | t.Fatal(err) 96 | } 97 | if assert.ObjectsAreEqual(v.Encoder, enc) { 98 | t.Errorf("Encoder got from Registry: %v, but want: %v", enc, v.Encoder) 99 | } 100 | 101 | // Validate Decoder 102 | dec, err := opts.Registry.LookupDecoder(v.ValueType) 103 | if err != nil { 104 | t.Fatal(err) 105 | } 106 | if assert.ObjectsAreEqual(v.Decoder, dec) { 107 | t.Errorf("Decoder got from Registry: %v, but want: %v", dec, v.Decoder) 108 | } 109 | } 110 | } 111 | -------------------------------------------------------------------------------- /stores/mongo/mon/trace.go: -------------------------------------------------------------------------------- 1 | package mon 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/zeromicro/go-zero/core/errorx" 7 | "github.com/zeromicro/go-zero/core/trace" 8 | "go.mongodb.org/mongo-driver/mongo" 9 | "go.opentelemetry.io/otel/attribute" 10 | "go.opentelemetry.io/otel/codes" 11 | oteltrace "go.opentelemetry.io/otel/trace" 12 | ) 13 | 14 | var mongoCmdAttributeKey = attribute.Key("mongo.cmd") 15 | 16 | func startSpan(ctx context.Context, cmd string) (context.Context, oteltrace.Span) { 17 | tracer := trace.TracerFromContext(ctx) 18 | ctx, span := tracer.Start(ctx, spanName, oteltrace.WithSpanKind(oteltrace.SpanKindClient)) 19 | span.SetAttributes(mongoCmdAttributeKey.String(cmd)) 20 | 21 | return ctx, span 22 | } 23 | 24 | func endSpan(span oteltrace.Span, err error) { 25 | defer span.End() 26 | 27 | if err == nil || errorx.In(err, mongo.ErrNoDocuments, mongo.ErrNilValue, mongo.ErrNilDocument) { 28 | span.SetStatus(codes.Ok, "") 29 | return 30 | } 31 | 32 | span.SetStatus(codes.Error, err.Error()) 33 | span.RecordError(err) 34 | } 35 | -------------------------------------------------------------------------------- /stores/mongo/mon/util.go: -------------------------------------------------------------------------------- 1 | package mon 2 | 3 | import ( 4 | "context" 5 | "encoding/json" 6 | "strings" 7 | "time" 8 | 9 | "github.com/zeromicro/go-zero/core/logx" 10 | "github.com/zeromicro/go-zero/core/timex" 11 | ) 12 | 13 | const mongoAddrSep = "," 14 | 15 | // FormatAddr formats mongo hosts to a string. 16 | func FormatAddr(hosts []string) string { 17 | return strings.Join(hosts, mongoAddrSep) 18 | } 19 | 20 | func logDuration(ctx context.Context, name, method string, startTime time.Duration, err error) { 21 | duration := timex.Since(startTime) 22 | logger := logx.WithContext(ctx).WithDuration(duration) 23 | if err != nil { 24 | logger.Errorf("mongo(%s) - %s - fail(%s)", name, method, err.Error()) 25 | return 26 | } 27 | 28 | if logSlowMon.True() && duration > slowThreshold.Load() { 29 | logger.Slowf("[MONGO] mongo(%s) - slowcall - %s - ok", name, method) 30 | } else if logMon.True() { 31 | logger.Infof("mongo(%s) - %s - ok", name, method) 32 | } 33 | } 34 | 35 | func logDurationWithDocs(ctx context.Context, name, method string, startTime time.Duration, 36 | err error, docs ...any) { 37 | duration := timex.Since(startTime) 38 | logger := logx.WithContext(ctx).WithDuration(duration) 39 | 40 | content, jerr := json.Marshal(docs) 41 | // jerr should not be non-nil, but we don't care much on this, 42 | // if non-nil, we just log without docs. 43 | if jerr != nil { 44 | if err != nil { 45 | logger.Errorf("mongo(%s) - %s - fail(%s)", name, method, err.Error()) 46 | } else if logSlowMon.True() && duration > slowThreshold.Load() { 47 | logger.Slowf("[MONGO] mongo(%s) - slowcall - %s - ok", name, method) 48 | } else if logMon.True() { 49 | logger.Infof("mongo(%s) - %s - ok", name, method) 50 | } 51 | return 52 | } 53 | 54 | if err != nil { 55 | logger.Errorf("mongo(%s) - %s - fail(%s) - %s", name, method, err.Error(), string(content)) 56 | } else if logSlowMon.True() && duration > slowThreshold.Load() { 57 | logger.Slowf("[MONGO] mongo(%s) - slowcall - %s - ok - %s", name, method, string(content)) 58 | } else if logMon.True() { 59 | logger.Infof("mongo(%s) - %s - ok - %s", name, method, string(content)) 60 | } 61 | } 62 | -------------------------------------------------------------------------------- /stores/mongo/mon/util_test.go: -------------------------------------------------------------------------------- 1 | package mon 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "testing" 7 | 8 | "github.com/stretchr/testify/assert" 9 | "github.com/zeromicro/go-zero/core/logx/logtest" 10 | "github.com/zeromicro/go-zero/core/timex" 11 | ) 12 | 13 | func TestFormatAddrs(t *testing.T) { 14 | tests := []struct { 15 | addrs []string 16 | expect string 17 | }{ 18 | { 19 | addrs: []string{"a", "b"}, 20 | expect: "a,b", 21 | }, 22 | { 23 | addrs: []string{"a", "b", "c"}, 24 | expect: "a,b,c", 25 | }, 26 | { 27 | addrs: []string{}, 28 | expect: "", 29 | }, 30 | { 31 | addrs: nil, 32 | expect: "", 33 | }, 34 | } 35 | 36 | for _, test := range tests { 37 | assert.Equal(t, test.expect, FormatAddr(test.addrs)) 38 | } 39 | } 40 | 41 | func Test_logDuration(t *testing.T) { 42 | buf := logtest.NewCollector(t) 43 | 44 | buf.Reset() 45 | logDuration(context.Background(), "foo", "bar", timex.Now()-slowThreshold.Load()*2, nil) 46 | assert.Contains(t, buf.String(), "foo") 47 | assert.Contains(t, buf.String(), "bar") 48 | assert.Contains(t, buf.String(), "slow") 49 | 50 | buf.Reset() 51 | logDuration(context.Background(), "foo", "bar", timex.Now(), nil) 52 | assert.Contains(t, buf.String(), "foo") 53 | assert.Contains(t, buf.String(), "bar") 54 | 55 | buf.Reset() 56 | logDuration(context.Background(), "foo", "bar", timex.Now(), errors.New("bar")) 57 | assert.Contains(t, buf.String(), "foo") 58 | assert.Contains(t, buf.String(), "bar") 59 | assert.Contains(t, buf.String(), "fail") 60 | 61 | defer func() { 62 | logMon.Set(true) 63 | logSlowMon.Set(true) 64 | }() 65 | 66 | buf.Reset() 67 | DisableInfoLog() 68 | logDuration(context.Background(), "foo", "bar", timex.Now(), nil) 69 | assert.Empty(t, buf.String()) 70 | 71 | buf.Reset() 72 | logDuration(context.Background(), "foo", "bar", timex.Now()-slowThreshold.Load()*2, nil) 73 | assert.Contains(t, buf.String(), "foo") 74 | assert.Contains(t, buf.String(), "bar") 75 | assert.Contains(t, buf.String(), "slow") 76 | 77 | buf.Reset() 78 | DisableLog() 79 | logDuration(context.Background(), "foo", "bar", timex.Now(), nil) 80 | assert.Empty(t, buf.String()) 81 | 82 | buf.Reset() 83 | logDuration(context.Background(), "foo", "bar", timex.Now()-slowThreshold.Load()*2, nil) 84 | assert.Empty(t, buf.String()) 85 | 86 | buf.Reset() 87 | logDuration(context.Background(), "foo", "bar", timex.Now(), errors.New("bar")) 88 | assert.Contains(t, buf.String(), "foo") 89 | assert.Contains(t, buf.String(), "bar") 90 | assert.Contains(t, buf.String(), "fail") 91 | } 92 | 93 | func Test_logDurationWithDoc(t *testing.T) { 94 | buf := logtest.NewCollector(t) 95 | buf.Reset() 96 | 97 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now()-slowThreshold.Load()*2, nil, make(chan int)) 98 | assert.Contains(t, buf.String(), "foo") 99 | assert.Contains(t, buf.String(), "bar") 100 | assert.Contains(t, buf.String(), "slow") 101 | 102 | buf.Reset() 103 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now()-slowThreshold.Load()*2, nil, "{'json': ''}") 104 | assert.Contains(t, buf.String(), "foo") 105 | assert.Contains(t, buf.String(), "bar") 106 | assert.Contains(t, buf.String(), "slow") 107 | assert.Contains(t, buf.String(), "json") 108 | 109 | buf.Reset() 110 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now(), nil, make(chan int)) 111 | assert.Contains(t, buf.String(), "foo") 112 | assert.Contains(t, buf.String(), "bar") 113 | 114 | buf.Reset() 115 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now(), nil, "{'json': ''}") 116 | assert.Contains(t, buf.String(), "foo") 117 | assert.Contains(t, buf.String(), "bar") 118 | assert.Contains(t, buf.String(), "json") 119 | 120 | buf.Reset() 121 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now(), errors.New("bar"), make(chan int)) 122 | assert.Contains(t, buf.String(), "foo") 123 | assert.Contains(t, buf.String(), "bar") 124 | assert.Contains(t, buf.String(), "fail") 125 | 126 | buf.Reset() 127 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now(), errors.New("bar"), "{'json': ''}") 128 | assert.Contains(t, buf.String(), "foo") 129 | assert.Contains(t, buf.String(), "bar") 130 | assert.Contains(t, buf.String(), "fail") 131 | assert.Contains(t, buf.String(), "json") 132 | 133 | defer func() { 134 | logMon.Set(true) 135 | logSlowMon.Set(true) 136 | }() 137 | 138 | buf.Reset() 139 | DisableInfoLog() 140 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now(), nil, make(chan int)) 141 | assert.Empty(t, buf.String()) 142 | 143 | buf.Reset() 144 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now(), nil, "{'json': ''}") 145 | assert.Empty(t, buf.String()) 146 | 147 | buf.Reset() 148 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now()-slowThreshold.Load()*2, nil, make(chan int)) 149 | assert.Contains(t, buf.String(), "foo") 150 | assert.Contains(t, buf.String(), "bar") 151 | assert.Contains(t, buf.String(), "slow") 152 | 153 | buf.Reset() 154 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now()-slowThreshold.Load()*2, nil, "{'json': ''}") 155 | assert.Contains(t, buf.String(), "foo") 156 | assert.Contains(t, buf.String(), "bar") 157 | assert.Contains(t, buf.String(), "slow") 158 | assert.Contains(t, buf.String(), "json") 159 | 160 | buf.Reset() 161 | DisableLog() 162 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now(), nil, make(chan int)) 163 | assert.Empty(t, buf.String()) 164 | 165 | buf.Reset() 166 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now(), nil, "{'json': ''}") 167 | assert.Empty(t, buf.String()) 168 | 169 | buf.Reset() 170 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now()-slowThreshold.Load()*2, nil, make(chan int)) 171 | assert.Empty(t, buf.String()) 172 | 173 | buf.Reset() 174 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now()-slowThreshold.Load()*2, nil, "{'json': ''}") 175 | assert.Empty(t, buf.String()) 176 | 177 | buf.Reset() 178 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now(), errors.New("bar"), make(chan int)) 179 | assert.Contains(t, buf.String(), "foo") 180 | assert.Contains(t, buf.String(), "bar") 181 | assert.Contains(t, buf.String(), "fail") 182 | 183 | buf.Reset() 184 | logDurationWithDocs(context.Background(), "foo", "bar", timex.Now(), errors.New("bar"), "{'json': ''}") 185 | assert.Contains(t, buf.String(), "foo") 186 | assert.Contains(t, buf.String(), "bar") 187 | assert.Contains(t, buf.String(), "fail") 188 | assert.Contains(t, buf.String(), "json") 189 | } 190 | -------------------------------------------------------------------------------- /zrpc/registry/consul/README.md: -------------------------------------------------------------------------------- 1 | ### Quick Start 2 | 3 | Prerequisites: 4 | 5 | Download the module: 6 | 7 | ```console 8 | go get -u github.com/zeromicro/zero-contrib/zrpc/registry/consul 9 | ``` 10 | 11 | For example: 12 | 13 | ## 修改RPC服务的代码 14 | 15 | - etc/\*.yaml 16 | 17 | ```yaml 18 | Consul: 19 | Host: 127.0.0.1:8500 # consul endpoint 20 | Token: 'f0512db6-76d6-f25e-f344-a98cc3484d42' # consul ACL token (optional) 21 | Key: add.rpc # service name registered to Consul 22 | Meta: 23 | Protocol: grpc 24 | Tag: 25 | - tag 26 | - rpc 27 | ``` 28 | 29 | - internal/config/config.go 30 | 31 | ```go 32 | type Config struct { 33 | zrpc.RpcServerConf 34 | Consul consul.Conf 35 | } 36 | ``` 37 | 38 | - main.go 39 | 40 | ```go 41 | import _ "github.com/zeromicro/zero-contrib/zrpc/registry/consul" 42 | 43 | func main() { 44 | flag.Parse() 45 | 46 | var c config.Config 47 | conf.MustLoad(*configFile, &c) 48 | 49 | server := zrpc.MustNewServer(c.RpcServerConf, func(grpcServer *grpc.Server) { 50 | 51 | }) 52 | // register service to consul 53 | _ = consul.RegisterService(c.ListenOn, c.Consul) 54 | 55 | server.Start() 56 | } 57 | ``` 58 | 59 | ### ACL Token 60 | 61 | the token need contain the **service** policy of `"write"` 62 | 63 | ``` 64 | service "add.rpc" { 65 | policy = "write" 66 | } 67 | service "check.rpc" { 68 | policy = "write" 69 | } 70 | ``` 71 | 72 | ## 修改API服务里的代码 73 | 74 | - main.go 75 | 76 | ```go 77 | import _ "github.com/zeromicro/zero-contrib/zrpc/registry/consul" 78 | ``` 79 | 80 | - etc/\*.yaml 81 | 82 | ```yaml 83 | # consul://[user:passwd]@host/service?param=value' 84 | # format like this 85 | Add: 86 | Target: consul://127.0.0.1:8500/add.rpc?wait=14s 87 | Check: 88 | Target: consul://127.0.0.1:8500/check.rpc?wait=14s 89 | 90 | # ACL token support 91 | Add: 92 | Target: consul://127.0.0.1:8500/add.rpc?wait=14s&token=f0512db6-76d6-f25e-f344-a98cc3484d42 93 | Check: 94 | Target: consul://127.0.0.1:8500/check.rpc?wait=14s&token=f0512db6-76d6-f25e-f344-a98cc3484d42 95 | ``` 96 | 97 | ### ACL Token 98 | 99 | the token need contain the **node** and **service** policy of `"read"` 100 | 101 | ``` 102 | node "consul-server" { 103 | policy = "read" 104 | } 105 | service "add.rpc" { 106 | policy = "read" 107 | } 108 | service "check.rpc" { 109 | policy = "read" 110 | } 111 | ``` -------------------------------------------------------------------------------- /zrpc/registry/consul/builder.go: -------------------------------------------------------------------------------- 1 | package consul 2 | 3 | import ( 4 | "context" 5 | 6 | "github.com/hashicorp/consul/api" 7 | "github.com/pkg/errors" 8 | "google.golang.org/grpc/resolver" 9 | ) 10 | 11 | // schemeName for the urls 12 | // All target URLs like 'consul://.../...' will be resolved by this resolver 13 | const schemeName = "consul" 14 | 15 | // builder implements resolver.Builder and use for constructing all consul resolvers 16 | type builder struct{} 17 | 18 | func (b *builder) Build(url resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) { 19 | tgt, err := parseURL(url.URL) 20 | if err != nil { 21 | return nil, errors.Wrap(err, "Wrong consul URL") 22 | } 23 | cli, err := api.NewClient(tgt.consulConfig()) 24 | if err != nil { 25 | return nil, errors.Wrap(err, "Couldn't connect to the Consul API") 26 | } 27 | cli.Health() 28 | ctx, cancel := context.WithCancel(context.Background()) 29 | pipe := make(chan []*consulAddr) 30 | go watchConsulService(ctx, cli.Health(), tgt, pipe) 31 | go populateEndpoints(ctx, cc, pipe) 32 | 33 | return &resolvr{cancelFunc: cancel}, nil 34 | } 35 | 36 | // Scheme returns the scheme supported by this resolver. 37 | // Scheme is defined at https://github.com/grpc/grpc/blob/master/doc/naming.md. 38 | func (b *builder) Scheme() string { 39 | return schemeName 40 | } 41 | -------------------------------------------------------------------------------- /zrpc/registry/consul/config.go: -------------------------------------------------------------------------------- 1 | package consul 2 | 3 | import ( 4 | "errors" 5 | ) 6 | 7 | const ( 8 | allEths = "0.0.0.0" 9 | envPodIP = "POD_IP" 10 | consulTags = "consul_tags" 11 | ) 12 | 13 | // Conf is the config item with the given key on etcd. 14 | type Conf struct { 15 | Host string 16 | Key string 17 | Token string `json:",optional"` 18 | Tag []string `json:",optional"` 19 | Meta map[string]string `json:",optional"` 20 | TTL int `json:"ttl,optional"` 21 | } 22 | 23 | // Validate validates c. 24 | func (c Conf) Validate() error { 25 | if len(c.Host) == 0 { 26 | return errors.New("empty consul hosts") 27 | } 28 | if len(c.Key) == 0 { 29 | return errors.New("empty consul key") 30 | } 31 | if c.TTL == 0 { 32 | c.TTL = 20 33 | } 34 | 35 | return nil 36 | } 37 | -------------------------------------------------------------------------------- /zrpc/registry/consul/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/zrpc/registry/consul 2 | 3 | go 1.15 4 | 5 | require ( 6 | github.com/hashicorp/consul/api v1.25.1 7 | github.com/jpillora/backoff v1.0.0 8 | github.com/pkg/errors v0.9.1 9 | github.com/zeromicro/go-zero v1.5.4 10 | google.golang.org/grpc v1.58.3 11 | ) 12 | 13 | require ( 14 | github.com/hashicorp/go-msgpack v1.1.5 // indirect 15 | github.com/miekg/dns v1.1.45 // indirect 16 | ) 17 | -------------------------------------------------------------------------------- /zrpc/registry/consul/register.go: -------------------------------------------------------------------------------- 1 | package consul 2 | 3 | import ( 4 | "fmt" 5 | "net" 6 | "os" 7 | "strconv" 8 | "strings" 9 | "time" 10 | 11 | "github.com/hashicorp/consul/api" 12 | "github.com/zeromicro/go-zero/core/logx" 13 | "github.com/zeromicro/go-zero/core/netx" 14 | "github.com/zeromicro/go-zero/core/proc" 15 | ) 16 | 17 | // RegisterService register service to consul 18 | func RegisterService(listenOn string, c Conf) error { 19 | pubListenOn := figureOutListenOn(listenOn) 20 | 21 | host, ports, err := net.SplitHostPort(pubListenOn) 22 | if err != nil { 23 | return fmt.Errorf("failed parsing address error: %v", err) 24 | } 25 | port, _ := strconv.ParseUint(ports, 10, 16) 26 | 27 | client, err := api.NewClient(&api.Config{Scheme: "http", Address: c.Host, Token: c.Token}) 28 | if err != nil { 29 | return fmt.Errorf("create consul client error: %v", err) 30 | } 31 | // 服务节点的名称 32 | serviceID := fmt.Sprintf("%s-%s-%d", c.Key, host, port) 33 | 34 | if c.TTL <= 0 { 35 | c.TTL = 20 36 | } 37 | 38 | ttl := fmt.Sprintf("%ds", c.TTL) 39 | expiredTTL := fmt.Sprintf("%ds", c.TTL*3) 40 | 41 | reg := &api.AgentServiceRegistration{ 42 | ID: serviceID, // 服务节点的名称 43 | Name: c.Key, // 服务名称 44 | Tags: c.Tag, // tag，可以为空 45 | Meta: c.Meta, // meta， 可以为空 46 | Port: int(port), // 服务端口 47 | Address: host, // 服务 IP 48 | Checks: []*api.AgentServiceCheck{ // 健康检查 49 | { 50 | CheckID: serviceID, // 服务节点的名称 51 | TTL: ttl, // 健康检查间隔 52 | Status: "passing", 53 | DeregisterCriticalServiceAfter: expiredTTL, // 注销时间，相当于过期时间 54 | }, 55 | }, 56 | } 57 | 58 | if err := client.Agent().ServiceRegister(reg); err != nil { 59 | return fmt.Errorf("initial register service '%s' host to consul error: %s", c.Key, err.Error()) 60 | } 61 | 62 | // initial register service check 63 | check := api.AgentServiceCheck{TTL: ttl, Status: "passing", DeregisterCriticalServiceAfter: expiredTTL} 64 | err = client.Agent().CheckRegister(&api.AgentCheckRegistration{ID: serviceID, Name: c.Key, ServiceID: serviceID, AgentServiceCheck: check}) 65 | if err != nil { 66 | return fmt.Errorf("initial register service check to consul error: %s", err.Error()) 67 | } 68 | 69 | ttlTicker := time.Duration(c.TTL-1) * time.Second 70 | if ttlTicker < time.Second { 71 | ttlTicker = time.Second 72 | } 73 | // routine to update ttl 74 | go func() { 75 | ticker := time.NewTicker(ttlTicker) 76 | defer ticker.Stop() 77 | for { 78 | <-ticker.C 79 | err = client.Agent().UpdateTTL(serviceID, "", "passing") 80 | logx.Info("update ttl") 81 | if err != nil { 82 | logx.Infof("update ttl of service error: %v", err.Error()) 83 | } 84 | } 85 | }() 86 | // consul deregister 87 | proc.AddShutdownListener(func() { 88 | err := client.Agent().ServiceDeregister(serviceID) 89 | if err != nil { 90 | logx.Info("deregister service error: ", err.Error()) 91 | } 92 | logx.Info("deregistered service from consul server.") 93 | }) 94 | 95 | return nil 96 | } 97 | 98 | func figureOutListenOn(listenOn string) string { 99 | fields := strings.Split(listenOn, ":") 100 | if len(fields) == 0 { 101 | return listenOn 102 | } 103 | 104 | host := fields[0] 105 | if len(host) > 0 && host != allEths { 106 | return listenOn 107 | } 108 | 109 | ip := os.Getenv(envPodIP) 110 | if len(ip) == 0 { 111 | ip = netx.InternalIp() 112 | } 113 | if len(ip) == 0 { 114 | return listenOn 115 | } 116 | 117 | return strings.Join(append([]string{ip}, fields[1:]...), ":") 118 | } 119 | -------------------------------------------------------------------------------- /zrpc/registry/consul/resolver.go: -------------------------------------------------------------------------------- 1 | package consul 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "sort" 7 | "strings" 8 | "time" 9 | 10 | "github.com/hashicorp/consul/api" 11 | "github.com/jpillora/backoff" 12 | "github.com/zeromicro/go-zero/core/logx" 13 | "google.golang.org/grpc/attributes" 14 | "google.golang.org/grpc/resolver" 15 | ) 16 | 17 | func init() { 18 | resolver.Register(&builder{}) 19 | } 20 | 21 | type resolvr struct { 22 | cancelFunc context.CancelFunc 23 | } 24 | 25 | type consulAddr struct { 26 | Addr string 27 | Port int 28 | Tags []string 29 | } 30 | 31 | func (r *resolvr) ResolveNow(resolver.ResolveNowOptions) {} 32 | 33 | // Close closes the resolver. 34 | func (r *resolvr) Close() { 35 | r.cancelFunc() 36 | } 37 | 38 | type servicer interface { 39 | Service(string, string, bool, *api.QueryOptions) ([]*api.ServiceEntry, *api.QueryMeta, error) 40 | } 41 | 42 | func watchConsulService(ctx context.Context, s servicer, tgt target, out chan<- []*consulAddr) { 43 | res := make(chan []*consulAddr) 44 | quit := make(chan struct{}) 45 | bck := &backoff.Backoff{ 46 | Factor: 2, 47 | Jitter: true, 48 | Min: 10 * time.Millisecond, 49 | Max: tgt.MaxBackoff, 50 | } 51 | go func() { 52 | var lastIndex uint64 53 | for { 54 | ss, meta, err := s.Service( 55 | tgt.Service, 56 | tgt.Tag, 57 | tgt.Healthy, 58 | &api.QueryOptions{ 59 | WaitIndex: lastIndex, 60 | Near: tgt.Near, 61 | WaitTime: tgt.Wait, 62 | Datacenter: tgt.Dc, 63 | AllowStale: tgt.AllowStale, 64 | RequireConsistent: tgt.RequireConsistent, 65 | }, 66 | ) 67 | if err != nil { 68 | logx.Errorf("[Consul resolver] Couldn't fetch endpoints. target={%s}; error={%v}", tgt.String(), err) 69 | 70 | time.Sleep(bck.Duration()) 71 | continue 72 | } 73 | 74 | bck.Reset() 75 | lastIndex = meta.LastIndex 76 | logx.Infof("[Consul resolver] %d endpoints fetched in(+wait) %s for target={%s}", 77 | len(ss), 78 | meta.RequestTime, 79 | tgt.String(), 80 | ) 81 | 82 | ee := make([]*consulAddr, 0, len(ss)) 83 | for _, s := range ss { 84 | address := s.Service.Address 85 | if s.Service.Address == "" { 86 | address = s.Node.Address 87 | } 88 | ee = append(ee, &consulAddr{ 89 | Addr: address, 90 | Port: s.Service.Port, 91 | Tags: s.Service.Tags, 92 | }) 93 | } 94 | 95 | if tgt.Limit != 0 && len(ee) > tgt.Limit { 96 | ee = ee[:tgt.Limit] 97 | } 98 | select { 99 | case res <- ee: 100 | continue 101 | case <-quit: 102 | return 103 | } 104 | } 105 | }() 106 | 107 | for { 108 | select { 109 | case ee := <-res: 110 | out <- ee 111 | case <-ctx.Done(): 112 | close(quit) 113 | return 114 | } 115 | } 116 | } 117 | 118 | func populateEndpoints(ctx context.Context, clientConn resolver.ClientConn, input <-chan []*consulAddr) { 119 | for { 120 | select { 121 | case cc := <-input: 122 | connsSet := make(map[string][]string, len(cc)) 123 | for _, c := range cc { 124 | connsSet[fmt.Sprintf("%s:%d", c.Addr, c.Port)] = c.Tags 125 | } 126 | conns := make([]resolver.Address, 0, len(connsSet)) 127 | for c, tags := range connsSet { 128 | rAddr := resolver.Address{Addr: c} 129 | if tags != nil { 130 | rAddr.Attributes = attributes.New(consulTags, strings.Join(tags, ",")) 131 | } 132 | conns = append(conns, rAddr) 133 | } 134 | 135 | sort.Sort(byAddressString(conns)) // Don't replace the same address list in the balancer 136 | _ = clientConn.UpdateState(resolver.State{Addresses: conns}) 137 | case <-ctx.Done(): 138 | logx.Info("[Consul resolver] Watch has been finished") 139 | return 140 | } 141 | } 142 | } 143 | 144 | // byAddressString sorts resolver.Address by Address Field sorting in increasing order. 145 | type byAddressString []resolver.Address 146 | 147 | func (p byAddressString) Len() int { return len(p) } 148 | func (p byAddressString) Less(i, j int) bool { return p[i].Addr < p[j].Addr } 149 | func (p byAddressString) Swap(i, j int) { p[i], p[j] = p[j], p[i] } 150 | -------------------------------------------------------------------------------- /zrpc/registry/consul/target.go: -------------------------------------------------------------------------------- 1 | package consul 2 | 3 | import ( 4 | "fmt" 5 | "net/http" 6 | "net/url" 7 | "strings" 8 | "time" 9 | 10 | "github.com/hashicorp/consul/api" 11 | "github.com/pkg/errors" 12 | "github.com/zeromicro/go-zero/core/mapping" 13 | ) 14 | 15 | type target struct { 16 | Addr string `key:",optional"` 17 | User string `key:",optional"` 18 | Password string `key:",optional"` 19 | Service string `key:",optional"` 20 | Wait time.Duration `key:"wait,optional"` 21 | Timeout time.Duration `key:"timeout,optional"` 22 | MaxBackoff time.Duration `key:"max-backoff,optional"` 23 | Tag string `key:"tag,optional"` 24 | Near string `key:"near,optional"` 25 | Limit int `key:"limit,optional"` 26 | Healthy bool `key:"healthy,optional"` 27 | TLSInsecure bool `key:"insecure,optional"` 28 | Token string `key:"token,optional"` 29 | Dc string `key:"dc,optional"` 30 | AllowStale bool `key:"allow-stale,optional"` 31 | RequireConsistent bool `key:"require-consistent,optional"` 32 | } 33 | 34 | func (t *target) String() string { 35 | return fmt.Sprintf("service='%s' healthy='%t' tag='%s'", t.Service, t.Healthy, t.Tag) 36 | } 37 | 38 | // parseURL with parameters 39 | func parseURL(rawURL url.URL) (target, error) { 40 | if rawURL.Scheme != schemeName || 41 | len(rawURL.Host) == 0 || len(strings.TrimLeft(rawURL.Path, "/")) == 0 { 42 | return target{}, 43 | errors.Errorf("Malformed URL('%s'). Must be in the next format: 'consul://[user:passwd]@host/service?param=value'", rawURL.String()) 44 | } 45 | 46 | var tgt target 47 | params := make(map[string]interface{}, len(rawURL.Query())) 48 | for name, value := range rawURL.Query() { 49 | params[name] = value[0] 50 | } 51 | err := mapping.UnmarshalKey(params, &tgt) 52 | if err != nil { 53 | return target{}, errors.Wrap(err, "Malformed URL parameters") 54 | } 55 | 56 | tgt.User = rawURL.User.Username() 57 | tgt.Password, _ = rawURL.User.Password() 58 | tgt.Addr = rawURL.Host 59 | tgt.Service = strings.TrimLeft(rawURL.Path, "/") 60 | 61 | if len(tgt.Near) == 0 { 62 | tgt.Near = "_agent" 63 | } 64 | if tgt.MaxBackoff == 0 { 65 | tgt.MaxBackoff = time.Second 66 | } 67 | 68 | return tgt, nil 69 | } 70 | 71 | // consulConfig returns config based on the parsed target. 72 | // It uses custom http-client. 73 | func (t *target) consulConfig() *api.Config { 74 | var creds *api.HttpBasicAuth 75 | if len(t.User) > 0 && len(t.Password) > 0 { 76 | creds = new(api.HttpBasicAuth) 77 | creds.Password = t.Password 78 | creds.Username = t.User 79 | } 80 | // custom http.Client 81 | c := &http.Client{ 82 | Timeout: t.Timeout, 83 | } 84 | return &api.Config{ 85 | Address: t.Addr, 86 | HttpAuth: creds, 87 | WaitTime: t.Wait, 88 | HttpClient: c, 89 | TLSConfig: api.TLSConfig{ 90 | InsecureSkipVerify: t.TLSInsecure, 91 | }, 92 | Token: t.Token, 93 | } 94 | } 95 | -------------------------------------------------------------------------------- /zrpc/registry/consul/tests/client_test.go: -------------------------------------------------------------------------------- 1 | package tests 2 | 3 | import ( 4 | "fmt" 5 | "testing" 6 | "time" 7 | 8 | _ "github.com/zeromicro/zero-contrib/zrpc/registry/consul" 9 | "google.golang.org/grpc" 10 | ) 11 | 12 | func TestCLient(t *testing.T) { 13 | svcCfg := fmt.Sprintf(`{"loadBalancingPolicy":"%s"}`, "round_robin") 14 | conn, err := grpc.Dial("consul://127.0.0.1:8500/gozero?wait=14s&tag=public", grpc.WithInsecure(), grpc.WithDefaultServiceConfig(svcCfg)) 15 | if err != nil { 16 | t.Fatal(err) 17 | } 18 | defer conn.Close() 19 | time.Sleep(29 * time.Second) 20 | } 21 | -------------------------------------------------------------------------------- /zrpc/registry/nacos/README.md: -------------------------------------------------------------------------------- 1 | ### Quick Start 2 | 3 | Prerequisites: 4 | 5 | Download the module: 6 | 7 | ```console 8 | go get -u github.com/zeromicro/zero-contrib/zrpc/registry/nacos 9 | ``` 10 | 11 | For example: 12 | 13 | ## Service 14 | 15 | - main.go 16 | 17 | ```go 18 | import _ "github.com/zeromicro/zero-contrib/zrpc/registry/nacos" 19 | 20 | func main() { 21 | flag.Parse() 22 | 23 | var c config.Config 24 | conf.MustLoad(*configFile, &c) 25 | 26 | server := zrpc.MustNewServer(c.RpcServerConf, func(grpcServer *grpc.Server) { 27 | 28 | }) 29 | // register service to nacos 30 | sc := []constant.ServerConfig{ 31 | *constant.NewServerConfig("192.168.100.15", 8848), 32 | } 33 | 34 | cc := &constant.ClientConfig{ 35 | NamespaceId: "public", 36 | TimeoutMs: 5000, 37 | NotLoadCacheAtStart: true, 38 | LogDir: "/tmp/nacos/log", 39 | CacheDir: "/tmp/nacos/cache", 40 | RotateTime: "1h", 41 | MaxAge: 3, 42 | LogLevel: "debug", 43 | } 44 | 45 | opts := nacos.NewNacosConfig("nacos.rpc", c.ListenOn, sc, cc) 46 | _ = nacos.RegisterService(opts) 47 | server.Start() 48 | } 49 | ``` 50 | 51 | ## Client 52 | 53 | - main.go 54 | 55 | ```go 56 | import _ "github.com/zeromicro/zero-contrib/zrpc/registry/nacos" 57 | ``` 58 | 59 | - etc/\*.yaml 60 | 61 | ```yaml 62 | # nacos://[user:passwd]@host/service?param=value' 63 | Target: nacos://192.168.100.15:8848/nacos.rpc?namespaceid=public&timeout=5000s 64 | ``` 65 | -------------------------------------------------------------------------------- /zrpc/registry/nacos/builder.go: -------------------------------------------------------------------------------- 1 | package nacos 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "net" 7 | "strconv" 8 | 9 | "github.com/nacos-group/nacos-sdk-go/v2/clients" 10 | "github.com/nacos-group/nacos-sdk-go/v2/common/constant" 11 | "github.com/nacos-group/nacos-sdk-go/v2/vo" 12 | "github.com/pkg/errors" 13 | "google.golang.org/grpc/resolver" 14 | ) 15 | 16 | func init() { 17 | resolver.Register(&builder{}) 18 | } 19 | 20 | // schemeName for the urls 21 | // All target URLs like 'nacos://.../...' will be resolved by this resolver 22 | const schemeName = "nacos" 23 | 24 | // builder implements resolver.Builder and use for constructing all consul resolvers 25 | type builder struct{} 26 | 27 | func (b *builder) Build(url resolver.Target, conn resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) { 28 | tgt, err := parseURL(url.URL) 29 | if err != nil { 30 | return nil, errors.Wrap(err, "Wrong nacos URL") 31 | } 32 | 33 | host, ports, err := net.SplitHostPort(tgt.Addr) 34 | if err != nil { 35 | return nil, fmt.Errorf("failed parsing address error: %v", err) 36 | } 37 | port, _ := strconv.ParseUint(ports, 10, 16) 38 | 39 | sc := []constant.ServerConfig{ 40 | *constant.NewServerConfig(host, port), 41 | } 42 | 43 | cc := &constant.ClientConfig{ 44 | AppName: tgt.AppName, 45 | NamespaceId: tgt.NamespaceID, 46 | Username: tgt.User, 47 | Password: tgt.Password, 48 | TimeoutMs: uint64(tgt.Timeout), 49 | NotLoadCacheAtStart: tgt.NotLoadCacheAtStart, 50 | UpdateCacheWhenEmpty: tgt.UpdateCacheWhenEmpty, 51 | } 52 | 53 | if tgt.CacheDir != "" { 54 | cc.CacheDir = tgt.CacheDir 55 | } 56 | if tgt.LogDir != "" { 57 | cc.LogDir = tgt.LogDir 58 | } 59 | if tgt.LogLevel != "" { 60 | cc.LogLevel = tgt.LogLevel 61 | } 62 | 63 | cli, err := clients.NewNamingClient(vo.NacosClientParam{ 64 | ServerConfigs: sc, 65 | ClientConfig: cc, 66 | }) 67 | if err != nil { 68 | return nil, errors.Wrap(err, "Couldn't connect to the nacos API") 69 | } 70 | 71 | ctx, cancel := context.WithCancel(context.Background()) 72 | pipe := make(chan []string) 73 | 74 | go cli.Subscribe(&vo.SubscribeParam{ 75 | ServiceName: tgt.Service, 76 | Clusters: tgt.Clusters, 77 | GroupName: tgt.GroupName, 78 | SubscribeCallback: newWatcher(ctx, cancel, pipe).CallBackHandle, // required 79 | }) 80 | 81 | go populateEndpoints(ctx, conn, pipe) 82 | 83 | return &resolvr{cancelFunc: cancel}, nil 84 | } 85 | 86 | // Scheme returns the scheme supported by this resolver. 87 | // Scheme is defined at https://github.com/grpc/grpc/blob/master/doc/naming.md. 88 | func (b *builder) Scheme() string { 89 | return schemeName 90 | } 91 | -------------------------------------------------------------------------------- /zrpc/registry/nacos/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/zrpc/registry/nacos 2 | 3 | go 1.15 4 | 5 | require ( 6 | github.com/pkg/errors v0.9.1 7 | github.com/zeromicro/go-zero v1.5.4 8 | google.golang.org/grpc v1.58.3 9 | ) 10 | 11 | require ( 12 | github.com/jmespath/go-jmespath v0.4.0 // indirect 13 | github.com/nacos-group/nacos-sdk-go/v2 v2.2.3 14 | ) 15 | -------------------------------------------------------------------------------- /zrpc/registry/nacos/options.go: -------------------------------------------------------------------------------- 1 | package nacos 2 | 3 | import "github.com/nacos-group/nacos-sdk-go/v2/common/constant" 4 | 5 | const ( 6 | allEths = "0.0.0.0" 7 | envPodIP = "POD_IP" 8 | ) 9 | 10 | // options 11 | type Options struct { 12 | ListenOn string 13 | ServiceName string 14 | Prefix string 15 | Weight float64 16 | Cluster string 17 | Group string 18 | Metadata map[string]string 19 | 20 | ServerConfig []constant.ServerConfig 21 | ClientConfig *constant.ClientConfig 22 | } 23 | 24 | type Option func(*Options) 25 | 26 | func NewNacosConfig(serviceName, listenOn string, sc []constant.ServerConfig, cc *constant.ClientConfig, opts ...Option) *Options { 27 | options := &Options{ 28 | ServiceName: serviceName, 29 | ListenOn: listenOn, 30 | ServerConfig: sc, 31 | ClientConfig: cc, 32 | Weight: 1.0, 33 | Metadata: make(map[string]string), 34 | } 35 | 36 | for _, opt := range opts { 37 | opt(options) 38 | } 39 | 40 | return options 41 | } 42 | 43 | func WithPrefix(prefix string) Option { 44 | return func(o *Options) { 45 | o.Prefix = prefix 46 | } 47 | } 48 | 49 | func WithWeight(weight float64) Option { 50 | return func(o *Options) { 51 | o.Weight = weight 52 | } 53 | } 54 | 55 | func WithCluster(cluster string) Option { 56 | return func(o *Options) { 57 | o.Cluster = cluster 58 | } 59 | } 60 | 61 | func WithGroup(group string) Option { 62 | return func(o *Options) { 63 | o.Group = group 64 | } 65 | } 66 | 67 | func WithMetadata(metadata map[string]string) Option { 68 | return func(o *Options) { 69 | o.Metadata = metadata 70 | } 71 | } 72 | -------------------------------------------------------------------------------- /zrpc/registry/nacos/register.go: -------------------------------------------------------------------------------- 1 | package nacos 2 | 3 | import ( 4 | "fmt" 5 | "log" 6 | "net" 7 | "os" 8 | "strconv" 9 | "strings" 10 | 11 | "github.com/nacos-group/nacos-sdk-go/v2/clients" 12 | "github.com/nacos-group/nacos-sdk-go/v2/vo" 13 | "github.com/zeromicro/go-zero/core/logx" 14 | "github.com/zeromicro/go-zero/core/proc" 15 | 16 | "github.com/zeromicro/go-zero/core/netx" 17 | ) 18 | 19 | // RegisterService register service to nacos 20 | func RegisterService(opts *Options) error { 21 | pubListenOn := figureOutListenOn(opts.ListenOn) 22 | 23 | host, ports, err := net.SplitHostPort(pubListenOn) 24 | if err != nil { 25 | return fmt.Errorf("failed parsing address error: %v", err) 26 | } 27 | port, _ := strconv.ParseUint(ports, 10, 16) 28 | 29 | client, err := clients.NewNamingClient( 30 | vo.NacosClientParam{ 31 | ServerConfigs: opts.ServerConfig, 32 | ClientConfig: opts.ClientConfig, 33 | }, 34 | ) 35 | if err != nil { 36 | log.Panic(err) 37 | } 38 | 39 | // service register 40 | _, err = client.RegisterInstance(vo.RegisterInstanceParam{ 41 | ServiceName: opts.ServiceName, 42 | Ip: host, 43 | Port: port, 44 | Weight: opts.Weight, 45 | Enable: true, 46 | Healthy: true, 47 | Ephemeral: true, 48 | Metadata: opts.Metadata, 49 | ClusterName: opts.Cluster, 50 | GroupName: opts.Group, 51 | }) 52 | 53 | if err != nil { 54 | return err 55 | } 56 | 57 | // service deregister 58 | proc.AddShutdownListener(func() { 59 | _, err := client.DeregisterInstance(vo.DeregisterInstanceParam{ 60 | Ip: host, 61 | Port: port, 62 | ServiceName: opts.ServiceName, 63 | Cluster: opts.Cluster, 64 | GroupName: opts.Group, 65 | Ephemeral: true, 66 | }) 67 | if err != nil { 68 | logx.Info("deregister service error: ", err.Error()) 69 | } else { 70 | logx.Info("deregistered service from nacos server.") 71 | } 72 | }) 73 | 74 | return nil 75 | } 76 | 77 | func figureOutListenOn(listenOn string) string { 78 | fields := strings.Split(listenOn, ":") 79 | if len(fields) == 0 { 80 | return listenOn 81 | } 82 | 83 | host := fields[0] 84 | if len(host) > 0 && host != allEths { 85 | return listenOn 86 | } 87 | 88 | ip := os.Getenv(envPodIP) 89 | if len(ip) == 0 { 90 | ip = netx.InternalIp() 91 | } 92 | if len(ip) == 0 { 93 | return listenOn 94 | } 95 | 96 | return strings.Join(append([]string{ip}, fields[1:]...), ":") 97 | } 98 | -------------------------------------------------------------------------------- /zrpc/registry/nacos/resolver.go: -------------------------------------------------------------------------------- 1 | package nacos 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "sort" 7 | 8 | "github.com/nacos-group/nacos-sdk-go/v2/common/logger" 9 | "github.com/nacos-group/nacos-sdk-go/v2/model" 10 | "github.com/zeromicro/go-zero/core/logx" 11 | "google.golang.org/grpc/resolver" 12 | ) 13 | 14 | type resolvr struct { 15 | cancelFunc context.CancelFunc 16 | } 17 | 18 | func (r *resolvr) ResolveNow(resolver.ResolveNowOptions) {} 19 | 20 | // Close closes the resolver. 21 | func (r *resolvr) Close() { 22 | r.cancelFunc() 23 | } 24 | 25 | type watcher struct { 26 | ctx context.Context 27 | cancel context.CancelFunc 28 | out chan<- []string 29 | } 30 | 31 | func newWatcher(ctx context.Context, cancel context.CancelFunc, out chan<- []string) *watcher { 32 | return &watcher{ 33 | ctx: ctx, 34 | cancel: cancel, 35 | out: out, 36 | } 37 | } 38 | 39 | func (nw *watcher) CallBackHandle(services []model.Instance, err error) { 40 | if err != nil { 41 | logger.Error("[Nacos resolver] watcher call back handle error:%v", err) 42 | return 43 | } 44 | ee := make([]string, 0, len(services)) 45 | for _, s := range services { 46 | if s.Metadata != nil && s.Metadata["gRPC_port"] != "" { 47 | ee = append(ee, fmt.Sprintf("%s:%s", s.Ip, s.Metadata["gRPC_port"])) 48 | } else { 49 | ee = append(ee, fmt.Sprintf("%s:%d", s.Ip, s.Port)) 50 | } 51 | } 52 | nw.out <- ee 53 | } 54 | 55 | func populateEndpoints(ctx context.Context, clientConn resolver.ClientConn, input <-chan []string) { 56 | for { 57 | select { 58 | case cc := <-input: 59 | connsSet := make(map[string]struct{}, len(cc)) 60 | for _, c := range cc { 61 | connsSet[c] = struct{}{} 62 | } 63 | conns := make([]resolver.Address, 0, len(connsSet)) 64 | for c := range connsSet { 65 | conns = append(conns, resolver.Address{Addr: c}) 66 | } 67 | sort.Sort(byAddressString(conns)) // Don't replace the same address list in the balancer 68 | _ = clientConn.UpdateState(resolver.State{Addresses: conns}) 69 | case <-ctx.Done(): 70 | logx.Info("[Nacos resolver] Watch has been finished") 71 | return 72 | } 73 | } 74 | } 75 | 76 | // byAddressString sorts resolver.Address by Address Field sorting in increasing order. 77 | type byAddressString []resolver.Address 78 | 79 | func (p byAddressString) Len() int { return len(p) } 80 | func (p byAddressString) Less(i, j int) bool { return p[i].Addr < p[j].Addr } 81 | func (p byAddressString) Swap(i, j int) { p[i], p[j] = p[j], p[i] } 82 | -------------------------------------------------------------------------------- /zrpc/registry/nacos/target.go: -------------------------------------------------------------------------------- 1 | package nacos 2 | 3 | import ( 4 | "net/url" 5 | "os" 6 | "strings" 7 | "time" 8 | 9 | "github.com/pkg/errors" 10 | "github.com/zeromicro/go-zero/core/mapping" 11 | ) 12 | 13 | type target struct { 14 | Addr string `key:",optional"` 15 | User string `key:",optional"` 16 | Password string `key:",optional"` 17 | Service string `key:",optional"` 18 | GroupName string `key:",optional"` 19 | Clusters []string `key:",optional"` 20 | NamespaceID string `key:"namespaceid,optional"` 21 | Timeout time.Duration `key:"timeout,optional"` 22 | AppName string `key:"appName,optional"` 23 | LogLevel string `key:",optional"` 24 | LogDir string `key:",optional"` 25 | CacheDir string `key:",optional"` 26 | NotLoadCacheAtStart bool `key:"notLoadCacheAtStart,optional"` 27 | UpdateCacheWhenEmpty bool `key:"updateCacheWhenEmpty,optional"` 28 | } 29 | 30 | // parseURL with parameters 31 | func parseURL(rawURL url.URL) (target, error) { 32 | if rawURL.Scheme != schemeName || 33 | len(rawURL.Host) == 0 || len(strings.TrimLeft(rawURL.Path, "/")) == 0 { 34 | return target{}, 35 | errors.Errorf("Malformed URL('%s'). Must be in the next format: 'nacos://[user:passwd]@host/service?param=value'", rawURL.String()) 36 | } 37 | 38 | var tgt target 39 | params := make(map[string]interface{}, len(rawURL.Query())) 40 | for name, value := range rawURL.Query() { 41 | params[name] = value[0] 42 | } 43 | 44 | err := mapping.UnmarshalKey(params, &tgt) 45 | if err != nil { 46 | return target{}, errors.Wrap(err, "Malformed URL parameters") 47 | } 48 | 49 | if tgt.NamespaceID == "" { 50 | tgt.NamespaceID = "public" 51 | } 52 | 53 | tgt.LogLevel = os.Getenv("NACOS_LOG_LEVEL") 54 | tgt.LogDir = os.Getenv("NACOS_LOG_DIR") 55 | tgt.CacheDir = os.Getenv("NACOS_CACHE_DIR") 56 | 57 | tgt.User = rawURL.User.Username() 58 | tgt.Password, _ = rawURL.User.Password() 59 | tgt.Addr = rawURL.Host 60 | tgt.Service = strings.TrimLeft(rawURL.Path, "/") 61 | 62 | if logLevel, exists := os.LookupEnv("NACOS_LOG_LEVEL"); exists { 63 | tgt.LogLevel = logLevel 64 | } 65 | 66 | if logDir, exists := os.LookupEnv("NACOS_LOG_DIR"); exists { 67 | tgt.LogDir = logDir 68 | } 69 | 70 | if notLoadCacheAtStart, exists := os.LookupEnv("NACOS_NOT_LOAD_CACHE_AT_START"); exists { 71 | tgt.NotLoadCacheAtStart = notLoadCacheAtStart == "true" 72 | } 73 | 74 | if updateCacheWhenEmpty, exists := os.LookupEnv("NACOS_UPDATE_CACHE_WHEN_EMPTY"); exists { 75 | tgt.UpdateCacheWhenEmpty = updateCacheWhenEmpty == "true" 76 | } 77 | 78 | return tgt, nil 79 | } 80 | -------------------------------------------------------------------------------- /zrpc/registry/polaris/README.md: -------------------------------------------------------------------------------- 1 | ### Quick Start 2 | 3 | Prerequisites: 4 | 5 | Download the module: 6 | 7 | ```console 8 | go get -u github.com/zeromicro/zero-contrib/zrpc/registry/polaris 9 | ``` 10 | 11 | For example: 12 | 13 | ## Service 14 | 15 | - ./polaris.yaml 16 | 17 | ```yaml 18 | global: 19 | serverConnector: 20 | addresses: 21 | - 127.0.0.1:8091 22 | ``` 23 | 24 | - main.go 25 | 26 | ```go 27 | import _ "github.com/zeromicro/zero-contrib/zrpc/registry/polaris" 28 | 29 | func main() { 30 | flag.Parse() 31 | 32 | var c config.Config 33 | conf.MustLoad(*configFile, &c) 34 | 35 | server := zrpc.MustNewServer(c.RpcServerConf, func(grpcServer *grpc.Server) { 36 | 37 | }) 38 | // register service to polaris 39 | opts := polaris.NewPolarisConfig(c.ListenOn) 40 | opts.ServiceName = "EchoServerZero" 41 | opts.Namespace = "default" 42 | opts.ServiceToken = "2af8fdf2534f451e8f01881d1b66f9ec" 43 | _ = polaris.RegisterService(opts) 44 | 45 | server.Start() 46 | } 47 | ``` 48 | 49 | ## Client 50 | 51 | - main.go 52 | 53 | ```go 54 | import _ "github.com/zeromicro/zero-contrib/zrpc/registry/polaris" 55 | ``` 56 | 57 | - etc/\*.yaml 58 | 59 | ```yaml 60 | # polaris://[user:passwd]@host/service?param=value' 61 | Target: polaris://127.0.0.1:8091/EchoServerZero?wait=14s 62 | ``` 63 | -------------------------------------------------------------------------------- /zrpc/registry/polaris/builder.go: -------------------------------------------------------------------------------- 1 | package polaris 2 | 3 | import ( 4 | "context" 5 | "strings" 6 | "sync" 7 | 8 | "github.com/pkg/errors" 9 | "google.golang.org/grpc/resolver" 10 | 11 | "github.com/polarismesh/polaris-go/api" 12 | "github.com/polarismesh/polaris-go/pkg/config" 13 | "github.com/polarismesh/polaris-go/pkg/model" 14 | ) 15 | 16 | var ( 17 | consumers map[string]api.ConsumerAPI = make(map[string]api.ConsumerAPI) 18 | lock *sync.Mutex = &sync.Mutex{} 19 | ) 20 | 21 | func init() { 22 | resolver.Register(&builder{}) 23 | } 24 | 25 | // builder implements resolver.Builder and use for constructing all consul resolvers 26 | type builder struct{} 27 | 28 | func (b *builder) Build(url resolver.Target, conn resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) { 29 | dsn := strings.Join([]string{schemeName + ":/", url.URL.Host, url.URL.Path}, "/") 30 | tgr, err := parseURL(dsn) 31 | if err != nil { 32 | return nil, errors.Wrap(err, "Wrong polaris URL") 33 | } 34 | 35 | var polarisErr error 36 | 37 | func() { 38 | lock.Lock() 39 | defer lock.Unlock() 40 | 41 | if _, exist := consumers[tgr.Addr]; exist { 42 | return 43 | } 44 | 45 | sdkCtx, err := api.InitContextByConfig(config.NewDefaultConfiguration([]string{tgr.Addr})) 46 | if err != nil { 47 | polarisErr = errors.Wrap(err, "Fail init polaris SDKContext") 48 | return 49 | } 50 | consumerAPI := api.NewConsumerAPIByContext(sdkCtx) 51 | consumers[tgr.Addr] = consumerAPI 52 | }() 53 | 54 | if polarisErr != nil { 55 | return nil, polarisErr 56 | } 57 | 58 | consumerAPI := consumers[tgr.Addr] 59 | ctx, cancel := context.WithCancel(context.Background()) 60 | pipe := make(chan []string, 4) 61 | 62 | go newWatcher(pipe).startWatch(ctx, consumerAPI, &api.WatchServiceRequest{ 63 | WatchServiceRequest: model.WatchServiceRequest{ 64 | Key: model.ServiceKey{ 65 | Namespace: tgr.Namespace, 66 | Service: tgr.Service, 67 | }, 68 | }, 69 | }) 70 | 71 | go populateEndpoints(ctx, conn, pipe) 72 | 73 | return &resolvr{cancelFunc: cancel}, nil 74 | } 75 | 76 | // Scheme returns the scheme supported by this resolver. 77 | // Scheme is defined at https://github.com/grpc/grpc/blob/master/doc/naming.md. 78 | func (b *builder) Scheme() string { 79 | return schemeName 80 | } 81 | -------------------------------------------------------------------------------- /zrpc/registry/polaris/constant.go: -------------------------------------------------------------------------------- 1 | package polaris 2 | 3 | const ( 4 | 5 | // schemeName for the urls 6 | // All target URLs like 'polaris://.../...' will be resolved by this resolver 7 | schemeName = "polaris" 8 | ) 9 | -------------------------------------------------------------------------------- /zrpc/registry/polaris/go.mod: -------------------------------------------------------------------------------- 1 | module github.com/zeromicro/zero-contrib/zrpc/registry/polaris 2 | 3 | go 1.15 4 | 5 | require ( 6 | github.com/pkg/errors v0.9.1 7 | github.com/polarismesh/polaris-go v1.5.3 8 | github.com/zeromicro/go-zero v1.5.4 9 | google.golang.org/grpc v1.58.0 10 | ) 11 | -------------------------------------------------------------------------------- /zrpc/registry/polaris/options.go: -------------------------------------------------------------------------------- 1 | package polaris 2 | 3 | const ( 4 | allEths = "0.0.0.0" 5 | envPodIP = "POD_IP" 6 | ) 7 | 8 | // options 9 | type Options struct { 10 | ListenOn string 11 | Namespace string 12 | ServiceToken string 13 | ServiceName string 14 | Weight float64 15 | Protocol string 16 | Version string 17 | HeartbeatInervalSec int 18 | Metadata map[string]string 19 | } 20 | 21 | type Option func(*Options) 22 | 23 | func NewPolarisConfig(listenOn string, opts ...Option) *Options { 24 | options := &Options{ 25 | ListenOn: listenOn, 26 | Namespace: "default", 27 | Protocol: "zrpc", 28 | Version: "1.0.0", 29 | HeartbeatInervalSec: 5, 30 | Metadata: make(map[string]string), 31 | } 32 | 33 | for _, opt := range opts { 34 | opt(options) 35 | } 36 | 37 | return options 38 | } 39 | 40 | func WithHeartbeatInervalSec(heartbeatInervalSec int) Option { 41 | return func(o *Options) { 42 | o.HeartbeatInervalSec = heartbeatInervalSec 43 | } 44 | } 45 | 46 | func WithWeight(weight float64) Option { 47 | return func(o *Options) { 48 | o.Weight = weight 49 | } 50 | } 51 | 52 | func WithNamespace(namespace string) Option { 53 | return func(o *Options) { 54 | o.Namespace = namespace 55 | } 56 | } 57 | 58 | func WithServiceName(serviceName string) Option { 59 | return func(o *Options) { 60 | o.ServiceName = serviceName 61 | } 62 | } 63 | 64 | func WithVersion(version string) Option { 65 | return func(o *Options) { 66 | o.Version = version 67 | } 68 | } 69 | 70 | func WithProtocol(protocol string) Option { 71 | return func(o *Options) { 72 | o.Protocol = protocol 73 | } 74 | } 75 | 76 | func WithMetadata(metadata map[string]string) Option { 77 | return func(o *Options) { 78 | o.Metadata = metadata 79 | } 80 | } 81 | -------------------------------------------------------------------------------- /zrpc/registry/polaris/register.go: -------------------------------------------------------------------------------- 1 | package polaris 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "log" 7 | "net" 8 | "os" 9 | "strconv" 10 | "strings" 11 | "sync" 12 | "time" 13 | 14 | "github.com/polarismesh/polaris-go/api" 15 | "github.com/zeromicro/go-zero/core/logx" 16 | "github.com/zeromicro/go-zero/core/netx" 17 | "github.com/zeromicro/go-zero/core/proc" 18 | ) 19 | 20 | var ( 21 | provider api.ProviderAPI 22 | beatlock *sync.RWMutex = &sync.RWMutex{} 23 | hearts map[string]context.CancelFunc = make(map[string]context.CancelFunc) 24 | ) 25 | 26 | func init() { 27 | p, err := api.NewProviderAPI() 28 | if err != nil { 29 | panic(err) 30 | } 31 | 32 | provider = p 33 | } 34 | 35 | // RegisterService register service to polaris 36 | func RegisterService(opts *Options) error { 37 | pubListenOn := figureOutListenOn(opts.ListenOn) 38 | 39 | host, ports, err := net.SplitHostPort(pubListenOn) 40 | if err != nil { 41 | return fmt.Errorf("failed parsing address error: %v", err) 42 | } 43 | port, _ := strconv.ParseInt(ports, 10, 64) 44 | 45 | if err != nil { 46 | log.Panic(err) 47 | } 48 | 49 | req := &api.InstanceRegisterRequest{} 50 | req.Service = opts.ServiceName 51 | req.ServiceToken = opts.ServiceToken 52 | req.Namespace = opts.Namespace 53 | req.Version = &opts.Version 54 | req.Protocol = &opts.Protocol 55 | req.Host = host 56 | req.Port = int(port) 57 | req.TTL = &opts.HeartbeatInervalSec 58 | 59 | resp, err := provider.Register(req) 60 | if err != nil { 61 | return err 62 | } 63 | 64 | ctx, cancel := context.WithCancel(context.Background()) 65 | beatlock.Lock() 66 | hearts[fmt.Sprintf("%s_%s_%s", opts.Namespace, opts.ServiceName, opts.ListenOn)] = cancel 67 | beatlock.Unlock() 68 | 69 | go doHeartbeat(ctx, req, opts, resp.InstanceID) 70 | 71 | addShutdownListener(req, opts, resp.InstanceID) 72 | return nil 73 | } 74 | 75 | func addShutdownListener(registerReq *api.InstanceRegisterRequest, opts *Options, instanceID string) { 76 | // service deregister 77 | proc.AddShutdownListener(func() { 78 | beatlock.Lock() 79 | cancel := hearts[fmt.Sprintf("%s_%s_%s", opts.Namespace, opts.ServiceName, opts.ListenOn)] 80 | beatlock.Unlock() 81 | cancel() 82 | 83 | req := &api.InstanceDeRegisterRequest{} 84 | req.Namespace = opts.Namespace 85 | req.Service = opts.ServiceName 86 | req.ServiceToken = opts.ServiceToken 87 | req.InstanceID = instanceID 88 | req.Host = registerReq.Host 89 | req.Port = registerReq.Port 90 | provider.Deregister(req) 91 | }) 92 | } 93 | 94 | // doHeartbeat 95 | func doHeartbeat(ctx context.Context, req *api.InstanceRegisterRequest, opts *Options, instanceID string) { 96 | ticker := time.NewTicker(time.Duration(opts.HeartbeatInervalSec * int(time.Second))) 97 | for { 98 | select { 99 | case <-ctx.Done(): 100 | return 101 | case <-ticker.C: 102 | beatreq := &api.InstanceHeartbeatRequest{} 103 | beatreq.Namespace = opts.Namespace 104 | beatreq.Service = opts.ServiceName 105 | beatreq.ServiceToken = opts.ServiceToken 106 | beatreq.InstanceID = instanceID 107 | beatreq.Host = req.Host 108 | beatreq.Port = req.Port 109 | 110 | if err := provider.Heartbeat(beatreq); err != nil { 111 | logx.Errorf("[Polaris provider] do heartbeat fail : %s, req : %#v", err.Error(), beatreq) 112 | } 113 | } 114 | } 115 | } 116 | 117 | func figureOutListenOn(listenOn string) string { 118 | fields := strings.Split(listenOn, ":") 119 | if len(fields) == 0 { 120 | return listenOn 121 | } 122 | 123 | host := fields[0] 124 | if len(host) > 0 && host != allEths { 125 | return listenOn 126 | } 127 | 128 | ip := os.Getenv(envPodIP) 129 | if len(ip) == 0 { 130 | ip = netx.InternalIp() 131 | } 132 | if len(ip) == 0 { 133 | return listenOn 134 | } 135 | 136 | return strings.Join(append([]string{ip}, fields[1:]...), ":") 137 | } 138 | -------------------------------------------------------------------------------- /zrpc/registry/polaris/resolver.go: -------------------------------------------------------------------------------- 1 | package polaris 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "sort" 7 | "time" 8 | 9 | "github.com/polarismesh/polaris-go/api" 10 | "github.com/polarismesh/polaris-go/pkg/model" 11 | "github.com/zeromicro/go-zero/core/logx" 12 | "google.golang.org/grpc/resolver" 13 | ) 14 | 15 | type resolvr struct { 16 | cancelFunc context.CancelFunc 17 | } 18 | 19 | func (r *resolvr) ResolveNow(resolver.ResolveNowOptions) {} 20 | 21 | // Close closes the resolver. 22 | func (r *resolvr) Close() { 23 | r.cancelFunc() 24 | } 25 | 26 | type polarisServiceWatcher struct { 27 | out chan<- []string 28 | } 29 | 30 | func newWatcher(out chan<- []string) *polarisServiceWatcher { 31 | return &polarisServiceWatcher{ 32 | out: out, 33 | } 34 | } 35 | 36 | func (watcher *polarisServiceWatcher) startWatch(ctx context.Context, consumer api.ConsumerAPI, subscribeParam *api.WatchServiceRequest) { 37 | for { 38 | resp, err := consumer.WatchService(subscribeParam) 39 | if err != nil { 40 | time.Sleep(time.Duration(500 * time.Millisecond)) 41 | continue 42 | } 43 | 44 | instances := resp.GetAllInstancesResp.Instances 45 | ee := make([]string, len(instances)+1) 46 | for i := range instances { 47 | ins := instances[i] 48 | ee[i] = fmt.Sprintf("%s:%d", ins.GetHost(), ins.GetPort()) 49 | } 50 | if len(ee) != 0 { 51 | watcher.out <- ee 52 | } 53 | 54 | logx.Infof("[Polaris resolver] Watch has been start, param : %#v", subscribeParam) 55 | 56 | select { 57 | case <-ctx.Done(): 58 | logx.Info("[Polaris resolver] Watch has been finished") 59 | return 60 | case event := <-resp.EventChannel: 61 | eType := event.GetSubScribeEventType() 62 | if eType == api.EventInstance { 63 | var insEvent, ok = event.(*model.InstanceEvent) 64 | if !ok { 65 | logx.Errorf("event not `*model.InstanceEvent`") 66 | continue 67 | } 68 | if insEvent == nil { 69 | logx.Errorf("insEvent is nil") 70 | continue 71 | } 72 | 73 | var ( 74 | insAddrList []string 75 | insCount int 76 | ) 77 | if insEvent.AddEvent != nil { 78 | insCount += len(insEvent.AddEvent.Instances) 79 | } 80 | if insEvent.UpdateEvent != nil { 81 | insCount += len(insEvent.UpdateEvent.UpdateList) 82 | } 83 | insAddrList = make([]string, insCount) 84 | 85 | if insEvent.AddEvent != nil { 86 | for _, s := range insEvent.AddEvent.Instances { 87 | insAddrList = append(insAddrList, fmt.Sprintf("%s:%d", s.GetHost(), s.GetPort())) 88 | } 89 | } 90 | if insEvent.UpdateEvent != nil { 91 | for _, s := range insEvent.UpdateEvent.UpdateList { 92 | insAddrList = append(insAddrList, fmt.Sprintf("%s:%d", s.After.GetHost(), s.After.GetPort())) 93 | } 94 | } 95 | 96 | if len(insAddrList) != 0 { 97 | watcher.out <- insAddrList 98 | } 99 | } 100 | } 101 | } 102 | } 103 | 104 | // populateEndpoints 105 | func populateEndpoints(ctx context.Context, clientConn resolver.ClientConn, input <-chan []string) { 106 | for { 107 | select { 108 | case cc := <-input: 109 | connsSet := make(map[string]struct{}, len(cc)) 110 | for _, c := range cc { 111 | connsSet[c] = struct{}{} 112 | } 113 | conns := make([]resolver.Address, 0, len(connsSet)) 114 | for c := range connsSet { 115 | conns = append(conns, resolver.Address{Addr: c}) 116 | } 117 | sort.Sort(byAddressString(conns)) // Don't replace the same address list in the balancer 118 | _ = clientConn.UpdateState(resolver.State{Addresses: conns}) 119 | case <-ctx.Done(): 120 | logx.Info("[Polaris resolver] Watch has been finished") 121 | return 122 | } 123 | } 124 | } 125 | 126 | // byAddressString sorts resolver.Address by Address Field sorting in increasing order. 127 | type byAddressString []resolver.Address 128 | 129 | func (p byAddressString) Len() int { return len(p) } 130 | func (p byAddressString) Less(i, j int) bool { return p[i].Addr < p[j].Addr } 131 | func (p byAddressString) Swap(i, j int) { p[i], p[j] = p[j], p[i] } 132 | -------------------------------------------------------------------------------- /zrpc/registry/polaris/target.go: -------------------------------------------------------------------------------- 1 | package polaris 2 | 3 | import ( 4 | "fmt" 5 | "net/url" 6 | "strings" 7 | "time" 8 | 9 | "github.com/pkg/errors" 10 | "github.com/zeromicro/go-zero/core/mapping" 11 | ) 12 | 13 | type target struct { 14 | Addr string `key:",optional"` 15 | Service string `key:",optional"` 16 | Namespace string `key:"namespace,optional"` 17 | Timeout time.Duration `key:"timeout,optional"` 18 | } 19 | 20 | // parseURL with parameters 21 | func parseURL(u string) (target, error) { 22 | rawURL, err := url.Parse(u) 23 | if err != nil { 24 | return target{}, errors.Wrap(err, "Malformed URL") 25 | } 26 | 27 | fmt.Printf("raw url : %s\n", rawURL) 28 | 29 | if rawURL.Scheme != schemeName || 30 | len(rawURL.Host) == 0 || len(strings.TrimLeft(rawURL.Path, "/")) == 0 { 31 | return target{}, 32 | errors.Errorf("Malformed URL('%s'). Must be in the next format: 'polaris://[user:passwd]@host/service?param=value'", u) 33 | } 34 | 35 | tgt := target{ 36 | Timeout: time.Duration(500 * time.Millisecond), 37 | Namespace: "default", 38 | } 39 | params := make(map[string]interface{}, len(rawURL.Query())) 40 | for name, value := range rawURL.Query() { 41 | params[name] = value[0] 42 | } 43 | 44 | err = mapping.UnmarshalKey(params, &tgt) 45 | if err != nil { 46 | return target{}, errors.Wrap(err, "Malformed URL parameters") 47 | } 48 | 49 | if tgt.Namespace == "" { 50 | tgt.Namespace = "default" 51 | } 52 | 53 | tgt.Addr = rawURL.Host 54 | tgt.Service = strings.TrimLeft(rawURL.Path, "/") 55 | 56 | fmt.Printf("tgt : %#v\n", tgt) 57 | return tgt, nil 58 | } 59 | --------------------------------------------------------------------------------