Fix accesstoken cache bugs (#13)

* Fix accesstoken cache bug Signed-off-by: taniwa <taniwa@yahoo-corp.jp> * Fix access.go Signed-off-by: taniwa <taniwa@yahoo-corp.jp> * Fix expiry in refresh Signed-off-by: taniwa <taniwa@yahoo-corp.jp> * Add test Signed-off-by: taniwa <taniwa@yahoo-corp.jp> * Add expiries_in to accessCacheData, fix variable names Signed-off-by: taniwa <taniwa@yahoo-corp.jp> * Add expiresIn as cache request parameter Signed-off-by: taniwa <taniwa@yahoo-corp.jp> * Add scope as cache request parameter Signed-off-by: taniwa <taniwa@yahoo-corp.jp> --------- Signed-off-by: taniwa <taniwa@yahoo-corp.jp>
AthenZ · Aug 23, 2023 · 8db0ea2 · 8db0ea2
1 parent 4aa1153
commit 8db0ea2
Show file tree

Hide file tree

Showing 4 changed files with 228 additions and 106 deletions.
diff --git a/go.mod b/go.mod
@@ -6,6 +6,7 @@ replace github.com/AthenZ/athenz => github.com/AthenZ/athenz v1.11.32
 
 require (
 	github.com/AthenZ/athenz v1.11.32
+	github.com/golang-jwt/jwt/v5 v5.0.0
 	github.com/kpango/fastime v1.1.9
 	github.com/kpango/gache v1.2.8
 	github.com/kpango/glg v1.6.15

diff --git a/go.sum b/go.sum
@@ -6,6 +6,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/dimfeld/httptreemux v5.0.1+incompatible h1:Qj3gVcDNoOthBAqftuD596rm4wg/adLLz5xh5CmpiCA=
 github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
 github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
+github.com/golang-jwt/jwt/v5 v5.0.0 h1:1n1XNM9hk7O9mnQoNBGolZvzebBQ7p93ULHRc28XJUE=
+github.com/golang-jwt/jwt/v5 v5.0.0/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
 github.com/klauspost/cpuid/v2 v2.2.5 h1:0E5MSMDEoAulmXNFquVs//DdoomxaoTY1kUhbc/qbZg=
 github.com/klauspost/cpuid/v2 v2.2.5/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws=
 github.com/kpango/fastime v1.1.9 h1:xVQHcqyPt5M69DyFH7g1EPRns1YQNap9d5eLhl/Jy84=

diff --git a/service/access.go b/service/access.go
@@ -31,6 +31,7 @@ import (
 	"time"
 
 	"github.com/AthenZ/athenz-client-sidecar/v2/config"
+	jwt "github.com/golang-jwt/jwt/v5"
 	"github.com/kpango/fastime"
 	"github.com/kpango/gache"
 	"github.com/kpango/glg"
@@ -66,11 +67,13 @@ type accessService struct {
 }
 
 type accessCacheData struct {
-	token             *AccessTokenResponse
+	token             string
 	domain            string
 	role              string
 	proxyForPrincipal string
-	expiresIn         int64
+	expiresIn         int64  // cache user request parameter
+	expiry            int64  // cache ZTS response
+	scope             string // cache ZTS response
 }
 
 // AccessTokenResponse represents the AccessTokenResponse from postAccessTokenRequest.
@@ -244,11 +247,20 @@ func (a *accessService) GetAccessProvider() AccessProvider {
 // getAccessToken returns AccessTokenResponse struct or error.
 // This function will return the access token stored inside the cache, or fetch the access token from Athenz when corresponding access token cannot be found in the cache.
 func (a *accessService) getAccessToken(ctx context.Context, domain, role, proxyForPrincipal string, expiresIn int64) (*AccessTokenResponse, error) {
-	tok, ok := a.getCache(domain, role, proxyForPrincipal)
+	cd, ok := a.getCache(domain, role, proxyForPrincipal)
 	if !ok {
 		return a.updateAccessToken(ctx, domain, role, proxyForPrincipal, expiresIn)
 	}
-	return tok, nil
+	atResponse := &AccessTokenResponse{
+		AccessToken: cd.token,
+		ExpiresIn:   int64(time.Unix(cd.expiry, 0).Sub(time.Now()).Seconds()),
+		TokenType:   "Bearer", // hardcoded in the same way as ZTS, https://github.com/AthenZ/athenz/blob/a85f48666763759ee28fda114acc4c8d2cafc28e/servers/zts/src/main/java/com/yahoo/athenz/zts/ZTSImpl.java#L2656C10-L2656C10
+	}
+	if cd.scope != "" {
+		atResponse.Scope = cd.scope // set scope ONLY when non-nil & non-empty, https://github.com/AthenZ/athenz/blob/a85f48666763759ee28fda114acc4c8d2cafc28e/core/zts/src/main/java/com/yahoo/athenz/zts/AccessTokenResponse.java#L21C14-L21C14
+	}
+
+	return atResponse, nil
 }
 
 // RefreshAccessTokenCache returns the error channel when it is updated.
@@ -307,15 +319,27 @@ func (a *accessService) updateAccessToken(ctx context.Context, domain, role, pro
 			return nil, e
 		}
 
+		tok, _, err := jwt.NewParser().ParseUnverified(at.AccessToken, &jwt.RegisteredClaims{})
+		if err != nil {
+			return nil, fmt.Errorf("jwt.ParseUnverified() err: %v", err)
+		}
+
+		expTime, err := tok.Claims.GetExpirationTime()
+		if err != nil {
+			return nil, fmt.Errorf("jwt.GetExpirationTime() err: %v", err)
+		}
+
 		a.tokenCache.SetWithExpire(key, &accessCacheData{
-			token:             at,
+			token:             at.AccessToken,
 			domain:            domain,
 			role:              role,
 			proxyForPrincipal: proxyForPrincipal,
 			expiresIn:         expiresIn,
-		}, time.Unix(at.ExpiresIn, 0).Sub(expTimeDelta))
+			expiry:            expTime.Unix(),
+			scope:             at.Scope,
+		}, expTime.Sub(expTimeDelta))
 
-		glg.Debugf("token is cached, domain: %s, role: %s, proxyForPrincipal: %s, expiry time: %v", domain, role, proxyForPrincipal, at.ExpiresIn)
+		glg.Debugf("token is cached, domain: %s, role: %s, proxyForPrincipal: %s, expiry time: %v", domain, role, proxyForPrincipal, expTime.Unix())
 		return at, nil
 	})
 	if err != nil {
@@ -404,12 +428,12 @@ func createScope(domain, role string) string {
 	return domain + ":domain"
 }
 
-func (a *accessService) getCache(domain, role, principal string) (*AccessTokenResponse, bool) {
+func (a *accessService) getCache(domain, role, principal string) (*accessCacheData, bool) {
 	val, ok := a.tokenCache.Get(encode(domain, role, principal))
 	if !ok {
 		return nil, false
 	}
-	return val.(*accessCacheData).token, ok
+	return val.(*accessCacheData), ok
 }
 
 // createGetAccessTokenRequest creates Athenz's postAccessTokenRequest.