Repository: go-kratos/kratos
Pull request #1817 (2022-02-14 06:33:33)

fix(jwt): parse server custom claims
<!--
🎉 Thanks for sending a pull request to Kratos! Here are some tips for you:

1. If this is your first time contributing to Kratos, please read our contribution guide: https://go-kratos.dev/en/docs/community/contribution/
2. Ensure you have added or ran the appropriate tests and lint for your PR, please use `make lint` and `make test` before filing your PR, use `make clean` to tidy your go mod.
3. If the PR is unfinished, you may need mark it as a WIP(Work In Progress) PR or draft PR
4. Please use a conventional commits format title: `<type>[optional scope]: <description>`
    
    Some suggestion on <type>:

    fix: A bug fix
    feat: A new feature
    test: Adding missing tests or correcting existing tests
    refactor: A code change that neither fixes a bug nor adds a feature
    break: Changes has break change

    docs: Documentation only changes
    deps: Changes external dependencies
    style: Changes that do not affect the meaning of the code (white-space, formatting, etc)
    chore Daily work, examples, etc.
    ci: Changes to our CI configuration files and scripts
-->

#### Description (what this PR does / why we need it):
<!--
* The description should include the motivation for this PR or contrast this with previous behavior
-->


#### Which issue(s) this PR fixes (resolves / be part of):
<!--
* Automatically closes linked issue when PR is merged.
* If you PR is not fully resolved issue, please use `part of #<issue number>` instead.

Usage: `fixes/resolves #<issue number>`, or `fixes/resolves (paste link of issue)`.
-->
fixes  #1720 #1815

#### Other special notes for reviewer:
<!--
* Somethings that need extra attention for reviewer
* Some additional notes, TODO list, etc.
-->


CryBecase: 或许可以考虑改成传一个 CopyableClaims 来代替 func() jwt.Claims？
```go
type Copyable interface {
	Copy() jwt.Claims
	jwt.Claims
}

// WithServerClaims with server claim
func WithServerClaims(claims Copyable) Option {
	return func(o *options) {
		o.serverClaims = func() jwt.Claims {
			return claims.Copy()
		}
	}
}
```
Casper-Mars: > 或许可以考虑改成传一个 CopyableClaims 来代替 func() jwt.Claims？
> 
> ```go
> type Copyable interface {
> 	Copy() jwt.Claims
> 	jwt.Claims
> }
> 
> // WithServerClaims with server claim
> func WithServerClaims(claims Copyable) Option {
> 	return func(o *options) {
> 		o.serverClaims = func() jwt.Claims {
> 			return claims.Copy()
> 		}
> 	}
> }
> ```

Copy()只是为了产生一个空的claim，并没有实际拷贝属性。这样使用建造者builder或者工厂factory会好一点吧？形如提供的代码那样。
CryBecase: > > 或许可以考虑改成传一个 CopyableClaims 来代替 func() jwt.Claims？
> > ```go
> > type Copyable interface {
> > 	Copy() jwt.Claims
> > 	jwt.Claims
> > }
> > 
> > // WithServerClaims with server claim
> > func WithServerClaims(claims Copyable) Option {
> > 	return func(o *options) {
> > 		o.serverClaims = func() jwt.Claims {
> > 			return claims.Copy()
> > 		}
> > 	}
> > }
> > ```
> 
> Copy()只是为了产生一个空的claim，并没有实际拷贝属性。这样使用建造者builder或者工厂factory会好一点吧？形如提交的代码那样。

确实是这样的，已经改成工厂模式了。
haiyux: use `make fix` to fix lint.
shenqidebaozi: @kagaya85 PTAL.
shenqidebaozi: @Casper-Mars PTAL.
codecov-commenter: # [Codecov](https://codecov.io/gh/go-kratos/kratos/pull/1817?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) Report
> Merging [#1817](https://codecov.io/gh/go-kratos/kratos/pull/1817?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) (5fbdc27) into [main](https://codecov.io/gh/go-kratos/kratos/commit/6ec695064b00cf7a79f4e022dae1bf1b9fb2f35e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) (6ec6950) will **decrease** coverage by `0.08%`.
> The diff coverage is `100.00%`.

> :exclamation: Current head 5fbdc27 differs from pull request most recent head 64c8b75. Consider uploading reports for the commit 64c8b75 to get more accurate results

[![Impacted file tree graph](https://codecov.io/gh/go-kratos/kratos/pull/1817/graphs/tree.svg?width=650&height=150&src=pr&token=XwnmBTz78x&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos)](https://codecov.io/gh/go-kratos/kratos/pull/1817?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos)

```diff
@@            Coverage Diff             @@
##             main    #1817      +/-   ##
==========================================
- Coverage   80.80%   80.71%   -0.09%     
==========================================
  Files          80       80              
  Lines        3558     3568      +10     
==========================================
+ Hits         2875     2880       +5     
- Misses        477      481       +4     
- Partials      206      207       +1     
```


| [Impacted Files](https://codecov.io/gh/go-kratos/kratos/pull/1817?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) | Coverage Δ | |
|---|---|---|
| [middleware/auth/jwt/jwt.go](https://codecov.io/gh/go-kratos/kratos/pull/1817/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos#diff-bWlkZGxld2FyZS9hdXRoL2p3dC9qd3QuZ28=) | `87.50% <100.00%> (+1.78%)` | :arrow_up: |
| [selector/node/ewma/node.go](https://codecov.io/gh/go-kratos/kratos/pull/1817/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos#diff-c2VsZWN0b3Ivbm9kZS9ld21hL25vZGUuZ28=) | `83.87% <0.00%> (-3.23%)` | :arrow_down: |
| [internal/context/context.go](https://codecov.io/gh/go-kratos/kratos/pull/1817/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos#diff-aW50ZXJuYWwvY29udGV4dC9jb250ZXh0Lmdv) | `97.01% <0.00%> (-2.99%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/go-kratos/kratos/pull/1817?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/go-kratos/kratos/pull/1817?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos). Last update [6ec6950...64c8b75](https://codecov.io/gh/go-kratos/kratos/pull/1817?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos).


Patch:
diff --git a/contrib/registry/nacos/registry.go b/contrib/registry/nacos/registry.go
index baf9b473..088cb8c2 100644
--- a/contrib/registry/nacos/registry.go
+++ b/contrib/registry/nacos/registry.go
@@ -8,6 +8,7 @@ import (
 	"strconv"
 
 	"github.com/nacos-group/nacos-sdk-go/clients/naming_client"
+	"github.com/nacos-group/nacos-sdk-go/common/constant"
 	"github.com/nacos-group/nacos-sdk-go/vo"
 
 	"github.com/go-kratos/kratos/v2/registry"
@@ -65,7 +66,7 @@ func New(cli naming_client.INamingClient, opts ...Option) (r *Registry) {
 	op := options{
 		prefix:  "/microservices",
 		cluster: "DEFAULT",
-		group:   "DEFAULT_GROUP",
+		group:   constant.DEFAULT_GROUP,
 		weight:  100,
 		kind:    "grpc",
 	}
@@ -167,6 +168,7 @@ func (r *Registry) Watch(ctx context.Context, serviceName string) (registry.Watc
 func (r *Registry) GetService(_ context.Context, serviceName string) ([]*registry.ServiceInstance, error) {
 	res, err := r.cli.SelectInstances(vo.SelectInstancesParam{
 		ServiceName: serviceName,
+		GroupName:   r.opts.group,
 		HealthyOnly: true,
 	})
 	if err != nil {
diff --git a/middleware/auth/jwt/jwt.go b/middleware/auth/jwt/jwt.go
index c7a7f78f..60dad2db 100644
--- a/middleware/auth/jwt/jwt.go
+++ b/middleware/auth/jwt/jwt.go
@@ -48,7 +48,7 @@ type Option func(*options)
 // Parser is a jwt parser
 type options struct {
 	signingMethod jwt.SigningMethod
-	claims        jwt.Claims
+	claims        func() jwt.Claims
 	tokenHeader   map[string]interface{}
 }
 
@@ -60,9 +60,11 @@ func WithSigningMethod(method jwt.SigningMethod) Option {
 }
 
 // WithClaims with customer claim
-func WithClaims(claims jwt.Claims) Option {
+// If you use it in Server, f needs to return a new jwt.Claims object each time to avoid concurrent write problems
+// If you use it in Client, f only needs to return a single object to provide performance
+func WithClaims(f func() jwt.Claims) Option {
 	return func(o *options) {
-		o.claims = claims
+		o.claims = f
 	}
 }
 
@@ -77,7 +79,6 @@ func WithTokenHeader(header map[string]interface{}) Option {
 func Server(keyFunc jwt.Keyfunc, opts ...Option) middleware.Middleware {
 	o := &options{
 		signingMethod: jwt.SigningMethodHS256,
-		claims:        jwt.RegisteredClaims{},
 	}
 	for _, opt := range opts {
 		opt(o)
@@ -93,7 +94,15 @@ func Server(keyFunc jwt.Keyfunc, opts ...Option) middleware.Middleware {
 					return nil, ErrMissingJwtToken
 				}
 				jwtToken := auths[1]
-				tokenInfo, err := jwt.Parse(jwtToken, keyFunc)
+				var (
+					tokenInfo *jwt.Token
+					err       error
+				)
+				if o.claims != nil {
+					tokenInfo, err = jwt.ParseWithClaims(jwtToken, o.claims(), keyFunc)
+				} else {
+					tokenInfo, err = jwt.Parse(jwtToken, keyFunc)
+				}
 				if err != nil {
 					if ve, ok := err.(*jwt.ValidationError); ok {
 						if ve.Errors&jwt.ValidationErrorMalformed != 0 {
@@ -120,9 +129,10 @@ func Server(keyFunc jwt.Keyfunc, opts ...Option) middleware.Middleware {
 
 // Client is a client jwt middleware.
 func Client(keyProvider jwt.Keyfunc, opts ...Option) middleware.Middleware {
+	claims := jwt.RegisteredClaims{}
 	o := &options{
 		signingMethod: jwt.SigningMethodHS256,
-		claims:        jwt.RegisteredClaims{},
+		claims:        func() jwt.Claims { return claims },
 	}
 	for _, opt := range opts {
 		opt(o)
@@ -132,7 +142,7 @@ func Client(keyProvider jwt.Keyfunc, opts ...Option) middleware.Middleware {
 			if keyProvider == nil {
 				return nil, ErrNeedTokenProvider
 			}
-			token := jwt.NewWithClaims(o.signingMethod, o.claims)
+			token := jwt.NewWithClaims(o.signingMethod, o.claims())
 			if o.tokenHeader != nil {
 				for k, v := range o.tokenHeader {
 					token.Header[k] = v


Test patch:
diff --git a/middleware/auth/jwt/jwt_test.go b/middleware/auth/jwt/jwt_test.go
index 7944d503..188744b0 100644
--- a/middleware/auth/jwt/jwt_test.go
+++ b/middleware/auth/jwt/jwt_test.go
@@ -4,8 +4,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math/rand"
 	"net/http"
 	"reflect"
+	"strconv"
+	"sync"
 	"testing"
 	"time"
 
@@ -63,6 +66,107 @@ func (tr *Transport) ReplyHeader() transport.Header {
 	return nil
 }
 
+type CustomerClaims struct {
+	Name string `json:"name"`
+	jwt.RegisteredClaims
+}
+
+func TestJWTServerParse(t *testing.T) {
+	var (
+		errConcurrentWrite = errors.New("concurrent write claims")
+		errParseClaims     = errors.New("bad result, token claims is not CustomerClaims")
+	)
+
+	testKey := "testKey"
+	tests := []struct {
+		name         string
+		token        func() string
+		claims       func() jwt.Claims
+		exceptErr    error
+		key          string
+		goroutineNum int
+	}{
+		{
+			name: "normal",
+			token: func() string {
+				token, err := jwt.NewWithClaims(jwt.SigningMethodHS256, &CustomerClaims{}).SignedString([]byte(testKey))
+				if err != nil {
+					panic(err)
+				}
+				return fmt.Sprintf(bearerFormat, token)
+			},
+			claims: func() jwt.Claims {
+				return &CustomerClaims{}
+			},
+			exceptErr:    nil,
+			key:          testKey,
+			goroutineNum: 1,
+		},
+		{
+			name: "concurrent request",
+			token: func() string {
+				token, err := jwt.NewWithClaims(jwt.SigningMethodHS256, &CustomerClaims{
+					Name: strconv.Itoa(rand.Int()),
+				}).SignedString([]byte(testKey))
+				if err != nil {
+					panic(err)
+				}
+				return fmt.Sprintf(bearerFormat, token)
+			},
+			claims: func() jwt.Claims {
+				return &CustomerClaims{}
+			},
+			exceptErr:    nil,
+			key:          testKey,
+			goroutineNum: 10000,
+		},
+	}
+
+	next := func(ctx context.Context, req interface{}) (interface{}, error) {
+		testToken, _ := FromContext(ctx)
+		var name string
+		if customerClaims, ok := testToken.(*CustomerClaims); ok {
+			name = customerClaims.Name
+		} else {
+			return nil, errParseClaims
+		}
+
+		// mock biz
+		time.Sleep(100 * time.Millisecond)
+
+		if customerClaims, ok := testToken.(*CustomerClaims); ok {
+			if name != customerClaims.Name {
+				return nil, errConcurrentWrite
+			}
+		} else {
+			return nil, errParseClaims
+		}
+		return "reply", nil
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			server := Server(
+				func(token *jwt.Token) (interface{}, error) { return []byte(testKey), nil },
+				WithClaims(test.claims),
+			)(next)
+			wg := sync.WaitGroup{}
+			for i := 0; i < test.goroutineNum; i++ {
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+					ctx := transport.NewServerContext(context.Background(), &Transport{reqHeader: newTokenHeader(authorizationKey, test.token())})
+					_, err2 := server(ctx, test.name)
+					if !errors.Is(test.exceptErr, err2) {
+						t.Errorf("except error %v, but got %v", test.exceptErr, err2)
+					}
+				}()
+			}
+			wg.Wait()
+		})
+	}
+}
+
 func TestServer(t *testing.T) {
 	testKey := "testKey"
 	mapClaims := jwt.MapClaims{}
@@ -279,6 +383,7 @@ func TestClientWithClaims(t *testing.T) {
 	testKey := "testKey"
 	mapClaims := jwt.MapClaims{}
 	mapClaims["name"] = "xiaoli"
+	mapClaimsFunc := func() jwt.Claims { return mapClaims }
 	claims := jwt.NewWithClaims(jwt.SigningMethodHS256, mapClaims)
 	token, err := claims.SignedString([]byte(testKey))
 	if err != nil {
@@ -301,7 +406,7 @@ func TestClientWithClaims(t *testing.T) {
 		next := func(ctx context.Context, req interface{}) (interface{}, error) {
 			return "reply", nil
 		}
-		handler := Client(test.tokenProvider, WithClaims(mapClaims))(next)
+		handler := Client(test.tokenProvider, WithClaims(mapClaimsFunc))(next)
 		header := &headerCarrier{}
 		_, err2 := handler(transport.NewClientContext(context.Background(), &Transport{reqHeader: header}), "ok")
 		if !errors.Is(test.expectError, err2) {
@@ -319,6 +424,7 @@ func TestClientWithHeader(t *testing.T) {
 	testKey := "testKey"
 	mapClaims := jwt.MapClaims{}
 	mapClaims["name"] = "xiaoli"
+	mapClaimsFunc := func() jwt.Claims { return mapClaims }
 	tokenHeader := map[string]interface{}{
 		"test": "test",
 	}
@@ -336,7 +442,7 @@ func TestClientWithHeader(t *testing.T) {
 	next := func(ctx context.Context, req interface{}) (interface{}, error) {
 		return "reply", nil
 	}
-	handler := Client(tProvider, WithClaims(mapClaims), WithTokenHeader(tokenHeader))(next)
+	handler := Client(tProvider, WithClaims(mapClaimsFunc), WithTokenHeader(tokenHeader))(next)
 	header := &headerCarrier{}
 	_, err2 := handler(transport.NewClientContext(context.Background(), &Transport{reqHeader: header}), "ok")
 	if err2 != nil {
@@ -351,6 +457,7 @@ func TestClientMissKey(t *testing.T) {
 	testKey := "testKey"
 	mapClaims := jwt.MapClaims{}
 	mapClaims["name"] = "xiaoli"
+	mapClaimsFunc := func() jwt.Claims { return mapClaims }
 	claims := jwt.NewWithClaims(jwt.SigningMethodHS256, mapClaims)
 	token, err := claims.SignedString([]byte(testKey))
 	if err != nil {
@@ -373,7 +480,7 @@ func TestClientMissKey(t *testing.T) {
 		next := func(ctx context.Context, req interface{}) (interface{}, error) {
 			return "reply", nil
 		}
-		handler := Client(test.tokenProvider, WithClaims(mapClaims))(next)
+		handler := Client(test.tokenProvider, WithClaims(mapClaimsFunc))(next)
 		header := &headerCarrier{}
 		_, err2 := handler(transport.NewClientContext(context.Background(), &Transport{reqHeader: header}), "ok")
 		if !errors.Is(test.expectError, err2) {