From 639ab6cca7b23be083f0e2476052baf370e82297 Mon Sep 17 00:00:00 2001
From: "lean.dev" <34773040+leandro-deveikis@users.noreply.github.com>
Date: Wed, 15 May 2024 09:39:02 -0300
Subject: [PATCH] Tempo: Implement PDC for gRPC client (#85867)

---
 pkg/tsdb/tempo/grpc.go  | 44 +++++++++++++++++++++++++++++++++++++----
 pkg/tsdb/tempo/tempo.go |  2 +-
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/pkg/tsdb/tempo/grpc.go b/pkg/tsdb/tempo/grpc.go
index 12ff0155471..3cac7850382 100644
--- a/pkg/tsdb/tempo/grpc.go
+++ b/pkg/tsdb/tempo/grpc.go
@@ -5,6 +5,7 @@ import (
 	"crypto/tls"
 	"encoding/base64"
 	"fmt"
+	"net"
 	"net/url"
 	"strings"
 
@@ -26,7 +27,7 @@ var logger = backend.NewLoggerWith("logger", "tsdb.tempo")
 // standard HTTP requests.
 // Using other library like connect-go isn't possible right now because Tempo uses non-standard proto compiler which
 // makes generating different client difficult. See https://github.com/grafana/grafana/pull/81683
-func newGrpcClient(settings backend.DataSourceInstanceSettings, opts httpclient.Options) (tempopb.StreamingQuerierClient, error) {
+func newGrpcClient(ctx context.Context, settings backend.DataSourceInstanceSettings, opts httpclient.Options) (tempopb.StreamingQuerierClient, error) {
 	parsedUrl, err := url.Parse(settings.URL)
 	if err != nil {
 		logger.Error("Error parsing URL for gRPC client", "error", err, "URL", settings.URL, "function", logEntrypoint())
@@ -43,18 +44,23 @@ func newGrpcClient(settings backend.DataSourceInstanceSettings, opts httpclient.
 		}
 	}
 
-	clientConn, err := grpc.Dial(onlyHost, getDialOpts(settings, opts)...)
+	dialOpts, err := getDialOpts(ctx, settings, opts)
+	if err != nil {
+		return nil, fmt.Errorf("error getting dial options: %w", err)
+	}
+	clientConn, err := grpc.Dial(onlyHost, dialOpts...)
 	if err != nil {
 		logger.Error("Error dialing gRPC client", "error", err, "URL", settings.URL, "function", logEntrypoint())
 		return nil, err
 	}
 
+	logger.Debug("Instantiating new gRPC client")
 	return tempopb.NewStreamingQuerierClient(clientConn), nil
 }
 
 // getDialOpts creates options and interceptors (middleware) this should roughly match what we do in
 // http_client_provider.go for standard http requests.
-func getDialOpts(settings backend.DataSourceInstanceSettings, opts httpclient.Options) []grpc.DialOption {
+func getDialOpts(ctx context.Context, settings backend.DataSourceInstanceSettings, opts httpclient.Options) ([]grpc.DialOption, error) {
 	// TODO: Missing middleware TracingMiddleware, DataSourceMetricsMiddleware, ContextualMiddleware,
 	//  ResponseLimitMiddleware RedirectLimitMiddleware.
 	// Also User agent but that is set before each rpc call as for decoupled DS we have to get it from request context
@@ -74,7 +80,37 @@ func getDialOpts(settings backend.DataSourceInstanceSettings, opts httpclient.Op
 		dialOps = append(dialOps, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	}
 
-	return dialOps
+	// The following code is required to make gRPC work with Grafana Cloud PDC
+	// (https://grafana.com/docs/grafana-cloud/connect-externally-hosted/private-data-source-connect/)
+	proxyClient, err := settings.ProxyClient(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("proxy client cannot be retrieved, it is not possible to check if secure socks proxy is enabled: %w", err)
+	}
+	if proxyClient.SecureSocksProxyEnabled() { // secure socks proxy is behind a feature flag
+		dialer, err := proxyClient.NewSecureSocksProxyContextDialer()
+		if err != nil {
+			return nil, fmt.Errorf("failure in creating dialer: %w", err)
+		}
+		logger.Debug("gRPC dialer instantiated. Appending gRPC dialer to dial options")
+		dialOps = append(dialOps, grpc.WithContextDialer(func(ctx context.Context, host string) (net.Conn, error) {
+			logger.Debug("Dialing secure socks proxy", "host", host)
+			conn, err := dialer.Dial("tcp", host)
+			if err != nil {
+				return nil, fmt.Errorf("not possible to dial secure socks proxy: %w", err)
+			}
+			select {
+			case <-ctx.Done():
+				logger.Debug("Context canceled")
+				// We return `conn` anyway since we need to better test how context cancellation works
+				return conn, fmt.Errorf("context canceled: %w", err)
+			default:
+				return conn, nil
+			}
+		}))
+	}
+
+	logger.Debug("Returning dial options")
+	return dialOps, nil
 }
 
 // CustomHeadersStreamInterceptor adds custom headers to the outgoing context for each RPC call. Should work similar
diff --git a/pkg/tsdb/tempo/tempo.go b/pkg/tsdb/tempo/tempo.go
index e5c3fdb599e..311993714da 100644
--- a/pkg/tsdb/tempo/tempo.go
+++ b/pkg/tsdb/tempo/tempo.go
@@ -66,7 +66,7 @@ func newInstanceSettings(httpClientProvider *httpclient.Provider) datasource.Ins
 			return nil, err
 		}
 
-		streamingClient, err := newGrpcClient(settings, opts)
+		streamingClient, err := newGrpcClient(ctx, settings, opts)
 		if err != nil {
 			ctxLogger.Error("Failed to get gRPC client", "error", err, "function", logEntrypoint())
 			return nil, err