grpc: refresh trusted headers from continuous auth responses

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * added `x-substreams-plan-tier` trusted header and `SubstreamsPlanTier()` to framework, documentation and null/secret implementations
 * add 'interval' param to grpc continuous auth, default changed from "10s" to "60s"
 * move from deprecated `github.com/bufbuild/connect-go` to `connectrpc.com/connect`
+* grpc continuous auth now merges each response's headers into the stream's trusted headers (added `ReplaceTrustedHeaders`), so a stream whose first authentication returned no headers picks them up on the next continuous call; `FromContext` returns a snapshot, read it again to see refreshed values
+* bump minimum Go version to 1.25
 
 ### Added
 

go.mod

@@ -1,8 +1,6 @@
 module github.com/streamingfast/dauth
 
-go 1.24.0
-
-toolchain go1.24.11
+go 1.25.0
 
 require (
 	connectrpc.com/connect v1.14.0

grpc/authenticator.go

@@ -3,6 +3,7 @@ package grpc
 import (
 	"context"
 	"fmt"
+	"maps"
 	"strings"
 	"time"
 	"unicode/utf8"
@@ -92,16 +93,20 @@ func (a *authenticatorPlugin) Authenticate(ctx context.Context, path string, hea
 		return nil, err
 	}
 
-	out := make(dauth.TrustedHeaders)
-	for _, authenticatedHeader := range resp.AuthenticatedHeaders {
-		out[strings.ToLower(authenticatedHeader.Key)] = authenticatedHeader.Value
-	}
-
 	ctx, cancel := context.WithCancelCause(ctx)
+	ctx = dauth.WithTrustedHeaders(ctx, toTrustedHeaders(resp))
 
 	go a.continuousAuth(ctx, req, cancel)
 
-	return dauth.WithTrustedHeaders(ctx, out), nil
+	return ctx, nil
+}
+
+func toTrustedHeaders(resp *pbauth.AuthResponse) dauth.TrustedHeaders {
+	out := make(dauth.TrustedHeaders)
+	for _, authenticatedHeader := range resp.AuthenticatedHeaders {
+		out[strings.ToLower(authenticatedHeader.Key)] = authenticatedHeader.Value
+	}
+	return out
 }
 
 func (a *authenticatorPlugin) continuousAuth(ctx context.Context, req *pbauth.AuthRequest, cancel context.CancelCauseFunc) {
@@ -118,9 +123,38 @@ func (a *authenticatorPlugin) continuousAuth(ctx context.Context, req *pbauth.Au
 
 		req.AuthCount++
 
-		if _, err := a.client.Authenticate(context.Background(), req); err != nil {
+		resp, err := a.client.Authenticate(context.Background(), req)
+		if err != nil {
 			cancel(err)
 			return
 		}
+
+		a.refreshTrustedHeaders(ctx, resp)
+	}
+}
+
+// refreshTrustedHeaders adds or updates headers from a continuous auth response, never removes them.
+// An empty response (e.g. a fail-open authenticator) is a no-op, it cannot be told apart from an
+// authoritative empty set, so a header the server stops sending stays for the stream lifetime.
+// Only continuousAuth writes the headers of its context, so the read-modify-write is not racy.
+func (a *authenticatorPlugin) refreshTrustedHeaders(ctx context.Context, resp *pbauth.AuthResponse) {
+	refreshed := toTrustedHeaders(resp)
+	if len(refreshed) == 0 {
+		return
+	}
+
+	current := dauth.FromContext(ctx)
+	merged := maps.Clone(current)
+	if merged == nil {
+		merged = make(dauth.TrustedHeaders, len(refreshed))
+	}
+	maps.Copy(merged, refreshed)
+	if maps.Equal(merged, current) {
+		return
+	}
+
+	if a.logger != nil {
+		a.logger.Debug("continuous authentication refreshed trusted headers", zap.Strings("headers", refreshed.Names()))
 	}
+	dauth.ReplaceTrustedHeaders(ctx, merged)
 }

grpc/authenticator_test.go

@@ -3,12 +3,15 @@ package grpc
 import (
 	"context"
 	"fmt"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"github.com/streamingfast/dauth"
 	pbauth "github.com/streamingfast/dauth/pb/sf/authentication/v1"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
-	"testing"
-	"time"
 )
 
 type mockClient struct {
@@ -28,36 +31,75 @@ func (m *mockClient) Authenticate(ctx context.Context, in *pbauth.AuthRequest, o
 }
 
 func TestAuthenticatorPlugin_ContinuousAuthenticate(t *testing.T) {
-	parentCtx := context.Background()
-	header := map[string][]string{
-		"x-user-id":   []string{"userid"},
-		"x-apikey-id": []string{"apiKey"},
-	}
-	ipAddress := "192.168.1.1"
-	continuousInternal := 10 * time.Millisecond
-	failOnCount := uint64(3)
-	// should be greater than (failOnCount * continuousInternal) + small bugger
-	testDuration := 35 * time.Millisecond
-
-	mockClient := &mockClient{failOnCount: failOnCount}
-	a := &authenticatorPlugin{
-		client:                mockClient,
-		continuousInterval:    continuousInternal,
-		enabledContinuousAuth: true,
-	}
+	synctest.Test(t, func(t *testing.T) {
+		header := map[string][]string{
+			"x-user-id":   {"userid"},
+			"x-apikey-id": {"apiKey"},
+		}
 
-	authenticatedCtx, err := a.Authenticate(parentCtx, "sf.firehose.v1/Blocks", header, ipAddress)
-	require.NoError(t, err)
-
-	for {
-		select {
-		case <-authenticatedCtx.Done():
-			require.Equal(t, context.Canceled, authenticatedCtx.Err())
-			require.Equal(t, "authentication failure", context.Cause(authenticatedCtx).Error())
-			return
-		case <-time.After(testDuration):
-			assert.Fail(t, "the context should have been canceled by now")
+		a := &authenticatorPlugin{
+			client:                &mockClient{failOnCount: 3},
+			continuousInterval:    time.Second,
+			enabledContinuousAuth: true,
 		}
+
+		authenticatedCtx, err := a.Authenticate(context.Background(), "sf.firehose.v1/Blocks", header, "192.168.1.1")
+		require.NoError(t, err)
+
+		// AuthCount 2 succeeds, AuthCount 3 fails
+		time.Sleep(time.Second)
+		synctest.Wait()
+		require.NoError(t, authenticatedCtx.Err())
+
+		time.Sleep(time.Second)
+		synctest.Wait()
+		require.Equal(t, context.Canceled, authenticatedCtx.Err())
+		require.Equal(t, "authentication failure", context.Cause(authenticatedCtx).Error())
+	})
+}
+
+type scriptedClient struct {
+	responses map[uint64]map[string]string
+}
+
+func (m *scriptedClient) Authenticate(ctx context.Context, in *pbauth.AuthRequest, opts ...grpc.CallOption) (*pbauth.AuthResponse, error) {
+	out := &pbauth.AuthResponse{}
+	for k, v := range m.responses[in.AuthCount] {
+		out.AuthenticatedHeaders = append(out.AuthenticatedHeaders, &pbauth.Header{Key: k, Value: v})
 	}
+	return out, nil
+}
+
+func TestAuthenticatorPlugin_ContinuousAuthenticate_RefreshesHeaders(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		client := &scriptedClient{responses: map[uint64]map[string]string{
+			// other counts get an empty response, like a fail-open authenticator
+			2: {"x-user-id": "userid", "x-meta": "meta"},
+		}}
+		a := &authenticatorPlugin{
+			client:                client,
+			continuousInterval:    time.Second,
+			enabledContinuousAuth: true,
+		}
+
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		authenticatedCtx, err := a.Authenticate(ctx, "sf.firehose.v2.Stream/Blocks", nil, "192.168.1.1")
+		require.NoError(t, err)
+		assert.Equal(t, "", dauth.FromContext(authenticatedCtx).Meta())
+
+		// AuthCount 2 returns the headers
+		time.Sleep(time.Second)
+		synctest.Wait()
+		assert.Equal(t, "meta", dauth.FromContext(authenticatedCtx).Meta())
 
+		// later empty responses must not remove them
+		time.Sleep(5 * time.Second)
+		synctest.Wait()
+		headers := dauth.FromContext(authenticatedCtx)
+		assert.Equal(t, "meta", headers.Meta())
+		assert.Equal(t, "userid", headers.UserID())
+		require.NoError(t, authenticatedCtx.Err())
+	})
 }

headers.go

@@ -5,6 +5,7 @@ import (
 	"maps"
 	"slices"
 	"strings"
+	"sync/atomic"
 
 	"google.golang.org/grpc/metadata"
 )
@@ -42,21 +43,50 @@ type trustedHeadersKeyType int
 
 const trustedHeadersKey trustedHeadersKeyType = iota
 
+type trustedHeadersHolder struct {
+	headers atomic.Pointer[TrustedHeaders]
+}
+
+// WithTrustedHeaders attaches h to a new context through a fresh holder. Calling it again on a
+// context that already carries trusted headers shadows them: the returned context and its
+// children no longer see a [ReplaceTrustedHeaders] made on the parent context.
 func WithTrustedHeaders(ctx context.Context, h TrustedHeaders) context.Context {
-	lowercased := make(TrustedHeaders)
-	for k, v := range h {
-		lowercased[strings.ToLower(k)] = v
-	}
+	holder := &trustedHeadersHolder{}
+	holder.headers.Store(lowercase(h))
 
-	return context.WithValue(ctx, trustedHeadersKey, lowercased)
+	return context.WithValue(ctx, trustedHeadersKey, holder)
 }
 
+// FromContext returns a snapshot of the trusted headers, call it again to see a later [ReplaceTrustedHeaders].
+//
+// The returned map is shared with every other reader of ctx and must not be modified, a
+// concurrent [ReplaceTrustedHeaders] (e.g. from continuous authentication) reads it.
 func FromContext(ctx context.Context) TrustedHeaders {
-	val := ctx.Value(trustedHeadersKey)
-	if val == nil {
+	holder, ok := ctx.Value(trustedHeadersKey).(*trustedHeadersHolder)
+	if !ok {
 		return nil
 	}
-	return val.(TrustedHeaders)
+	return *holder.headers.Load()
+}
+
+// ReplaceTrustedHeaders swaps the trusted headers of ctx and every context derived from it,
+// returning false if ctx carries no trusted headers. The swap is atomic but a read-modify-write
+// built on [FromContext] is not, callers must ensure a single writer per context.
+func ReplaceTrustedHeaders(ctx context.Context, h TrustedHeaders) bool {
+	holder, ok := ctx.Value(trustedHeadersKey).(*trustedHeadersHolder)
+	if !ok {
+		return false
+	}
+	holder.headers.Store(lowercase(h))
+	return true
+}
+
+func lowercase(h TrustedHeaders) *TrustedHeaders {
+	lowercased := make(TrustedHeaders, len(h))
+	for k, v := range h {
+		lowercased[strings.ToLower(k)] = v
+	}
+	return &lowercased
 }
 
 // Deprecated: use [OrganizationID] instead, the [HeaderUserID] now carries the organization id.

headers_test.go

@@ -49,3 +49,21 @@ func TestNilDoesntPanic(t *testing.T) {
 	assert.Equal(t, "", h.SubstreamsPlanTier())
 	assert.Equal(t, "", h.Get("something"))
 }
+
+func TestReplaceTrustedHeaders(t *testing.T) {
+	ctx := WithTrustedHeaders(context.Background(), TrustedHeaders{HeaderUserID: "my-user-id"})
+	derived, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	before := FromContext(derived)
+
+	assert.True(t, ReplaceTrustedHeaders(ctx, TrustedHeaders{HeaderUserID: "my-user-id", "X-Meta": "some-meta"}))
+
+	assert.Equal(t, "some-meta", FromContext(derived).Meta(), "derived contexts see the replaced headers")
+	assert.Equal(t, "", before.Meta(), "a snapshot taken before the replace is unchanged")
+}
+
+func TestReplaceTrustedHeaders_NoHeaders(t *testing.T) {
+	assert.False(t, ReplaceTrustedHeaders(context.Background(), TrustedHeaders{HeaderUserID: "my-user-id"}))
+	assert.Nil(t, FromContext(context.Background()))
+}