feat(cache): unified cache abstraction with in-memory and NATS KV backends (#2952)

Signed-off-by: Miguel Martinez Trivino <miguel@chainloop.dev>

Author: migmartri

.gitignore

@@ -55,7 +55,7 @@ values.local.yaml
 deployment/chainloop/charts/**/*.tgz
 
 # Local development overrides
-app/controlplane/configs/config.local.yaml
+config.local.yaml
 devel/dex/config.dev.local.yaml
 devel/compose.override.yml
-.claude/worktrees/
+.claude/worktrees/

app/controlplane/cmd/wire.go

@@ -21,10 +21,13 @@
 package main
 
 import (
+	"time"
+
 	conf "github.com/chainloop-dev/chainloop/app/controlplane/internal/conf/controlplane/config/v1"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/dispatcher"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/server"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/service"
+	"github.com/chainloop-dev/chainloop/app/controlplane/internal/usercontext/entities"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/auditor"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/authz"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/biz"
@@ -33,9 +36,12 @@ import (
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/policies"
 	"github.com/chainloop-dev/chainloop/app/controlplane/plugins/sdk/v1"
 	"github.com/chainloop-dev/chainloop/pkg/blobmanager/loader"
+	"github.com/chainloop-dev/chainloop/pkg/cache"
 	"github.com/chainloop-dev/chainloop/pkg/credentials"
 	"github.com/go-kratos/kratos/v2/log"
+	"github.com/golang-jwt/jwt/v4"
 	"github.com/google/wire"
+	"github.com/nats-io/nats.go"
 )
 
 func wireApp(*conf.Bootstrap, credentials.ReaderWriter, log.Logger, sdk.AvailablePlugins) (*app, func(), error) {
@@ -60,6 +66,7 @@ func wireApp(*conf.Bootstrap, credentials.ReaderWriter, log.Logger, sdk.Availabl
 			newDataConf,
 			newPolicyProviderConfig,
 			newNatsConnection,
+			cacheProviderSet,
 			auditor.NewAuditLogPublisher,
 			newCASServerOptions,
 			newAuthAllowList,
@@ -127,3 +134,48 @@ func newCASServerOptions(in *conf.Bootstrap_CASServer) *biz.CASServerDefaultOpts
 func newAuthAllowList(conf *conf.Bootstrap) *pkgConf.AllowList {
 	return conf.Auth.GetAllowList()
 }
+
+var cacheProviderSet = wire.NewSet(
+	newMembershipsCache,
+	newClaimsCache,
+)
+
+func newClaimsCache(conn *nats.Conn, logger log.Logger) (cache.Cache[*jwt.MapClaims], error) {
+	l := log.NewHelper(logger)
+	backend := "memory"
+	opts := []cache.Option{cache.WithTTL(10 * time.Second), cache.WithLogger(&kratosLogAdapter{h: l})}
+	if conn != nil {
+		backend = "nats"
+		opts = append(opts, cache.WithNATS(conn, "chainloop-jwt-claims"))
+	}
+	l.Infow("msg", "cache initialized", "bucket", "chainloop-jwt-claims", "backend", backend, "ttl", "10s")
+	return cache.New[*jwt.MapClaims](opts...)
+}
+
+func newMembershipsCache(conn *nats.Conn, logger log.Logger) (cache.Cache[*entities.Membership], error) {
+	l := log.NewHelper(logger)
+	backend := "memory"
+	opts := []cache.Option{cache.WithTTL(time.Second), cache.WithLogger(&kratosLogAdapter{h: l})}
+	if conn != nil {
+		backend = "nats"
+		opts = append(opts, cache.WithNATS(conn, "chainloop-memberships"))
+	}
+	l.Infow("msg", "cache initialized", "bucket", "chainloop-memberships", "backend", backend, "ttl", "1s")
+	return cache.New[*entities.Membership](opts...)
+}
+
+// kratosLogAdapter adapts kratos log.Helper (Debugw(...interface{})) to cache.Logger (Debugw(string, ...any)).
+type kratosLogAdapter struct{ h *log.Helper }
+
+func (a *kratosLogAdapter) Debugw(msg string, keyvals ...any) {
+	a.h.Debugw(append([]any{"msg", msg}, keyvals...)...)
+}
+func (a *kratosLogAdapter) Infow(msg string, keyvals ...any) {
+	a.h.Infow(append([]any{"msg", msg}, keyvals...)...)
+}
+func (a *kratosLogAdapter) Warnw(msg string, keyvals ...any) {
+	a.h.Warnw(append([]any{"msg", msg}, keyvals...)...)
+}
+func (a *kratosLogAdapter) Errorw(msg string, keyvals ...any) {
+	a.h.Errorw(append([]any{"msg", msg}, keyvals...)...)
+}

app/controlplane/cmd/wire_gen.go

@@ -25,9 +25,11 @@ import (
 	conf "github.com/chainloop-dev/chainloop/app/controlplane/internal/conf/controlplane/config/v1"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/sentrycontext"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/usercontext/attjwtmiddleware"
+	"github.com/chainloop-dev/chainloop/app/controlplane/internal/usercontext/entities"
 	authzMiddleware "github.com/chainloop-dev/chainloop/app/controlplane/pkg/authz/middleware"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/biz"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/jwt/user"
+	"github.com/chainloop-dev/chainloop/pkg/cache"
 
 	"buf.build/go/protovalidate"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/service"
@@ -61,6 +63,8 @@ type Opts struct {
 	OrganizationUseCase *biz.OrganizationUseCase
 	WorkflowUseCase     *biz.WorkflowUseCase
 	MembershipUseCase   *biz.MembershipUseCase
+	MembershipsCache    cache.Cache[*entities.Membership]
+	ClaimsCache         cache.Cache[*jwt.MapClaims]
 	// Services
 	WorkflowSvc         *service.WorkflowService
 	AuthSvc             *service.AuthService
@@ -203,7 +207,7 @@ func craftMiddleware(opts *Opts) []middleware.Middleware {
 			// 2.c - Set its user
 			usercontext.WithCurrentUserMiddleware(opts.UserUseCase, logHelper),
 			// Store all memberships in the context
-			usercontext.WithCurrentMembershipsMiddleware(opts.MembershipUseCase),
+			usercontext.WithCurrentMembershipsMiddleware(opts.MembershipUseCase, opts.MembershipsCache),
 			selector.Server(
 				// 2.d- Set its organization
 				usercontext.WithCurrentOrganizationMiddleware(opts.UserUseCase, opts.OrganizationUseCase, logHelper),
@@ -235,6 +239,7 @@ func craftMiddleware(opts *Opts) []middleware.Middleware {
 				attjwtmiddleware.NewUserTokenProvider(opts.AuthConfig.GeneratedJwsHmacSecret),
 				// Delegated Federated provider
 				attjwtmiddleware.WithFederatedProvider(opts.FederatedConfig),
+				attjwtmiddleware.WithClaimsCache(opts.ClaimsCache),
 			),
 			// 2.a - Set its workflow and organization in the context
 			usercontext.WithAttestationContextFromRobotAccount(opts.RobotAccountUseCase, opts.OrganizationUseCase, logHelper),
@@ -245,7 +250,7 @@ func craftMiddleware(opts *Opts) []middleware.Middleware {
 			// 2.d - Set its robot account from federated delegation
 			usercontext.WithAttestationContextFromFederatedInfo(opts.OrganizationUseCase, logHelper),
 			// Store all memberships in the context
-			usercontext.WithCurrentMembershipsMiddleware(opts.MembershipUseCase),
+			usercontext.WithCurrentMembershipsMiddleware(opts.MembershipUseCase, opts.MembershipsCache),
 			// 3 - Update API Token last usage
 			usercontext.WithAPITokenUsageUpdater(opts.APITokenUseCase, logHelper),
 			// 4 - Validate the CAS Backend is fully configured and valid

app/controlplane/internal/server/grpc.go

@@ -1,5 +1,5 @@
 //
-// Copyright 2024-2025 The Chainloop Authors.
+// Copyright 2024-2026 The Chainloop Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -34,6 +34,7 @@ import (
 	casJWT "github.com/chainloop-dev/chainloop/internal/robotaccount/cas"
 	"github.com/chainloop-dev/chainloop/pkg/attestation"
 	"github.com/chainloop-dev/chainloop/pkg/attestation/renderer/chainloop"
+	"github.com/chainloop-dev/chainloop/pkg/cache"
 	"github.com/chainloop-dev/chainloop/pkg/credentials"
 	errors "github.com/go-kratos/kratos/v2/errors"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
@@ -62,6 +63,7 @@ type AttestationService struct {
 	signingUseCase          *biz.SigningUseCase
 	userUseCase             *biz.UserUseCase
 	bootstrapConfig         *conf.Bootstrap
+	membershipsCache        cache.Cache[*entities.Membership]
 }
 
 type NewAttestationServiceOpts struct {
@@ -83,6 +85,7 @@ type NewAttestationServiceOpts struct {
 	SigningUseCase     *biz.SigningUseCase
 	UserUC             *biz.UserUseCase
 	BootstrapConfig    *conf.Bootstrap
+	MembershipsCache   cache.Cache[*entities.Membership]
 	Opts               []NewOpt
 }
 
@@ -107,6 +110,7 @@ func NewAttestationService(opts *NewAttestationServiceOpts) *AttestationService
 		signingUseCase:          opts.SigningUseCase,
 		userUseCase:             opts.UserUC,
 		bootstrapConfig:         opts.BootstrapConfig,
+		membershipsCache:        opts.MembershipsCache,
 	}
 }
 
@@ -797,7 +801,9 @@ func (s *AttestationService) FindOrCreateWorkflow(ctx context.Context, req *cpAP
 	}
 
 	// reset RBAC cache, since we might have created a new project
-	usercontext.ResetMembershipsCache()
+	if s.membershipsCache != nil {
+		_ = s.membershipsCache.Purge(ctx)
+	}
 
 	return &cpAPI.FindOrCreateWorkflowResponse{Result: bizWorkflowToPb(wf)}, nil
 }

app/controlplane/internal/service/attestation.go

@@ -1,5 +1,5 @@
 //
-// Copyright 2024-2025 The Chainloop Authors.
+// Copyright 2024-2026 The Chainloop Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -24,19 +24,18 @@ import (
 	"io"
 	"net/http"
 	"strings"
-	"time"
 
 	v1 "github.com/chainloop-dev/chainloop/app/controlplane/api/controlplane/v1"
 	conf "github.com/chainloop-dev/chainloop/app/controlplane/internal/conf/controlplane/config/v1"
 	"github.com/chainloop-dev/chainloop/app/controlplane/internal/usercontext/entities"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/jwt/apitoken"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/jwt/robotaccount"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/jwt/user"
+	"github.com/chainloop-dev/chainloop/pkg/cache"
 	errorsAPI "github.com/go-kratos/kratos/v2/errors"
 	"github.com/go-kratos/kratos/v2/log"
 	"github.com/go-kratos/kratos/v2/middleware"
 	"github.com/go-kratos/kratos/v2/transport"
-	"github.com/hashicorp/golang-lru/v2/expirable"
 
 	"github.com/golang-jwt/jwt/v4"
 )
@@ -199,6 +198,12 @@ func WithVerifyAudienceFunc(f VerifyAudienceFunc) TokenProviderOption {
 type options struct {
 	tokenProviders   []providerOption
 	federatedAuthURL string
+	claimsCache      cache.Cache[*jwt.MapClaims]
+}
+
+// WithClaimsCache sets an external cache for federated JWT claims.
+func WithClaimsCache(c cache.Cache[*jwt.MapClaims]) JWTOption {
+	return func(o *options) { o.claimsCache = c }
 }
 
 func withTokenProvider(providerKey string, opts ...TokenProviderOption) JWTOption {
@@ -251,9 +256,6 @@ func WithJWTMulti(l log.Logger, opts ...JWTOption) middleware.Middleware {
 		logger.Infof("federated authentication enabled, using URL: %s", o.federatedAuthURL)
 	}
 
-	// claims cache with 10s TTL and unlimited keys
-	claimsCache := expirable.NewLRU[string, *jwt.MapClaims](0, nil, time.Second*10)
-
 	return func(handler middleware.Handler) middleware.Handler {
 		return func(ctx context.Context, req interface{}) (interface{}, error) {
 			if header, ok := transport.FromServerContext(ctx); ok {
@@ -286,7 +288,7 @@ func WithJWTMulti(l log.Logger, opts ...JWTOption) middleware.Middleware {
 							}
 
 							logger.Infof("calling federated provider, orgName: %s", orgName)
-							claims, err := callFederatedProvider(o.federatedAuthURL, jwtToken, orgName, claimsCache)
+							claims, err := callFederatedProvider(ctx, o.federatedAuthURL, jwtToken, orgName, o.claimsCache)
 							if err != nil {
 								// if we receive an error from upstream we want to expose it to the user, for example if the federated provider
 								// is saying that the token is invalid
@@ -333,9 +335,9 @@ func WithJWTMulti(l log.Logger, opts ...JWTOption) middleware.Middleware {
 
 // callFederatedProvider calls the federated provider to verify the token
 // it returns the claims of the token if the token is valid and verified
-func callFederatedProvider(verifyURL string, jwtToken, orgName string, cache *expirable.LRU[string, *jwt.MapClaims]) (*jwt.MapClaims, error) {
+func callFederatedProvider(ctx context.Context, verifyURL string, jwtToken, orgName string, claimsCache cache.Cache[*jwt.MapClaims]) (*jwt.MapClaims, error) {
 	cacheKey := fmt.Sprintf("%s:%s", jwtToken, orgName)
-	if claims, ok := cache.Get(cacheKey); ok {
+	if claims, ok, _ := claimsCache.Get(ctx, cacheKey); ok {
 		return claims, nil
 	}
 
@@ -391,7 +393,7 @@ func callFederatedProvider(verifyURL string, jwtToken, orgName string, cache *ex
 		"orgName":    response.OrgName,
 	}
 
-	cache.Add(cacheKey, claims)
+	_ = claimsCache.Set(ctx, cacheKey, claims)
 
 	return claims, nil
 }