From 6388cc8d171b508aefa13314b12264767641a487 Mon Sep 17 00:00:00 2001 From: "Jose I. Paris" Date: Fri, 25 Sep 2026 12:01:29 +0200 Subject: [PATCH] fix(cli): refresh the CAS upload token before it expires The CAS upload token expires after 30 seconds. The CLI requested it before crafting, so a material that takes longer to craft, such as a big AI coding session under secret redaction, failed to upload with a "JWT expired" error. The CAS connection now gets its token from a provider on each call. The CLI caches the token and requests a new one from the control plane when the cached token is close to expiry. Fixes PFM-7428 Assisted-by: Claude Code Signed-off-by: Jose I. Paris Chainloop-Trace-Sessions: a2cde44d-e76c-48c0-8395-818659409bcb --- app/cli/pkg/action/action.go | 18 ++++- app/cli/pkg/action/cas_token.go | 73 ++++++++++++++++++ app/cli/pkg/action/cas_token_test.go | 111 +++++++++++++++++++++++++++ pkg/grpcconn/grpcconn.go | 32 +++++++- pkg/grpcconn/grpcconn_test.go | 69 +++++++++++++++++ 5 files changed, 298 insertions(+), 5 deletions(-) create mode 100644 app/cli/pkg/action/cas_token.go create mode 100644 app/cli/pkg/action/cas_token_test.go diff --git a/app/cli/pkg/action/action.go b/app/cli/pkg/action/action.go index 0d9c7bf19..d8de0205c 100644 --- a/app/cli/pkg/action/action.go +++ b/app/cli/pkg/action/action.go @@ -176,7 +176,23 @@ func getCASBackend(ctx context.Context, client pb.AttestationServiceClient, work } } - artifactCASConn, err := grpcconn.New(casURI, result.Token, opts...) + // The CAS token expires quickly, and crafting can take longer than that. + // Request a new token when the cached one is close to expiry. + tokenSource := newCASTokenSource(result.Token, func(ctx context.Context) (string, error) { + resp, err := client.GetUploadCreds(ctx, &pb.AttestationServiceGetUploadCredsRequest{ + WorkflowRunId: workflowRunID, + }) + if err != nil { + return "", fmt.Errorf("refreshing upload creds: %w", err) + } + if resp.GetResult().GetToken() == "" { + return "", fmt.Errorf("refreshing upload creds: empty token") + } + return resp.GetResult().GetToken(), nil + }) + opts = append(opts, grpcconn.WithTokenProvider(tokenSource.Token)) + + artifactCASConn, err := grpcconn.New(casURI, "", opts...) if err != nil { logger.Warn().Err(err).Msg("failed to create CAS connection, will store inline") return nil, nil, fmt.Errorf("creating CAS connection: %w", err) diff --git a/app/cli/pkg/action/cas_token.go b/app/cli/pkg/action/cas_token.go new file mode 100644 index 000000000..9397c42c5 --- /dev/null +++ b/app/cli/pkg/action/cas_token.go @@ -0,0 +1,73 @@ +// +// Copyright 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. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package action + +import ( + "context" + "sync" + "time" + + "github.com/golang-jwt/jwt/v5" +) + +// casTokenRefreshMargin is the minimum validity a cached CAS token must have +// left to be reused. It leaves time for the request to reach the CAS. +const casTokenRefreshMargin = 10 * time.Second + +// casTokenSource caches the short-lived CAS token and requests a new one from +// the control plane when the cached token is close to expiry. Crafting a +// material (e.g. redacting a big AI coding session) can take longer than the +// token lifetime, so the token must be fresh when the upload starts. +type casTokenSource struct { + mu sync.Mutex + token string + fetch func(ctx context.Context) (string, error) +} + +func newCASTokenSource(initial string, fetch func(ctx context.Context) (string, error)) *casTokenSource { + return &casTokenSource{token: initial, fetch: fetch} +} + +// Token returns the cached token, or a new one if the cached token expires +// within casTokenRefreshMargin +func (s *casTokenSource) Token(ctx context.Context) (string, error) { + s.mu.Lock() + defer s.mu.Unlock() + + if tokenValidFor(s.token) > casTokenRefreshMargin { + return s.token, nil + } + + token, err := s.fetch(ctx) + if err != nil { + return "", err + } + + s.token = token + return token, nil +} + +// tokenValidFor returns how long the token stays valid. The signature is not +// verified: the CAS verifies it, and the CLI only needs the expiration. +// A token that cannot be parsed or has no expiration counts as expired. +func tokenValidFor(token string) time.Duration { + claims := &jwt.RegisteredClaims{} + if _, _, err := jwt.NewParser().ParseUnverified(token, claims); err != nil || claims.ExpiresAt == nil { + return 0 + } + + return time.Until(claims.ExpiresAt.Time) +} diff --git a/app/cli/pkg/action/cas_token_test.go b/app/cli/pkg/action/cas_token_test.go new file mode 100644 index 000000000..d5f58a912 --- /dev/null +++ b/app/cli/pkg/action/cas_token_test.go @@ -0,0 +1,111 @@ +// +// Copyright 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. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package action + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/golang-jwt/jwt/v5" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func signedTestToken(t *testing.T, expiresIn time.Duration) string { + t.Helper() + claims := jwt.RegisteredClaims{ExpiresAt: jwt.NewNumericDate(time.Now().Add(expiresIn))} + token, err := jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString([]byte("test")) + require.NoError(t, err) + return token +} + +func TestCASTokenSource(t *testing.T) { + fresh := signedTestToken(t, time.Minute) + + testCases := []struct { + name string + initial string + fetchErr error + wantToken string + wantFetches int + wantErr bool + }{ + { + name: "valid token is reused", + initial: signedTestToken(t, time.Minute), + wantFetches: 0, + }, + { + name: "expired token is refreshed", + initial: signedTestToken(t, -time.Second), + wantToken: fresh, + wantFetches: 1, + }, + { + name: "token close to expiry is refreshed", + initial: signedTestToken(t, casTokenRefreshMargin/2), + wantToken: fresh, + wantFetches: 1, + }, + { + name: "unparseable token is refreshed", + initial: "not-a-jwt", + wantToken: fresh, + wantFetches: 1, + }, + { + name: "refresh error is returned", + initial: signedTestToken(t, -time.Second), + fetchErr: errors.New("boom"), + wantFetches: 1, + wantErr: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + fetches := 0 + src := newCASTokenSource(tc.initial, func(context.Context) (string, error) { + fetches++ + if tc.fetchErr != nil { + return "", tc.fetchErr + } + return fresh, nil + }) + + got, err := src.Token(context.Background()) + assert.Equal(t, tc.wantFetches, fetches) + if tc.wantErr { + assert.Error(t, err) + return + } + require.NoError(t, err) + + want := tc.wantToken + if want == "" { + want = tc.initial + } + assert.Equal(t, want, got) + + // A second call reuses the token that is now cached. + _, err = src.Token(context.Background()) + require.NoError(t, err) + assert.Equal(t, tc.wantFetches, fetches) + }) + } +} diff --git a/pkg/grpcconn/grpcconn.go b/pkg/grpcconn/grpcconn.go index 43a95d2e8..b5b2dc3bb 100644 --- a/pkg/grpcconn/grpcconn.go +++ b/pkg/grpcconn/grpcconn.go @@ -50,8 +50,13 @@ type newOptionalArg struct { orgName string cliVersion string maxRecvMsgSize int + tokenProvider TokenProvider } +// TokenProvider returns the bearer token for the next request. Use it when the +// token expires quickly and the connection must get a new one per request. +type TokenProvider func(ctx context.Context) (string, error) + type Option func(*newOptionalArg) func WithCAFile(caFilePath string) Option { @@ -87,6 +92,14 @@ func WithCLIVersion(version string) Option { } } +// WithTokenProvider gets the bearer token from the provider on each request. +// It replaces the static token given to New. +func WithTokenProvider(p TokenProvider) Option { + return func(opt *newOptionalArg) { + opt.tokenProvider = p + } +} + // WithMaxRecvMsgSize overrides the client-side cap on incoming unary message // size. Values <= 0 fall back to DefaultMaxRecvMsgSize. func WithMaxRecvMsgSize(size int) Option { @@ -121,8 +134,9 @@ func New(uri, authToken string, opt ...Option) (*grpc.ClientConn, error) { unaryInterceptors := []grpc.UnaryClientInterceptor{} streamInterceptors := []grpc.StreamClientInterceptor{} - if authToken != "" { + if authToken != "" || optionalArgs.tokenProvider != nil { grpcCreds := newTokenAuth(authToken, optionalArgs.insecure, optionalArgs.orgName) + grpcCreds.provider = optionalArgs.tokenProvider opts = append(opts, grpc.WithPerRPCCredentials(grpcCreds)) unaryInterceptors = append(unaryInterceptors, grpc_retry.UnaryClientInterceptor()) } @@ -221,19 +235,29 @@ type tokenAuth struct { token string insecure bool orgName string + // provider, when set, supplies the token for each request instead of token + provider TokenProvider } // Implementation of PerRPCCredentials interface that sends a bearer token in each request. // https://pkg.go.dev/google.golang.org/grpc/credentials#PerRPCCredentials func newTokenAuth(token string, insecure bool, orgName string) *tokenAuth { - return &tokenAuth{token, insecure, orgName} + return &tokenAuth{token: token, insecure: insecure, orgName: orgName} } // Return value is mapped to request headers. -func (t tokenAuth) GetRequestMetadata(_ context.Context, _ ...string) (map[string]string, error) { +func (t tokenAuth) GetRequestMetadata(ctx context.Context, _ ...string) (map[string]string, error) { const OrganizationHeader = "Chainloop-Organization" + token := t.token + if t.provider != nil { + var err error + if token, err = t.provider(ctx); err != nil { + return nil, fmt.Errorf("getting auth token: %w", err) + } + } + return map[string]string{ - "authorization": "Bearer " + t.token, + "authorization": "Bearer " + token, OrganizationHeader: t.orgName, }, nil } diff --git a/pkg/grpcconn/grpcconn_test.go b/pkg/grpcconn/grpcconn_test.go index fdc6d70e0..c59654f53 100644 --- a/pkg/grpcconn/grpcconn_test.go +++ b/pkg/grpcconn/grpcconn_test.go @@ -19,12 +19,17 @@ import ( "context" "crypto/x509" "encoding/base64" + "errors" + "fmt" + "net" "os" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "google.golang.org/grpc" + "google.golang.org/grpc/health" + healthpb "google.golang.org/grpc/health/grpc_health_v1" "google.golang.org/grpc/metadata" ) @@ -258,3 +263,67 @@ func TestBackwardCompatibility_OldClientNewConfig(t *testing.T) { } assert.NoError(t, err, "new client should load file path via detection") } + +func TestGetRequestMetadataWithTokenProvider(t *testing.T) { + testCases := []struct { + name string + provider TokenProvider + wantHeaders map[string]string + wantErr bool + }{ + { + name: "provider token is sent", + provider: func(context.Context) (string, error) { return "fresh", nil }, + wantHeaders: map[string]string{"authorization": "Bearer fresh", "Chainloop-Organization": ""}, + }, + { + name: "provider error fails the request", + provider: func(context.Context) (string, error) { return "", errors.New("boom") }, + wantErr: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + auth := &tokenAuth{provider: tc.provider} + got, err := auth.GetRequestMetadata(context.TODO()) + if tc.wantErr { + assert.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tc.wantHeaders, got) + }) + } +} + +func TestNewWithTokenProviderSendsTokenPerCall(t *testing.T) { + lis, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + var gotAuth []string + srv := grpc.NewServer(grpc.UnaryInterceptor(func(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { + md, _ := metadata.FromIncomingContext(ctx) + gotAuth = append(gotAuth, md.Get("authorization")...) + return handler(ctx, req) + })) + healthpb.RegisterHealthServer(srv, health.NewServer()) + go func() { _ = srv.Serve(lis) }() + t.Cleanup(srv.Stop) + + calls := 0 + conn, err := New(lis.Addr().String(), "", WithInsecure(true), WithTokenProvider(func(context.Context) (string, error) { + calls++ + return fmt.Sprintf("token-%d", calls), nil + })) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + + client := healthpb.NewHealthClient(conn) + for range 2 { + _, err := client.Check(context.Background(), &healthpb.HealthCheckRequest{}) + require.NoError(t, err) + } + + assert.Equal(t, []string{"Bearer token-1", "Bearer token-2"}, gotAuth) +}