postgres: simplify proxy code (#80121)

This commit is contained in:
Gábor Farkas
2024-01-17 09:20:27 +01:00
committed by GitHub
parent e062fa1d7e
commit 2bf56b12cf
3 changed files with 21 additions and 89 deletions

View File

@ -2,12 +2,7 @@ package postgres
import (
"context"
"crypto/md5"
"database/sql"
"database/sql/driver"
"fmt"
"net"
"slices"
"time"
sdkproxy "github.com/grafana/grafana-plugin-sdk-go/backend/proxy"
@ -15,52 +10,20 @@ import (
"golang.org/x/net/proxy"
)
// 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(cnnstr string, opts *sdkproxy.Options) (string, error) {
// create a unique driver per connection string
hash := fmt.Sprintf("%x", md5.Sum([]byte(cnnstr)))
driverName := "postgres-proxy-" + hash
// only register the driver once
if !slices.Contains(sql.Drivers(), driverName) {
connector, err := pq.NewConnector(cnnstr)
if err != nil {
return "", err
}
driver, err := newPostgresProxyDriver(connector, opts)
if err != nil {
return "", err
}
sql.Register(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)
// newPostgresProxyDriver updates the dialer for a postgres connector with a dialer that proxies connections through the secure socks proxy
// and returns a new postgres driver to register
func newPostgresProxyDriver(connector *pq.Connector, opts *sdkproxy.Options) (*postgresProxyDriver, error) {
func newPostgresProxyDialer(opts *sdkproxy.Options) (pq.Dialer, error) {
dialer, err := sdkproxy.New(opts).NewSecureSocksProxyContextDialer()
if err != nil {
return nil, err
}
// update the postgres dialer with the proxy dialer
connector.Dialer(&postgresProxyDialer{d: dialer})
return &postgresProxyDriver{connector}, nil
return &postgresProxyDialer{d: dialer}, nil
}
var _ pq.Dialer = (&postgresProxyDialer{})
// postgresProxyDialer implements the postgres dialer using a proxy dialer, as their functions differ slightly
type postgresProxyDialer struct {
d proxy.Dialer
@ -78,13 +41,3 @@ func (p *postgresProxyDialer) DialTimeout(network, address string, timeout time.
return p.d.(proxy.ContextDialer).DialContext(ctx, network, address)
}
// 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())
}