Plugins: Add sql support for the secure socks proxy (#64630)

pkg/tsdb/postgres/postgres.go

@@ -15,6 +15,7 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/data/sqlutil"
 
 	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/tsdb/sqleng"
 )
@@ -60,6 +61,7 @@ func (s *Service) newInstanceSettings(cfg *setting.Cfg) datasource.InstanceFacto
 			ConnMaxLifetime:     14400,
 			Timescaledb:         false,
 			ConfigurationMethod: "file-path",
+			SecureDSProxy:       false,
 		}
 
 		err := json.Unmarshal(settings.JSONData, &jsonData)
@@ -92,8 +94,17 @@ func (s *Service) newInstanceSettings(cfg *setting.Cfg) datasource.InstanceFacto
 			logger.Debug("GetEngine", "connection", cnnstr)
 		}
 
+		driverName := "postgres"
+		// register a proxy driver if the secure socks proxy is enabled
+		if cfg.IsFeatureToggleEnabled(featuremgmt.FlagSecureSocksDatasourceProxy) && cfg.SecureSocksDSProxy.Enabled && jsonData.SecureDSProxy {
+			driverName, err = createPostgresProxyDriver(&cfg.SecureSocksDSProxy, cnnstr)
+			if err != nil {
+				return "", nil
+			}
+		}
+
 		config := sqleng.DataPluginConfiguration{
-			DriverName:        "postgres",
+			DriverName:        driverName,
 			ConnectionString:  cnnstr,
 			DSInfo:            dsInfo,
 			MetricColumnTypes: []string{"UNKNOWN", "TEXT", "VARCHAR", "CHAR"},

pkg/tsdb/postgres/proxy.go

@@ -0,0 +1,107 @@
+package postgres
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"net"
+	"time"
+
+	iproxy "github.com/grafana/grafana/pkg/infra/proxy"
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/grafana/grafana/pkg/tsdb/sqleng"
+	"github.com/grafana/grafana/pkg/util"
+	"github.com/lib/pq"
+	"golang.org/x/net/proxy"
+	"xorm.io/core"
+)
+
+// createPostgresProxyDriver creates and registers a new sql driver that uses a postgres connector and updates the dialer to
+// route connections through the secure socks proxy
+func createPostgresProxyDriver(settings *setting.SecureSocksDSProxySettings, cnnstr string) (string, error) {
+	sqleng.XormDriverMu.Lock()
+	defer sqleng.XormDriverMu.Unlock()
+
+	// create a unique driver per connection string
+	hash, err := util.Md5SumString(cnnstr)
+	if err != nil {
+		return "", err
+	}
+	driverName := "postgres-proxy-" + hash
+
+	// only register the driver once
+	if core.QueryDriver(driverName) == nil {
+		connector, err := pq.NewConnector(cnnstr)
+		if err != nil {
+			return "", err
+		}
+
+		driver, err := newPostgresProxyDriver(settings, connector)
+		if err != nil {
+			return "", err
+		}
+
+		sql.Register(driverName, driver)
+		core.RegisterDriver(driverName, driver)
+	}
+	return driverName, nil
+}
+
+// postgresProxyDriver is a regular postgres driver with an updated dialer.
+// This is done because there is no way to save a dialer to the postgres driver in xorm
+type postgresProxyDriver struct {
+	c *pq.Connector
+}
+
+var _ driver.DriverContext = (*postgresProxyDriver)(nil)
+var _ core.Driver = (*postgresProxyDriver)(nil)
+
+// newPostgresProxyDriver updates the dialer for a postgres connector with a dialer that proxys connections through the secure socks proxy
+// and returns a new postgres driver to register
+func newPostgresProxyDriver(cfg *setting.SecureSocksDSProxySettings, connector *pq.Connector) (*postgresProxyDriver, error) {
+	dialer, err := iproxy.NewSecureSocksProxyContextDialer(cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	// update the postgres dialer with the proxy dialer
+	connector.Dialer(&postgresProxyDialer{d: dialer})
+
+	return &postgresProxyDriver{connector}, nil
+}
+
+// postgresProxyDialer implements the postgres dialer using a proxy dialer, as their functions differ slightly
+type postgresProxyDialer struct {
+	d proxy.Dialer
+}
+
+// Dial uses the normal proxy dial function with the updated dialer
+func (p *postgresProxyDialer) Dial(network, addr string) (c net.Conn, err error) {
+	return p.d.Dial(network, addr)
+}
+
+// DialTimeout uses the normal postgres dial timeout function with the updated dialer
+func (p *postgresProxyDialer) DialTimeout(network, address string, timeout time.Duration) (net.Conn, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	return p.d.(proxy.ContextDialer).DialContext(ctx, network, address)
+}
+
+// Parse uses the xorm postgres dialect for the driver (this has to be implemented to register the driver with xorm)
+func (d *postgresProxyDriver) Parse(a string, b string) (*core.Uri, error) {
+	sqleng.XormDriverMu.RLock()
+	defer sqleng.XormDriverMu.RUnlock()
+
+	return core.QueryDriver("postgres").Parse(a, b)
+}
+
+// OpenConnector returns the normal postgres connector that has the updated dialer context
+func (d *postgresProxyDriver) OpenConnector(name string) (driver.Connector, error) {
+	return d.c, nil
+}
+
+// Open uses the connector with the updated dialer to open a new connection
+func (d *postgresProxyDriver) Open(dsn string) (driver.Conn, error) {
+	return d.c.Connect(context.Background())
+}

pkg/tsdb/postgres/proxy_test.go

@@ -0,0 +1,70 @@
+package postgres
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/infra/proxy/proxyutil"
+	"github.com/lib/pq"
+	"github.com/stretchr/testify/require"
+	"xorm.io/core"
+)
+
+func TestPostgresProxyDriver(t *testing.T) {
+	dialect := "postgres"
+	settings := proxyutil.SetupTestSecureSocksProxySettings(t)
+	dbURL := "localhost:5432"
+	cnnstr := fmt.Sprintf("postgres://auser:password@%s/db?sslmode=disable", dbURL)
+	driverName, err := createPostgresProxyDriver(settings, cnnstr)
+	require.NoError(t, err)
+
+	t.Run("Driver should not be registered more than once", func(t *testing.T) {
+		testDriver, err := createPostgresProxyDriver(settings, cnnstr)
+		require.NoError(t, err)
+		require.Equal(t, driverName, testDriver)
+	})
+
+	t.Run("A new driver should be created for a new connection string", func(t *testing.T) {
+		testDriver, err := createPostgresProxyDriver(settings, "server=localhost;user id=sa;password=yourStrong(!)Password;database=db2")
+		require.NoError(t, err)
+		require.NotEqual(t, driverName, testDriver)
+	})
+
+	t.Run("Parse should have the same result as xorm mssql parse", func(t *testing.T) {
+		xormDriver := core.QueryDriver(dialect)
+		xormResult, err := xormDriver.Parse(dialect, cnnstr)
+		require.NoError(t, err)
+
+		xormNewDriver := core.QueryDriver(driverName)
+		xormNewResult, err := xormNewDriver.Parse(dialect, cnnstr)
+		require.NoError(t, err)
+		require.Equal(t, xormResult, xormNewResult)
+	})
+
+	t.Run("Connector should use dialer context that routes through the socks proxy to db", func(t *testing.T) {
+		connector, err := pq.NewConnector(cnnstr)
+		require.NoError(t, err)
+		driver, err := newPostgresProxyDriver(settings, connector)
+		require.NoError(t, err)
+
+		conn, err := driver.OpenConnector(cnnstr)
+		require.NoError(t, err)
+
+		_, err = conn.Connect(context.Background())
+		require.Contains(t, err.Error(), fmt.Sprintf("socks connect %s %s->%s", "tcp", settings.ProxyAddress, dbURL))
+	})
+
+	t.Run("Connector should use dialer context that routes through the socks proxy to db", func(t *testing.T) {
+		connector, err := pq.NewConnector(cnnstr)
+		require.NoError(t, err)
+		driver, err := newPostgresProxyDriver(settings, connector)
+		require.NoError(t, err)
+
+		conn, err := driver.OpenConnector(cnnstr)
+		require.NoError(t, err)
+
+		_, err = conn.Connect(context.Background())
+		require.Contains(t, err.Error(), fmt.Sprintf("socks connect %s %s->%s", "tcp", settings.ProxyAddress, dbURL))
+	})
+}