Repository: go-kratos/kratos
Pull request #1752 (2022-01-06 15:03:15)

feat(middleware/auth/jwt): add customer header
#### Description (what this PR does / why we need it):
Allow user put a customer header of JWT token while using the client auth middleware 





haiyux: 可以试着升一下jwt版本
codecov-commenter: # [Codecov](https://codecov.io/gh/go-kratos/kratos/pull/1752?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) Report
> Merging [#1752](https://codecov.io/gh/go-kratos/kratos/pull/1752?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) (271dbc9) into [main](https://codecov.io/gh/go-kratos/kratos/commit/da70e22201155541619269ab4f52fd1cd49a15c3?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) (da70e22) will **increase** coverage by `0.03%`.
> The diff coverage is `85.71%`.

> :exclamation: Current head 271dbc9 differs from pull request most recent head 189fb3b. Consider uploading reports for the commit 189fb3b to get more accurate results
[![Impacted file tree graph](https://codecov.io/gh/go-kratos/kratos/pull/1752/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/1752?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    #1752      +/-   ##
==========================================
+ Coverage   81.83%   81.86%   +0.03%     
==========================================
  Files          78       78              
  Lines        3474     3480       +6     
==========================================
+ Hits         2843     2849       +6     
  Misses        427      427              
  Partials      204      204              
```


| [Impacted Files](https://codecov.io/gh/go-kratos/kratos/pull/1752?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/1752/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos#diff-bWlkZGxld2FyZS9hdXRoL2p3dC9qd3QuZ28=) | `85.71% <85.71%> (+1.33%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/go-kratos/kratos/pull/1752?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/1752?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos). Last update [da70e22...189fb3b](https://codecov.io/gh/go-kratos/kratos/pull/1752?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).

Casper-Mars: 


> 可以试着升一下jwt版本

嗯嗯，版本升级打算迟一点再升，还有一个自定义claim的能力在考虑怎么做进来

Patch:
diff --git a/middleware/auth/jwt/jwt.go b/middleware/auth/jwt/jwt.go
index 5c77dc22..59d670ce 100644
--- a/middleware/auth/jwt/jwt.go
+++ b/middleware/auth/jwt/jwt.go
@@ -22,21 +22,24 @@ const (
 	// bearerFormat authorization token format
 	bearerFormat string = "Bearer %s"
 
-	// authorizationKey holds the key used to store the JWT Token in the request header.
+	// authorizationKey holds the key used to store the JWT Token in the request tokenHeader.
 	authorizationKey string = "Authorization"
+
+	// reason holds the error reason.
+	reason string = "UNAUTHORIZED"
 )
 
 var (
-	ErrMissingJwtToken        = errors.Unauthorized("UNAUTHORIZED", "JWT token is missing")
-	ErrMissingKeyFunc         = errors.Unauthorized("UNAUTHORIZED", "keyFunc is missing")
-	ErrTokenInvalid           = errors.Unauthorized("UNAUTHORIZED", "Token is invalid")
-	ErrTokenExpired           = errors.Unauthorized("UNAUTHORIZED", "JWT token has expired")
-	ErrTokenParseFail         = errors.Unauthorized("UNAUTHORIZED", "Fail to parse JWT token ")
-	ErrUnSupportSigningMethod = errors.Unauthorized("UNAUTHORIZED", "Wrong signing method")
-	ErrWrongContext           = errors.Unauthorized("UNAUTHORIZED", "Wrong context for middleware")
-	ErrNeedTokenProvider      = errors.Unauthorized("UNAUTHORIZED", "Token provider is missing")
-	ErrSignToken              = errors.Unauthorized("UNAUTHORIZED", "Can not sign token.Is the key correct?")
-	ErrGetKey                 = errors.Unauthorized("UNAUTHORIZED", "Can not get key while signing token")
+	ErrMissingJwtToken        = errors.Unauthorized(reason, "JWT token is missing")
+	ErrMissingKeyFunc         = errors.Unauthorized(reason, "keyFunc is missing")
+	ErrTokenInvalid           = errors.Unauthorized(reason, "Token is invalid")
+	ErrTokenExpired           = errors.Unauthorized(reason, "JWT token has expired")
+	ErrTokenParseFail         = errors.Unauthorized(reason, "Fail to parse JWT token ")
+	ErrUnSupportSigningMethod = errors.Unauthorized(reason, "Wrong signing method")
+	ErrWrongContext           = errors.Unauthorized(reason, "Wrong context for middleware")
+	ErrNeedTokenProvider      = errors.Unauthorized(reason, "Token provider is missing")
+	ErrSignToken              = errors.Unauthorized(reason, "Can not sign token.Is the key correct?")
+	ErrGetKey                 = errors.Unauthorized(reason, "Can not get key while signing token")
 )
 
 // Option is jwt option.
@@ -46,6 +49,7 @@ type Option func(*options)
 type options struct {
 	signingMethod jwt.SigningMethod
 	claims        jwt.Claims
+	tokenHeader   map[string]interface{}
 }
 
 // WithSigningMethod with signing method option.
@@ -62,6 +66,13 @@ func WithClaims(claims jwt.Claims) Option {
 	}
 }
 
+// WithTokenHeader withe customer tokenHeader for client side
+func WithTokenHeader(header map[string]interface{}) Option {
+	return func(o *options) {
+		o.tokenHeader = header
+	}
+}
+
 // Server is a server auth middleware. Check the token and extract the info from token.
 func Server(keyFunc jwt.Keyfunc, opts ...Option) middleware.Middleware {
 	o := &options{
@@ -93,7 +104,7 @@ func Server(keyFunc jwt.Keyfunc, opts ...Option) middleware.Middleware {
 							return nil, ErrTokenParseFail
 						}
 					}
-					return nil, errors.Unauthorized("UNAUTHORIZED", err.Error())
+					return nil, errors.Unauthorized(reason, err.Error())
 				} else if !tokenInfo.Valid {
 					return nil, ErrTokenInvalid
 				} else if tokenInfo.Method != o.signingMethod {
@@ -122,6 +133,11 @@ func Client(keyProvider jwt.Keyfunc, opts ...Option) middleware.Middleware {
 				return nil, ErrNeedTokenProvider
 			}
 			token := jwt.NewWithClaims(o.signingMethod, o.claims)
+			if o.tokenHeader != nil {
+				for k, v := range o.tokenHeader {
+					token.Header[k] = v
+				}
+			}
 			key, err := keyProvider(token)
 			if err != nil {
 				return nil, ErrGetKey


Test patch:
diff --git a/middleware/auth/jwt/jwt_test.go b/middleware/auth/jwt/jwt_test.go
index 627b585e..3f02cacc 100644
--- a/middleware/auth/jwt/jwt_test.go
+++ b/middleware/auth/jwt/jwt_test.go
@@ -294,6 +294,34 @@ func TestClientWithClaims(t *testing.T) {
 	})
 }
 
+func TestClientWithHeader(t *testing.T) {
+	testKey := "testKey"
+	mapClaims := jwt.MapClaims{}
+	mapClaims["name"] = "xiaoli"
+	tokenHeader := map[string]interface{}{
+		"test": "test",
+	}
+	claims := jwt.NewWithClaims(jwt.SigningMethodHS256, mapClaims)
+	for k, v := range tokenHeader {
+		claims.Header[k] = v
+	}
+	token, err := claims.SignedString([]byte(testKey))
+	if err != nil {
+		panic(err)
+	}
+	tProvider := func(*jwt.Token) (interface{}, error) {
+		return []byte(testKey), nil
+	}
+	next := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return "reply", nil
+	}
+	handler := Client(tProvider, WithClaims(mapClaims), WithTokenHeader(tokenHeader))(next)
+	header := &headerCarrier{}
+	_, err2 := handler(transport.NewClientContext(context.Background(), &Transport{reqHeader: header}), "ok")
+	assert.Equal(t, nil, err2)
+	assert.Equal(t, fmt.Sprintf(bearerFormat, token), header.Get(authorizationKey))
+}
+
 func TestClientMissKey(t *testing.T) {
 	testKey := "testKey"
 	mapClaims := jwt.MapClaims{}