From 6d77c0e187eafd6c37bd00684a5b5ca88fad9217 Mon Sep 17 00:00:00 2001
From: Gabriel MABILLE <gamab@users.noreply.github.com>
Date: Mon, 25 Nov 2024 14:17:52 +0100
Subject: [PATCH] AuthZ client: Add tracing (#96983)

* AuthZ client: Add tracing

* InProc as well
---
 pkg/services/authz/client.go | 27 +++++++++++++++++----------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/pkg/services/authz/client.go b/pkg/services/authz/client.go
index 4727f9d56d4..5095706ce8f 100644
--- a/pkg/services/authz/client.go
+++ b/pkg/services/authz/client.go
@@ -53,17 +53,17 @@ func ProvideAuthZClient(
 
 	switch authCfg.mode {
 	case ModeInProc:
-		client, err = newInProcLegacyClient(server)
+		client, err = newInProcLegacyClient(server, tracer)
 		if err != nil {
 			return nil, err
 		}
 	case ModeGRPC:
-		client, err = newGrpcLegacyClient(authCfg)
+		client, err = newGrpcLegacyClient(authCfg, tracer)
 		if err != nil {
 			return nil, err
 		}
 	case ModeCloud:
-		client, err = newCloudLegacyClient(authCfg)
+		client, err = newCloudLegacyClient(authCfg, tracer)
 		if err != nil {
 			return nil, err
 		}
@@ -87,12 +87,12 @@ func ProvideStandaloneAuthZClient(
 	}
 
 	if authCfg.mode == ModeGRPC {
-		return newGrpcLegacyClient(authCfg)
+		return newGrpcLegacyClient(authCfg, tracer)
 	}
-	return newCloudLegacyClient(authCfg)
+	return newCloudLegacyClient(authCfg, tracer)
 }
 
-func newInProcLegacyClient(server *legacyServer) (authzlib.AccessChecker, error) {
+func newInProcLegacyClient(server *legacyServer, tracer tracing.Tracer) (authzlib.AccessChecker, error) {
 	noAuth := func(ctx context.Context) (context.Context, error) {
 		return ctx, nil
 	}
@@ -111,13 +111,18 @@ func newInProcLegacyClient(server *legacyServer) (authzlib.AccessChecker, error)
 		&authzlib.ClientConfig{},
 		authzlib.WithGrpcConnectionClientOption(channel),
 		authzlib.WithDisableAccessTokenClientOption(),
+		authzlib.WithTracerClientOption(tracer),
 	)
 }
 
-func newGrpcLegacyClient(authCfg *Cfg) (authzlib.AccessChecker, error) {
+func newGrpcLegacyClient(authCfg *Cfg, tracer tracing.Tracer) (authzlib.AccessChecker, error) {
 	// This client interceptor is a noop, as we don't send an access token
 	clientConfig := authnlib.GrpcClientConfig{}
-	clientInterceptor, err := authnlib.NewGrpcClientInterceptor(&clientConfig, authnlib.WithDisableAccessTokenOption())
+	clientInterceptor, err := authnlib.NewGrpcClientInterceptor(
+		&clientConfig,
+		authnlib.WithDisableAccessTokenOption(),
+		authnlib.WithTracerOption(tracer),
+	)
 	if err != nil {
 		return nil, err
 	}
@@ -127,6 +132,7 @@ func newGrpcLegacyClient(authCfg *Cfg) (authzlib.AccessChecker, error) {
 		authzlib.WithGrpcDialOptionsClientOption(
 			getDialOpts(clientInterceptor, authCfg.allowInsecure)...,
 		),
+		authzlib.WithTracerClientOption(tracer),
 		// TODO: remove this once access tokens are supported on-prem
 		authzlib.WithDisableAccessTokenClientOption(),
 	)
@@ -137,7 +143,7 @@ func newGrpcLegacyClient(authCfg *Cfg) (authzlib.AccessChecker, error) {
 	return client, nil
 }
 
-func newCloudLegacyClient(authCfg *Cfg) (authzlib.AccessChecker, error) {
+func newCloudLegacyClient(authCfg *Cfg, tracer tracing.Tracer) (authzlib.AccessChecker, error) {
 	grpcClientConfig := authnlib.GrpcClientConfig{
 		TokenClientConfig: &authnlib.TokenExchangeConfig{
 			Token:            authCfg.token,
@@ -149,7 +155,7 @@ func newCloudLegacyClient(authCfg *Cfg) (authzlib.AccessChecker, error) {
 		},
 	}
 
-	clientInterceptor, err := authnlib.NewGrpcClientInterceptor(&grpcClientConfig)
+	clientInterceptor, err := authnlib.NewGrpcClientInterceptor(&grpcClientConfig, authnlib.WithTracerOption(tracer))
 	if err != nil {
 		return nil, err
 	}
@@ -159,6 +165,7 @@ func newCloudLegacyClient(authCfg *Cfg) (authzlib.AccessChecker, error) {
 		authzlib.WithGrpcDialOptionsClientOption(
 			getDialOpts(clientInterceptor, authCfg.allowInsecure)...,
 		),
+		authzlib.WithTracerClientOption(tracer),
 	)
 	if err != nil {
 		return nil, err