derived token cache

AthenZ · Sep 18, 2024 · ceb4ff1 · ceb4ff1
1 parent 965ab91
commit ceb4ff1
Show file tree

Hide file tree

Showing 8 changed files with 50 additions and 81 deletions.
diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -108,7 +108,7 @@ func (idCfg *IdentityConfig) loadFromENV() error {
 	loadEnv("TOKEN_SERVER_TLS_CA_PATH", &idCfg.TokenServerTLSCAPath)
 	loadEnv("TOKEN_SERVER_TLS_CERT_PATH", &idCfg.TokenServerTLSCertPath)
 	loadEnv("TOKEN_SERVER_TLS_KEY_PATH", &idCfg.TokenServerTLSKeyPath)
-	loadEnv("TOKEN_DIR", &idCfg.TokenDir)
+	loadEnv("TOKEN_DIR", &idCfg.tokenDir)
 	loadEnv("METRICS_SERVER_ADDR", &idCfg.MetricsServerAddr)
 	loadEnv("DELETE_INSTANCE_ID", &idCfg.rawDeleteInstanceID)
 	loadEnv("USE_TOKEN_SERVER", &idCfg.rawUseTokenServer)
@@ -142,11 +142,11 @@ func (idCfg *IdentityConfig) loadFromENV() error {
 	if err != nil {
 		return fmt.Errorf("Invalid ROLE_CERT_OUTPUT_KEY_FILE [%q], %w", idCfg.rawRoleCertKeyFileOutput, err)
 	}
-	idCfg.TokenRefresh, err = time.ParseDuration(idCfg.rawTokenRefresh)
+	idCfg.tokenRefresh, err = time.ParseDuration(idCfg.rawTokenRefresh)
 	if err != nil {
 		return fmt.Errorf("Invalid TOKEN_REFRESH_INTERVAL [%q], %w", idCfg.rawTokenRefresh, err)
 	}
-	idCfg.TokenExpiry, err = time.ParseDuration(idCfg.rawTokenExpiry)
+	idCfg.tokenExpiry, err = time.ParseDuration(idCfg.rawTokenExpiry)
 	if err != nil {
 		return fmt.Errorf("Invalid TOKEN_EXPIRY [%q], %w", idCfg.rawTokenExpiry, err)
 	}
@@ -209,15 +209,15 @@ func (idCfg *IdentityConfig) loadFromFlag(program string, args []string) error {
 	f.StringVar(&idCfg.roleCertKeyNamingFormat, "role-cert-key-naming-format", idCfg.roleCertKeyNamingFormat, "The file name format when outputting the role cert key to a file")
 	// RoleAuthHeader
 	f.StringVar(&idCfg.TokenType, "token-type", idCfg.TokenType, "type of the role token to request (\"roletoken\", \"accesstoken\" or \"roletoken+accesstoken\")")
-	f.DurationVar(&idCfg.TokenRefresh, "token-refresh-interval", idCfg.TokenRefresh, "token refresh interval")
-	f.DurationVar(&idCfg.TokenExpiry, "token-expiry", idCfg.TokenExpiry, "token expiry duration (0 to use Athenz server's default expiry)")
+	f.DurationVar(&idCfg.tokenRefresh, "token-refresh-interval", idCfg.tokenRefresh, "token refresh interval")
+	f.DurationVar(&idCfg.tokenExpiry, "token-expiry", idCfg.tokenExpiry, "token expiry duration (0 to use Athenz server's default expiry)")
 	f.StringVar(&idCfg.TokenServerAddr, "token-server-addr", idCfg.TokenServerAddr, "HTTP server address to provide tokens (required for token provisioning)")
 	f.BoolVar(&idCfg.TokenServerRESTAPI, "token-server-rest-api", idCfg.TokenServerRESTAPI, "enable token server RESTful API (true/false)")
 	f.DurationVar(&idCfg.TokenServerTimeout, "token-server-timeout", idCfg.TokenServerTimeout, "token server timeout (default 3s)")
 	f.StringVar(&idCfg.TokenServerTLSCAPath, "token-server-tls-ca-path", idCfg.TokenServerTLSCAPath, "token server TLS CA path (if set, enable TLS Client Authentication)")
 	f.StringVar(&idCfg.TokenServerTLSCertPath, "token-server-tls-cert-path", idCfg.TokenServerTLSCertPath, "token server TLS certificate path (if empty, disable TLS)")
 	f.StringVar(&idCfg.TokenServerTLSKeyPath, "token-server-tls-key-path", idCfg.TokenServerTLSKeyPath, "token server TLS certificate key path (if empty, disable TLS)")
-	f.StringVar(&idCfg.TokenDir, "token-dir", idCfg.TokenDir, "directory to write token files")
+	f.StringVar(&idCfg.tokenDir, "token-dir", idCfg.tokenDir, "directory to write token files")
 	f.StringVar(&idCfg.MetricsServerAddr, "metrics-server-addr", idCfg.MetricsServerAddr, "HTTP server address to provide metrics")
 	f.BoolVar(&idCfg.DeleteInstanceID, "delete-instance-id", idCfg.DeleteInstanceID, "delete x509 certificate record from identity provider on shutdown (true/false)")
 	// Token Server
@@ -249,8 +249,8 @@ func (idCfg *IdentityConfig) parseRawValues() (err error) {
 
 func (idCfg *IdentityConfig) validateAndInit() (err error) {
 
-	if idCfg.TokenExpiry != 0 && idCfg.TokenRefresh >= idCfg.TokenExpiry {
-		return fmt.Errorf("Invalid TokenRefresh[%s] >= TokenExpiry[%s]", idCfg.TokenRefresh.String(), idCfg.TokenExpiry.String())
+	if idCfg.tokenExpiry != 0 && idCfg.tokenRefresh >= idCfg.tokenExpiry {
+		return fmt.Errorf("Invalid TokenRefresh[%s] >= TokenExpiry[%s]", idCfg.tokenRefresh.String(), idCfg.tokenExpiry.String())
 	}
 
 	// TODO: clarify unused logic

diff --git a/pkg/config/default.go b/pkg/config/default.go
@@ -90,23 +90,22 @@ func DefaultIdentityConfig() *IdentityConfig {
 		PodName:                   "",
 		Reloader:                  nil,
 		ServerCACert:              "",
-		TokenTargetDomainRoles:    []DomainRole{},
 		roleCertDir:               "",
 		roleCertFilenameDelimiter: DEFAULT_ROLE_CERT_FILENAME_DELIMITER,
 		roleCertKeyFileOutput:     false,
 		roleCertNamingFormat:      "",
 		roleCertKeyNamingFormat:   "",
 		RoleAuthHeader:            DEFAULT_ROLE_AUTH_HEADER,
 		TokenType:                 "accesstoken",
-		TokenRefresh:              DEFAULT_TOKEN_REFRESH,
-		TokenExpiry:               DEFAULT_TOKEN_EXPIRY,
+		tokenDir:                  "",
+		tokenRefresh:              DEFAULT_TOKEN_REFRESH,
+		tokenExpiry:               DEFAULT_TOKEN_EXPIRY,
 		TokenServerAddr:           "",
 		TokenServerRESTAPI:        false,
 		TokenServerTimeout:        DEFAULT_TOKEN_SERVER_TIMEOUT,
 		TokenServerTLSCAPath:      "",
 		TokenServerTLSCertPath:    "",
 		TokenServerTLSKeyPath:     "",
-		TokenDir:                  "",
 		MetricsServerAddr:         "",
 		HealthCheckAddr:           "",
 		HealthCheckEndpoint:       "",

diff --git a/pkg/config/derived-access-token.go b/pkg/config/derived-access-token.go
diff --git a/pkg/config/derived-target-domain-roles.go b/pkg/config/derived-target-domain-roles.go
@@ -24,7 +24,7 @@ import (
 
 type DerivedTargetDomainRoles struct {
 	roleCerts []DomainRole // private as the derived state is used only within the config package
-	// tokens    []DomainRole // private as the derived state is used only within the config package
+	tokens    []DomainRole // private as the derived state is used only within the config package
 }
 
 // derivedTargetDomainRoles sets the DerivedTargetDomainRoles with the given rawTargetDomainRoles.
@@ -63,10 +63,9 @@ func (idCfg *IdentityConfig) derivedTargetDomainRoles() error {
 		})
 	}
 
-	idCfg.TokenTargetDomainRoles = tokenDomainRoles // TODO: Delete me and refactor by using the type DerivedTargetDomainRoles below:
 	idCfg.targetDomainRoles = DerivedTargetDomainRoles{
 		roleCerts: roleCertDomainRoles,
-		// tokens: tokenDomainRoles,
+		tokens:    tokenDomainRoles,
 	}
 
 	return nil

diff --git a/pkg/config/dervied.go b/pkg/config/dervied.go
@@ -39,12 +39,11 @@ func (idCfg *IdentityConfig) loadDerivedConfig() error {
 		return err
 	}
 
-	// TODO:
 	// depends on the following:
 	// - derivedTargetDomainRoles()
-	// if err := idCfg.derivedAccessTokenConfig(); err != nil {
-	// 	return err
-	// }
+	if err := idCfg.derivedTokenChacheConfig(); err != nil {
+		return err
+	}
 
 	// TODO:
 	// depends on the following:

diff --git a/pkg/config/model.go b/pkg/config/model.go
@@ -49,7 +49,6 @@ type IdentityConfig struct {
 	ServerCACert           string
 	K8sSecretBackup        DerivedK8sSecretBackup
 	ServiceCert            DerivedServiceCert
-	TokenTargetDomainRoles []DomainRole             // TODO: Will be migrated into DerivedTargetDomainRoles
 	targetDomainRoles      DerivedTargetDomainRoles // private as the derived state is used only within the config package
 	// RoleCerts Derived State and its related fields:
 	RoleCert                  DerivedRoleCert
@@ -59,17 +58,20 @@ type IdentityConfig struct {
 	roleCertNamingFormat      string
 	roleCertKeyNamingFormat   string
 	//
+	// Token Cache Derived State and its related fields:
+	TokenCache   DerivedTokenCache
+	tokenDir     string
+	tokenRefresh time.Duration
+	tokenExpiry  time.Duration
+	//
 	RoleAuthHeader         string
 	TokenType              string
-	TokenRefresh           time.Duration
-	TokenExpiry            time.Duration
 	TokenServerAddr        string
 	TokenServerRESTAPI     bool
 	TokenServerTimeout     time.Duration
 	TokenServerTLSCAPath   string
 	TokenServerTLSCertPath string
 	TokenServerTLSKeyPath  string
-	TokenDir               string
 	MetricsServerAddr      string
 	HealthCheckAddr        string
 	HealthCheckEndpoint    string

diff --git a/pkg/token/server.go b/pkg/token/server.go
@@ -90,8 +90,8 @@ func postRoleToken(ts *tokenService, w http.ResponseWriter, r *http.Request) {
 	// if rtRequest.MaxExpiry != nil && *rtRequest.MaxExpiry > 0{
 	// 	k.MaxExpiry = *rtRequest.MaxExpiry
 	// }
-	if k.MinExpiry == 0 && ts.tokenExpiryInSecond > 0 {
-		k.MinExpiry = ts.tokenExpiryInSecond
+	if k.MinExpiry == 0 && ts.idCfg.TokenCache.ExpirySeconds > 0 {
+		k.MinExpiry = ts.idCfg.TokenCache.ExpirySeconds
 	}
 
 	// cache lookup (token TTL must >= 1 minute)
@@ -173,8 +173,8 @@ func postAccessToken(ts *tokenService, w http.ResponseWriter, r *http.Request) {
 	if atRequest.Expiry != nil && *atRequest.Expiry > 0 {
 		k.MaxExpiry = *atRequest.Expiry
 	}
-	if k.MaxExpiry == 0 && ts.tokenExpiryInSecond > 0 {
-		k.MaxExpiry = ts.tokenExpiryInSecond
+	if k.MaxExpiry == 0 && ts.idCfg.TokenCache.ExpirySeconds > 0 {
+		k.MaxExpiry = ts.idCfg.TokenCache.ExpirySeconds
 	}
 
 	// cache lookup (token TTL must >= 1 minute)
@@ -259,7 +259,7 @@ func newHandlerFunc(ts *tokenService, timeout time.Duration) http.Handler {
 		} else {
 			// TODO: Since the specifications are not yet decided, the value of WriteFileRequired is undetermined.
 			// TODO: Maybe we need to separate the cache keys for RT and AT?
-			k := CacheKey{Domain: domain, Role: role, MinExpiry: ts.tokenExpiryInSecond}
+			k := CacheKey{Domain: domain, Role: role, MinExpiry: ts.idCfg.TokenCache.ExpirySeconds}
 			if ts.tokenType&mACCESS_TOKEN != 0 {
 				k, aToken = ts.accessTokenCache.Search(k)
 				if aToken == nil {

diff --git a/pkg/token/service.go b/pkg/token/service.go
@@ -51,10 +51,8 @@ type tokenService struct {
 	ztsClient *zts.ZTSClient
 	saService string
 
-	idCfg *config.IdentityConfig
-	// TODO: move to derived token
-	tokenExpiryInSecond int
-	tokenType           mode
+	idCfg     *config.IdentityConfig
+	tokenType mode
 
 	tokenServer        *http.Server
 	tokenServerRunning bool
@@ -66,25 +64,27 @@ func New(ctx context.Context, idCfg *config.IdentityConfig) (daemon.Daemon, erro
 		return nil, nil
 	}
 	// TODO: In the next PR, the determination will be made on a per Access Token and Role Token basis.
-	enableWriteFiles := idCfg.TokenDir != ""
+	enableAccessToken := idCfg.TokenCache.AccessToken.Use
+	enableRoleToken := idCfg.TokenCache.RoleToken.Use
+	enableWriteFiles := enableAccessToken || enableRoleToken
 	if !enableWriteFiles {
-		log.Debugf("Skipping to write token files to directory with empty TOKEN_DIR [%s]", idCfg.TokenDir)
+		// The Dir settings for the access token and role token are the same.
+		log.Debugf("Skipping to write token files to directory with empty TOKEN_DIR [%s]", idCfg.TokenCache.AccessToken.Dir)
 	}
 
 	// initialize token cache with placeholder
 	tt := newType(idCfg.TokenType)
-	tokenExpiryInSecond := int(idCfg.TokenExpiry.Seconds())
 	accessTokenCache := NewLockedTokenCache("accesstoken", idCfg.Namespace, idCfg.PodName)
 	roleTokenCache := NewLockedTokenCache("roletoken", idCfg.Namespace, idCfg.PodName)
-	for _, dr := range idCfg.TokenTargetDomainRoles {
+	for _, dr := range idCfg.TokenCache.TargetDomainRoles {
 		domain, role := dr.Domain, dr.Role
 		// TODO: Rewrite the following if statement as "if tt.isAccessTokenEnabled()..."
 		if tt&mACCESS_TOKEN != 0 {
-			accessTokenCache.Store(CacheKey{Domain: domain, Role: role, MaxExpiry: tokenExpiryInSecond, WriteFileRequired: enableWriteFiles}, &AccessToken{})
+			accessTokenCache.Store(CacheKey{Domain: domain, Role: role, MaxExpiry: idCfg.TokenCache.ExpirySeconds, WriteFileRequired: enableWriteFiles}, &AccessToken{})
 		}
 		// TODO: Rewrite the following if statement as "if tt.isRoleTokenEnabled()..."
 		if tt&mROLE_TOKEN != 0 {
-			roleTokenCache.Store(CacheKey{Domain: domain, Role: role, MinExpiry: tokenExpiryInSecond, WriteFileRequired: enableWriteFiles}, &RoleToken{})
+			roleTokenCache.Store(CacheKey{Domain: domain, Role: role, MinExpiry: idCfg.TokenCache.ExpirySeconds, WriteFileRequired: enableWriteFiles}, &RoleToken{})
 		}
 	}
 
@@ -110,14 +110,13 @@ func New(ctx context.Context, idCfg *config.IdentityConfig) (daemon.Daemon, erro
 
 	// setup token service
 	ts := &tokenService{
-		shutdownChan:        make(chan struct{}, 1),
-		accessTokenCache:    accessTokenCache,
-		roleTokenCache:      roleTokenCache,
-		ztsClient:           ztsClient,
-		saService:           saService,
-		idCfg:               idCfg,
-		tokenType:           tt,
-		tokenExpiryInSecond: tokenExpiryInSecond,
+		shutdownChan:     make(chan struct{}, 1),
+		accessTokenCache: accessTokenCache,
+		roleTokenCache:   roleTokenCache,
+		ztsClient:        ztsClient,
+		saService:        saService,
+		idCfg:            idCfg,
+		tokenType:        tt,
 	}
 
 	// write tokens as files only if it is non-init mode OR TOKEN_DIR is set
@@ -189,15 +188,15 @@ func (ts *tokenService) Start(ctx context.Context) error {
 	}
 
 	// refreshes tokens periodically
-	if ts.idCfg.TokenRefresh > 0 {
-		t := time.NewTicker(ts.idCfg.TokenRefresh)
+	if ts.idCfg.TokenCache.Refresh > 0 {
+		t := time.NewTicker(ts.idCfg.TokenCache.Refresh)
 		ts.shutdownWg.Add(1)
 		go func() {
 			defer t.Stop()
 			defer ts.shutdownWg.Done()
 
 			for {
-				log.Infof("Will refresh cached tokens within %s", ts.idCfg.TokenRefresh.String())
+				log.Infof("Will refresh cached tokens within %s", ts.idCfg.TokenCache.Refresh.String())
 
 				select {
 				case <-ts.shutdownChan:
@@ -211,7 +210,7 @@ func (ts *tokenService) Start(ctx context.Context) error {
 					}
 
 					// backoff retry until TOKEN_REFRESH_INTERVAL / 4 OR context is done
-					for _, err := range ts.updateTokenCachesAndWriteFiles(ctx, ts.idCfg.TokenRefresh/4) {
+					for _, err := range ts.updateTokenCachesAndWriteFiles(ctx, ts.idCfg.TokenCache.Refresh/4) {
 						log.Errorf("Failed to refresh tokens after multiple retries: %s", err.Error())
 					}
 				}
@@ -394,7 +393,7 @@ func (d *tokenService) updateAndWriteFileToken(key CacheKey, tt mode) error {
 		// File output processing
 		domain, role := key.Domain, key.Role
 		token := d.accessTokenCache.Load(key)
-		outPath := filepath.Join(d.idCfg.TokenDir, domain+":role."+role+".accesstoken")
+		outPath := filepath.Join(d.idCfg.TokenCache.AccessToken.Dir, domain+d.idCfg.TokenCache.AccessToken.Delimiter+role+".accesstoken")
 		return d.writeFile(token, outPath, mACCESS_TOKEN)
 	}
 	updateAndWriteFileRoleToken := func(key CacheKey) error {
@@ -405,7 +404,7 @@ func (d *tokenService) updateAndWriteFileToken(key CacheKey, tt mode) error {
 		// File output processing
 		domain, role := key.Domain, key.Role
 		token := d.roleTokenCache.Load(key)
-		outPath := filepath.Join(d.idCfg.TokenDir, domain+":role."+role+".roletoken")
+		outPath := filepath.Join(d.idCfg.TokenCache.RoleToken.Dir, domain+d.idCfg.TokenCache.RoleToken.Delimiter+role+".roletoken")
 		return d.writeFile(token, outPath, mROLE_TOKEN)
 	}
 	switch tt {