Initialize usage reporting via plugin registry after auth server setup (#64265)

* Initialize usage reporting via plugin registry after auth server setup

Usage reporting was previously initialized in enterprise process
constructors (cloud.NewTeleport, pro.NewTeleport), which ran after the
OSS auth server completed initialization. This caused enterprise auth
extensions to miss the UsageReporter service because they were initialized
after the OSS auth server was created but before the usage reporting was set up.

To fix the ordering, add a callback mechanism to plugin.Registry:
- SetUsageReportingInitFunc registers a callback to be called once the
  auth server is ready.
- InitUsageReporting triggers that callback from initAuthService, after
  setLocalAuth makes the server available to callers.

Signed-off-by: Tiago Silva <tiago.silva@goteleport.com>

* update comment

* handle review feedback

* remove dynamic function load

* fix another addition

---------

Signed-off-by: Tiago Silva <tiago.silva@goteleport.com>

lib/auth/auth.go

@@ -1498,6 +1498,10 @@ type Server struct {
 	// EncryptedIO provides encryption for session related data such as
 	// recordings, thumbnails, and metadata.
 	EncryptedIO *recordingencryption.EncryptedIO
+
+	// anonymizationKey is used to anonymize sensitive data in a consistent way for
+	// use in telemetry and logging without exposing the original values.
+	anonymizationKey []byte
 }
 
 // SetSAMLService registers svc as the SAMLService that provides the SAML
@@ -2552,21 +2556,65 @@ func (a *Server) GetClusterID(ctx context.Context) (string, error) {
 	return clusterName.GetClusterID(), nil
 }
 
-// GetAnonymizationKey returns the anonymization key that identifies this client.
+// InitializeAnonymizationKey initializes the HMAC anonymization key for this auth server.
+// The anonymization key is used for hashing any PII data (like usernames, hostnames, etc.)
+// that may be included in events emitted by this auth server.
 // The anonymization key may be any of the following, in order of precedence:
 // - (Teleport Cloud) a key provided by the Teleport Cloud API
 // - a key embedded in the license file
 // - the cluster's UUID
-func (a *Server) GetAnonymizationKey(ctx context.Context) (string, error) {
+// If the anonymization key was already initialized, this function is a no-op.
+func (a *Server) InitializeAnonymizationKey() error {
+	a.lock.Lock()
+	defer a.lock.Unlock()
+	if len(a.anonymizationKey) > 0 {
+		// already initialized, no need to do anything
+		return nil
+	}
+
 	if key := modules.GetModules().Features().CloudAnonymizationKey; len(key) > 0 {
-		return string(key), nil
+		a.anonymizationKey = key
+		return nil
 	}
 
 	if a.license != nil && len(a.license.AnonymizationKey) > 0 {
-		return string(a.license.AnonymizationKey), nil
+		a.anonymizationKey = a.license.AnonymizationKey
+		return nil
 	}
-	id, err := a.GetClusterID(ctx)
-	return id, trace.Wrap(err)
+
+	// load the cluster name from the backend to get the cluster ID for anonymization key generation
+	clusterName, err := a.Services.GetClusterName(a.closeCtx)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+
+	id := clusterName.GetClusterID()
+	if len(id) == 0 {
+		// this should never happen since cluster ID is a required field,
+		// but we check just in case to avoid panics and to provide a more helpful error message.
+		return trace.BadParameter("cluster ID is empty, cannot initialize anonymization key")
+	}
+	a.anonymizationKey = []byte(id)
+	return nil
+}
+
+// GetAnonymizationKey returns the anonymization key for this auth server.
+// Before calling this function, InitializeAnonymizationKey should have been called
+// to ensure the key is properly initialized.
+func (a *Server) GetAnonymizationKey() []byte {
+	a.lock.Lock()
+	defer a.lock.Unlock()
+
+	return a.anonymizationKey
+}
+
+// SetAnonymizationKey sets the anonymization key for this auth server.
+// This is only used by Teleport Cloud where the anonymization key can be
+// refreshed.
+func (a *Server) SetAnonymizationKey(key []byte) {
+	a.lock.Lock()
+	defer a.lock.Unlock()
+	a.anonymizationKey = key
 }
 
 // GetDomainName returns the domain name that identifies this authority server.

lib/auth/auth_test.go

@@ -5009,9 +5009,17 @@ func TestServer_GetAnonymizationKey(t *testing.T) {
 
 			testTLSServer.AuthServer.AuthServer.SetLicense(tt.license)
 
-			got, err := testTLSServer.AuthServer.AuthServer.GetAnonymizationKey(context.Background())
+			err = testTLSServer.AuthServer.AuthServer.InitializeAnonymizationKey()
+			require.NoError(t, err)
+
+			got := testTLSServer.AuthServer.AuthServer.GetAnonymizationKey()
 			tt.errCheck(t, err)
-			require.Equal(t, tt.want, got)
+			require.Equal(t, tt.want, string(got))
+
+			testTLSServer.AuthServer.AuthServer.SetAnonymizationKey([]byte("somethingelse"))
+			got = testTLSServer.AuthServer.AuthServer.GetAnonymizationKey()
+			tt.errCheck(t, err)
+			require.Equal(t, "somethingelse", string(got))
 		})
 	}
 }

lib/auth/discoveryconfig/discoveryconfigv1/service.go

@@ -65,7 +65,7 @@ type ServiceConfig struct {
 	Emitter apievents.Emitter
 
 	// UsageReporter is the reporter for sending usage events.
-	UsageReporter func() usagereporter.UsageReporter
+	UsageReporter usagereporter.UsageReporter
 }
 
 // CheckAndSetDefaults checks the ServiceConfig fields and returns an error if
@@ -105,7 +105,7 @@ type Service struct {
 	backend       services.DiscoveryConfigs
 	clock         clockwork.Clock
 	emitter       apievents.Emitter
-	usageReporter func() usagereporter.UsageReporter
+	usageReporter usagereporter.UsageReporter
 }
 
 // NewService returns a new DiscoveryConfigs gRPC service.
@@ -432,7 +432,7 @@ func (s *Service) emitUsageEvent(dc *discoveryconfig.DiscoveryConfig, action pre
 		creationMethod = CreationMethodIAC
 	}
 
-	s.usageReporter().AnonymizeAndSubmit(&usagereporter.DiscoveryConfigEvent{
+	s.usageReporter.AnonymizeAndSubmit(&usagereporter.DiscoveryConfigEvent{
 		Action:              action,
 		DiscoveryConfigName: dc.GetName(),
 		ResourceTypes:       resourceTypes,

lib/auth/discoveryconfig/discoveryconfigv1/service_test.go

@@ -548,7 +548,7 @@ func initSvc(t *testing.T, clusterName string) (context.Context, localClient, *S
 		Backend:       localResourceService,
 		Authorizer:    authorizer,
 		Emitter:       emitter,
-		UsageReporter: func() usagereporter.UsageReporter { return usagereporter.DiscardUsageReporter{} },
+		UsageReporter: usagereporter.DiscardUsageReporter{},
 	})
 	require.NoError(t, err)
 

lib/auth/grpcserver.go

@@ -156,7 +156,6 @@ import (
 	"github.com/gravitational/teleport/lib/services/local/generic"
 	"github.com/gravitational/teleport/lib/session"
 	"github.com/gravitational/teleport/lib/srv/server/installer"
-	usagereporter "github.com/gravitational/teleport/lib/usagereporter/teleport"
 	"github.com/gravitational/teleport/lib/utils"
 	logutils "github.com/gravitational/teleport/lib/utils/log"
 	"github.com/gravitational/teleport/lib/utils/slices"
@@ -6390,12 +6389,10 @@ func NewGRPCServer(cfg GRPCServerConfig) (*GRPCServer, error) {
 	integrationv1pb.RegisterAWSRolesAnywhereServiceServer(server, integrationAWSRolesAnywhereServiceServer)
 
 	userTask, err := usertasksv1.NewService(usertasksv1.ServiceConfig{
-		Authorizer: cfg.Authorizer,
-		Backend:    cfg.AuthServer.Services,
-		Cache:      cfg.AuthServer.Cache,
-		// This must be a function because cfg.AuthServer.UsageReporter is changed after `NewGRPCServer` is called.
-		// It starts as a DiscardUsageReporter, but when running in Cloud, gets replaced by a real reporter.
-		UsageReporter: func() usagereporter.UsageReporter { return cfg.AuthServer.UsageReporter },
+		Authorizer:    cfg.Authorizer,
+		Backend:       cfg.AuthServer.Services,
+		Cache:         cfg.AuthServer.Cache,
+		UsageReporter: cfg.AuthServer.UsageReporter,
 		Emitter:       cfg.Emitter,
 	})
 	if err != nil {
@@ -6404,13 +6401,11 @@ func NewGRPCServer(cfg GRPCServerConfig) (*GRPCServer, error) {
 	usertaskv1pb.RegisterUserTaskServiceServer(server, userTask)
 
 	discoveryConfig, err := discoveryconfigv1.NewService(discoveryconfigv1.ServiceConfig{
-		Authorizer: cfg.Authorizer,
-		Backend:    cfg.AuthServer.Services,
-		Clock:      cfg.AuthServer.clock,
-		Emitter:    cfg.Emitter,
-		// This must be a function because cfg.AuthServer.UsageReporter is changed after `NewGRPCServer` is called.
-		// It starts as a DiscardUsageReporter, but when running in Cloud, gets replaced by a real reporter.
-		UsageReporter: func() usagereporter.UsageReporter { return cfg.AuthServer.UsageReporter },
+		Authorizer:    cfg.Authorizer,
+		Backend:       cfg.AuthServer.Services,
+		Clock:         cfg.AuthServer.clock,
+		Emitter:       cfg.Emitter,
+		UsageReporter: cfg.AuthServer.UsageReporter,
 		Logger:        cfg.AuthServer.logger.With(teleport.ComponentKey, "discoveryconfig_crud_service"),
 	})
 	if err != nil {

lib/auth/usertasks/usertasksv1/service.go

@@ -54,7 +54,7 @@ type ServiceConfig struct {
 	Clock clockwork.Clock
 
 	// UsageReporter is the reporter for sending usage without it be related to an API call.
-	UsageReporter func() usagereporter.UsageReporter
+	UsageReporter usagereporter.UsageReporter
 
 	// Emitter is the event emitter.
 	Emitter apievents.Emitter
@@ -100,7 +100,7 @@ type Service struct {
 	backend       services.UserTasks
 	cache         Reader
 	clock         clockwork.Clock
-	usageReporter func() usagereporter.UsageReporter
+	usageReporter usagereporter.UsageReporter
 	emitter       apievents.Emitter
 }
 
@@ -139,7 +139,7 @@ func (s *Service) CreateUserTask(ctx context.Context, req *usertasksv1.CreateUse
 		return nil, trace.Wrap(err)
 	}
 
-	s.usageReporter().AnonymizeAndSubmit(userTaskToUserTaskStateEvent(req.GetUserTask()))
+	s.usageReporter.AnonymizeAndSubmit(userTaskToUserTaskStateEvent(req.GetUserTask()))
 
 	return rsp, nil
 }
@@ -279,7 +279,7 @@ func (s *Service) UpdateUserTask(ctx context.Context, req *usertasksv1.UpdateUse
 	}
 
 	if stateChanged {
-		s.usageReporter().AnonymizeAndSubmit(userTaskToUserTaskStateEvent(req.GetUserTask()))
+		s.usageReporter.AnonymizeAndSubmit(userTaskToUserTaskStateEvent(req.GetUserTask()))
 	}
 
 	return rsp, nil
@@ -345,7 +345,7 @@ func (s *Service) UpsertUserTask(ctx context.Context, req *usertasksv1.UpsertUse
 	}
 
 	if stateChanged {
-		s.usageReporter().AnonymizeAndSubmit(userTaskToUserTaskStateEvent(req.GetUserTask()))
+		s.usageReporter.AnonymizeAndSubmit(userTaskToUserTaskStateEvent(req.GetUserTask()))
 	}
 
 	return rsp, nil

lib/auth/usertasks/usertasksv1/service_test.go

@@ -349,7 +349,7 @@ func newService(t *testing.T, checker services.AccessChecker, usageReporter usag
 		Authorizer:    authorizer,
 		Backend:       backendService,
 		Cache:         backendService,
-		UsageReporter: func() usagereporter.UsageReporter { return usageReporter },
+		UsageReporter: usageReporter,
 		Emitter:       emitter,
 		Clock:         clock,
 	})

lib/plugin/registry.go

@@ -51,6 +51,16 @@ type Registry interface {
 	RegisterAuthWebHandlers(handler any) error
 	// RegisterAuthServices registers Teleport AuthServer services
 	RegisterAuthServices(ctx context.Context, server any, getClientCert getCertFunc) error
+	// SetUsageReportingInitFunc stores a callback that will be called once the OSS auth server is
+	// fully initialized. The callback receives a *service.TeleportProcess and is responsible for
+	// starting all usage-reporting pipelines. It must be set before service.NewTeleport is called.
+	SetUsageReportingInitFunc(func(process any) error)
+	// InitUsageReporting invokes the callback registered with SetUsageReportingInitFunc. It is
+	// called by initAuthService after setLocalAuth, ensuring the UsageReporter is available to
+	// auth-server enterprise extensions. The function process parameter is a *service.TeleportProcess
+	// instance.
+	// A no-op if no callback was registered.
+	InitUsageReporting(process any) error
 }
 
 // NewRegistry creates an instance of the Registry
@@ -61,7 +71,19 @@ func NewRegistry() Registry {
 }
 
 type registry struct {
-	plugins map[string]Plugin
+	plugins              map[string]Plugin
+	usageReportingInitFn func(process any) error
+}
+
+func (r *registry) SetUsageReportingInitFunc(f func(process any) error) {
+	r.usageReportingInitFn = f
+}
+
+func (r *registry) InitUsageReporting(process any) error {
+	if r.usageReportingInitFn != nil {
+		return r.usageReportingInitFn(process)
+	}
+	return nil
 }
 
 // IsRegistered returns whether a plugin with the give name exists.

lib/service/service.go

@@ -1136,7 +1136,6 @@ func NewTeleport(cfg *servicecfg.Config) (_ *TeleportProcess, err error) {
 		return nil, trace.BadParameter("binary not compiled against BoringCrypto, check " +
 			"that Enterprise FIPS release was downloaded from " +
 			"a Teleport account https://teleport.sh")
-
 	}
 
 	if cfg.Auth.Preference.GetPrivateKeyPolicy().IsHardwareKeyPolicy() && cfg.Modules.BuildType() != modules.BuildEnterprise {
@@ -2514,6 +2513,16 @@ func (process *TeleportProcess) initAuthService() error {
 	if err != nil {
 		return trace.Wrap(err)
 	}
+
+	// setLocalAuth must be called before InitUsageReporting so that enterprise
+	// auth extensions can access the UsageReporter via GetAuthServer.
+	process.setLocalAuth(authServer)
+	if process.Config.PluginRegistry != nil {
+		if err := process.Config.PluginRegistry.InitUsageReporting(process); err != nil {
+			return trace.Wrap(err, "initializing usage reporting")
+		}
+	}
+
 	authServer.EncryptedIO = encryptedIO
 
 	lockWatcher, err := services.NewLockWatcher(process.ExitContext(), services.LockWatcherConfig{
@@ -2592,8 +2601,6 @@ func (process *TeleportProcess) initAuthService() error {
 	}
 	recordingMetadataProvider.SetService(recordingMetadataService)
 
-	process.setLocalAuth(authServer)
-
 	// The auth server runs its own upload completer, which is necessary in sync recording modes where
 	// a node can abandon an upload before it is competed.
 	// (In async recording modes, auth only ever sees completed uploads, as the node's upload completer

lib/usagereporter/daemon/usagereporter.go

@@ -181,5 +181,5 @@ func newClusterAnonymizer(authClusterID string) (utils.Anonymizer, error) {
 	if err != nil {
 		return nil, trace.BadParameter("Invalid auth cluster ID %s", authClusterID)
 	}
-	return utils.NewHMACAnonymizer(authClusterID)
+	return utils.NewHMACAnonymizer(utils.AnonymizationKeyString(authClusterID))
 }

lib/usagereporter/teleport/aggregating/reporter_test.go

@@ -72,7 +72,7 @@ func TestReporter(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	anonymizer, err := utils.NewHMACAnonymizer("0123456789abcdef")
+	anonymizer, err := utils.NewHMACAnonymizer(utils.AnonymizationKeyString("0123456789abcdef"))
 	require.NoError(t, err)
 
 	r, err := NewReporter(ctx, ReporterConfig{
@@ -249,7 +249,7 @@ func TestReporterMachineWorkloadIdentityActivity(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	anonymizer, err := utils.NewHMACAnonymizer("0123456789abcdef")
+	anonymizer, err := utils.NewHMACAnonymizer(utils.AnonymizationKeyString("0123456789abcdef"))
 	require.NoError(t, err)
 
 	r, err := NewReporter(ctx, ReporterConfig{
@@ -416,7 +416,7 @@ func TestReporterSessionSummariesAccessed(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		anonymizer, err := utils.NewHMACAnonymizer("0123456789abcdef")
+		anonymizer, err := utils.NewHMACAnonymizer(utils.AnonymizationKeyString("0123456789abcdef"))
 		require.NoError(t, err)
 
 		r, err := NewReporter(ctx, ReporterConfig{
@@ -523,7 +523,7 @@ func TestReporterIdentitySecuritySummariesGenerated(t *testing.T) {
 		})
 		require.NoError(t, err)
 
-		anonymizer, err := utils.NewHMACAnonymizer("0123456789abcdef")
+		anonymizer, err := utils.NewHMACAnonymizer(utils.AnonymizationKeyString("0123456789abcdef"))
 		require.NoError(t, err)
 
 		r, err := NewReporter(ctx, ReporterConfig{

lib/usagereporter/teleport/audit_test.go

@@ -30,7 +30,7 @@ import (
 )
 
 func TestConvertAuditEvent(t *testing.T) {
-	anonymizer, err := utils.NewHMACAnonymizer("anon-key-or-cluster-id")
+	anonymizer, err := utils.NewHMACAnonymizer(utils.AnonymizationKeyString("anon-key-or-cluster-id"))
 	require.NoError(t, err)
 
 	cases := []struct {

lib/usagereporter/teleport/usagereporter_test.go

@@ -33,7 +33,7 @@ import (
 )
 
 func TestConvertUsageEvent(t *testing.T) {
-	anonymizer, err := utils.NewHMACAnonymizer("anon-key-or-cluster-id")
+	anonymizer, err := utils.NewHMACAnonymizer(utils.AnonymizationKeyString("anon-key-or-cluster-id"))
 	require.NoError(t, err)
 
 	expectedAnonymizedUserString := anonymizer.AnonymizeString("myuser")
@@ -726,15 +726,16 @@ func TestConvertUsageEvent(t *testing.T) {
 			}},
 			identityUsername: "myuser",
 			errCheck:         require.NoError,
-			expected: &prehogv1a.SubmitEventRequest{Event: &prehogv1a.SubmitEventRequest_AccessListReviewDelete{
-				AccessListReviewDelete: &prehogv1a.AccessListReviewDeleteEvent{
-					UserName: expectedAnonymizedUserString,
-					Metadata: &prehogv1a.AccessListMetadata{
-						Id: expectedAnonymizedAccessListIDString,
+			expected: &prehogv1a.SubmitEventRequest{
+				Event: &prehogv1a.SubmitEventRequest_AccessListReviewDelete{
+					AccessListReviewDelete: &prehogv1a.AccessListReviewDeleteEvent{
+						UserName: expectedAnonymizedUserString,
+						Metadata: &prehogv1a.AccessListMetadata{
+							Id: expectedAnonymizedAccessListIDString,
+						},
 					},
 				},
 			},
-			},
 		},
 		{
 			name: "discovery fetch event",

lib/utils/anonymizer.go

@@ -23,7 +23,6 @@ import (
 	"crypto/sha256"
 	"encoding/base64"
 	"strings"
-	"sync"
 
 	"github.com/gravitational/trace"
 )
@@ -39,39 +38,51 @@ type Anonymizer interface {
 	// AnonymizeNonEmpty anonymizes the given string into bytes if the string is
 	// nonempty, otherwise returns an empty slice.
 	AnonymizeNonEmpty(s string) []byte
+}
 
-	// SetAnonymizationKey updates the underlying anonymization key.
-	SetAnonymizationKey(k []byte)
+var _ AnonymizationKeyProvider = (AnonymizationKeyString)("")
+
+// AnonymizationKeyString is a simple implementation of AnonymizationKeyProvider that uses a string as the key.
+type AnonymizationKeyString string
+
+func (h AnonymizationKeyString) GetAnonymizationKey() []byte {
+	return []byte(h)
+}
+
+func (h AnonymizationKeyString) InitializeAnonymizationKey() error {
+	return nil
 }
 
 // HMACAnonymizer implements anonymization using HMAC
 type HMACAnonymizer struct {
 	// key is the HMAC key
-	key []byte
-	mu  sync.RWMutex
+	keyProvider AnonymizationKeyProvider
 }
 
 var _ Anonymizer = (*HMACAnonymizer)(nil)
 
-func (a *HMACAnonymizer) SetAnonymizationKey(k []byte) {
-	a.mu.Lock()
-	a.key = k
-	a.mu.Unlock()
+type AnonymizationKeyProvider interface {
+	// InitializeAnonymizationKey initializes the anonymization key if needed.
+	InitializeAnonymizationKey() error
+	// GetHMACAnonymizerKey returns the HMAC anonymizer key.
+	GetAnonymizationKey() []byte
 }
 
 // NewHMACAnonymizer returns a new HMAC-based anonymizer
-func NewHMACAnonymizer(key string) (*HMACAnonymizer, error) {
-	if strings.TrimSpace(key) == "" {
+func NewHMACAnonymizer(keyProvider AnonymizationKeyProvider) (*HMACAnonymizer, error) {
+	if err := keyProvider.InitializeAnonymizationKey(); err != nil {
+		return nil, trace.Wrap(err, "failed to initialize anonymization key")
+	}
+	key := keyProvider.GetAnonymizationKey()
+	if strings.TrimSpace(string(key)) == "" {
 		return nil, trace.BadParameter("HMAC key must not be empty")
 	}
-	return &HMACAnonymizer{key: []byte(key)}, nil
+	return &HMACAnonymizer{keyProvider: keyProvider}, nil
 }
 
 // Anonymize anonymizes the provided data using HMAC
 func (a *HMACAnonymizer) Anonymize(data []byte) string {
-	a.mu.RLock()
-	k := a.key
-	a.mu.RUnlock()
+	k := a.keyProvider.GetAnonymizationKey()
 
 	h := hmac.New(sha256.New, k)
 	h.Write(data)
@@ -89,9 +100,7 @@ func (a *HMACAnonymizer) AnonymizeNonEmpty(s string) []byte {
 		return nil
 	}
 
-	a.mu.RLock()
-	k := a.key
-	a.mu.RUnlock()
+	k := a.keyProvider.GetAnonymizationKey()
 
 	h := hmac.New(sha256.New, k)
 	h.Write([]byte(s))

lib/utils/anonymizer_test.go

@@ -28,11 +28,11 @@ import (
 func TestHMACAnonymizer(t *testing.T) {
 	t.Parallel()
 
-	a, err := NewHMACAnonymizer(" ")
+	a, err := NewHMACAnonymizer(AnonymizationKeyString(" "))
 	require.ErrorAs(t, err, new(*trace.BadParameterError))
 	require.Nil(t, a)
 
-	a, err = NewHMACAnonymizer("key")
+	a, err = NewHMACAnonymizer(AnonymizationKeyString("key"))
 	require.NoError(t, err)
 	require.NotNil(t, a)