@@ -5,8 +5,10 @@ import (
55 "encoding/json"
66 "fmt"
77 "log/slog"
8+ "net/http"
89 "slices"
910 "strings"
11+ "sync"
1012 "time"
1113
1214 "github.com/Layr-Labs/eigenx-kms-go/pkg/types"
@@ -90,21 +92,46 @@ type AttestationVerifier struct {
9092 intelJwksCache jwk.Set
9193 projectID string
9294 debugMode bool
95+ cancel context.CancelFunc
96+ httpTransport * http.Transport
97+ closeOnce sync.Once
9398}
9499
95100func NewAttestationVerifier (ctx context.Context , logger * slog.Logger , projectID string , refreshInterval time.Duration , debugMode bool ) (* AttestationVerifier , error ) {
96101 avLogger := logger .With ("component" , "attestation_verifier" )
97102 avLogger .Debug ("Initializing attestation verifier" , "project_id" , projectID , "refresh_interval" , refreshInterval )
98103
104+ // Derive a cancellable context owned by the verifier so JWK cache workers
105+ // (httprc/v3) can be shut down via Close() in addition to parent cancellation.
106+ cacheCtx , cancel := context .WithCancel (ctx )
107+
108+ // Own the HTTP transport so Close() can tear down idle connections and the
109+ // associated read/write loop goroutines. Using DefaultClient (httprc's default)
110+ // leaks persistent HTTP/2 read loops past context cancellation.
111+ // Guard the type assertion: an application can install a custom
112+ // http.DefaultTransport (e.g. an instrumented transport) that is not a
113+ // *http.Transport, and blind assertion would panic.
114+ var transport * http.Transport
115+ if dt , ok := http .DefaultTransport .(* http.Transport ); ok {
116+ transport = dt .Clone ()
117+ } else {
118+ transport = & http.Transport {}
119+ }
120+ httpClient := & http.Client {Transport : transport }
121+
99122 avLogger .Debug ("Creating Google Confidential Space JWK cache" , "jwk_url" , confidentialSpaceJWKURL )
100- googleJwksCache , err := NewJWKCache ( ctx , confidentialSpaceJWKURL , refreshInterval )
123+ googleJwksCache , err := newJWKCacheWithClient ( cacheCtx , confidentialSpaceJWKURL , refreshInterval , httpClient )
101124 if err != nil {
125+ cancel ()
126+ transport .CloseIdleConnections ()
102127 return nil , fmt .Errorf ("failed to create Google JWK cache: %w" , err )
103128 }
104129
105130 avLogger .Debug ("Creating Intel Trust Authority JWK cache" , "jwk_url" , intelTrustAuthorityJWKURL )
106- intelJwksCache , err := NewJWKCache ( ctx , intelTrustAuthorityJWKURL , refreshInterval )
131+ intelJwksCache , err := newJWKCacheWithClient ( cacheCtx , intelTrustAuthorityJWKURL , refreshInterval , httpClient )
107132 if err != nil {
133+ cancel ()
134+ transport .CloseIdleConnections ()
108135 return nil , fmt .Errorf ("failed to create Intel JWK cache: %w" , err )
109136 }
110137
@@ -116,9 +143,25 @@ func NewAttestationVerifier(ctx context.Context, logger *slog.Logger, projectID
116143 googleJwksCache : googleJwksCache ,
117144 intelJwksCache : intelJwksCache ,
118145 debugMode : debugMode ,
146+ cancel : cancel ,
147+ httpTransport : transport ,
119148 }, nil
120149}
121150
151+ // Close stops the background JWK cache workers spawned by the verifier and
152+ // releases pooled HTTP connections. Safe to call multiple times and safe to
153+ // call concurrently from multiple goroutines.
154+ func (av * AttestationVerifier ) Close () {
155+ av .closeOnce .Do (func () {
156+ if av .cancel != nil {
157+ av .cancel ()
158+ }
159+ if av .httpTransport != nil {
160+ av .httpTransport .CloseIdleConnections ()
161+ }
162+ })
163+ }
164+
122165func (av * AttestationVerifier ) VerifyAttestation (ctx context.Context , tokenString string , provider AttestationProvider ) (* types.AttestationClaims , error ) {
123166 av .logger .Debug ("Starting attestation verification" , "token_length" , len (tokenString ), "provider" , provider )
124167
@@ -364,14 +407,34 @@ func extractAppIDFromInstanceName(instanceName string) (string, error) {
364407 return instanceNameParts [len (instanceNameParts )- 1 ], nil
365408}
366409
410+ // NewJWKCache creates a JWK set backed by a periodic refresh worker.
411+ //
412+ // The background worker goroutines and their HTTP idle connections are tied
413+ // to ctx — cancel it to stop them. Callers that need deterministic teardown
414+ // of pooled HTTP connections (e.g. short-lived tests or a short-lived verifier
415+ // lifecycle separate from process shutdown) should not use this function
416+ // directly; instead they should manage the HTTP transport themselves and use
417+ // AttestationVerifier, which closes idle connections on Close().
367418func NewJWKCache (ctx context.Context , jwkUrl string , refreshInterval time.Duration ) (jwk.Set , error ) {
419+ return newJWKCacheWithClient (ctx , jwkUrl , refreshInterval , nil )
420+ }
421+
422+ // newJWKCacheWithClient is like NewJWKCache but lets the caller supply the HTTP
423+ // client used to fetch the JWK set. Passing a non-nil client lets the caller
424+ // own its transport and close idle connections when the cache is no longer
425+ // needed, preventing leaked HTTP read/write loop goroutines.
426+ func newJWKCacheWithClient (ctx context.Context , jwkUrl string , refreshInterval time.Duration , httpClient * http.Client ) (jwk.Set , error ) {
368427 cache , err := jwk .NewCache (ctx , httprc .NewClient ())
369428 if err != nil {
370429 return nil , fmt .Errorf ("failed to create jwk cache: %w" , err )
371430 }
372431
373432 // register a constant refresh interval for this URL.
374- err = cache .Register (ctx , jwkUrl , jwk .WithConstantInterval (refreshInterval ))
433+ registerOpts := []jwk.RegisterOption {jwk .WithConstantInterval (refreshInterval )}
434+ if httpClient != nil {
435+ registerOpts = append (registerOpts , jwk .WithHTTPClient (httpClient ))
436+ }
437+ err = cache .Register (ctx , jwkUrl , registerOpts ... )
375438 if err != nil {
376439 return nil , fmt .Errorf ("failed to register jwk location: %w" , err )
377440 }
0 commit comments