Support setting up proxy via provider attribute proxy_url

Previously, the only way to configure proxying requests to Postgres
through a SOCKS5 proxy server was to set `ALL_PROXY` environment
variable. However, this could have unintended side effects and impact
other providers.

With `proxy_url` provider attribute, users can configure proxy for each
provider instance individually and make use of e.g. Terraform variables.

It is also possible to use `PGPROXY` environment variable to configure
SOCKS5 proxy for all instances of PostgreSQL provider in the current
root module. This environment variable is also used in acceptance tests
for this provider.

Signed-off-by: David van der Spek <[email protected]>

Author: take-five

postgresql/config.go

@@ -183,6 +183,7 @@ type Config struct {
 	SSLClientCert                   *ClientCertificateConfig
 	SSLRootCertPath                 string
 	GCPIAMImpersonateServiceAccount string
+	ProxyURL                        string
 }
 
 // Client struct holding connection string
@@ -290,7 +291,10 @@ func (c *Client) Connect() (*DBConnection, error) {
 		var db *sql.DB
 		var err error
 		if c.config.Scheme == "postgres" {
-			db, err = sql.Open(proxyDriverName, dsn)
+			db = sql.OpenDB(proxyConnector{
+				dsn:      dsn,
+				proxyURL: c.config.ProxyURL,
+			})
 		} else if c.config.Scheme == "gcppostgres" && c.config.GCPIAMImpersonateServiceAccount != "" {
 			db, err = openImpersonatedGCPDBConnection(context.Background(), dsn, c.config.GCPIAMImpersonateServiceAccount)
 		} else {

postgresql/provider.go

@@ -3,9 +3,10 @@ package postgresql
 import (
 	"context"
 	"fmt"
+	"os"
+
 	"github.com/aws/aws-sdk-go-v2/credentials"
 	"github.com/aws/aws-sdk-go-v2/service/sts"
-	"os"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
@@ -199,6 +200,12 @@ func Provider() *schema.Provider {
 				Description:  "Specify the expected version of PostgreSQL.",
 				ValidateFunc: validateExpectedVersion,
 			},
+			"proxy_url": {
+				Type:         schema.TypeString,
+				Optional:     true,
+				Description:  "SOCKS5 proxy URL.",
+				ValidateFunc: validation.IsURLWithScheme([]string{"socks5", "socks5h"}),
+			},
 		},
 
 		ResourcesMap: map[string]*schema.Resource{
@@ -382,6 +389,7 @@ func providerConfigure(d *schema.ResourceData) (interface{}, error) {
 		ExpectedVersion:                 version,
 		SSLRootCertPath:                 d.Get("sslrootcert").(string),
 		GCPIAMImpersonateServiceAccount: d.Get("gcp_iam_impersonate_service_account").(string),
+		ProxyURL:                        d.Get("proxy_url").(string),
 	}
 
 	if value, ok := d.GetOk("clientcert"); ok {

postgresql/proxy_driver.go

@@ -4,7 +4,10 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"fmt"
 	"net"
+	"net/url"
+	"os"
 	"time"
 
 	"github.com/lib/pq"
@@ -13,21 +16,87 @@ import (
 
 const proxyDriverName = "postgresql-proxy"
 
-type proxyDriver struct{}
+type proxyDriver struct {
+	proxyURL string
+}
 
 func (d proxyDriver) Open(name string) (driver.Conn, error) {
 	return pq.DialOpen(d, name)
 }
 
 func (d proxyDriver) Dial(network, address string) (net.Conn, error) {
-	dialer := proxy.FromEnvironment()
+	dialer, err := d.dialer()
+	if err != nil {
+		return nil, err
+	}
 	return dialer.Dial(network, address)
 }
 
 func (d proxyDriver) DialTimeout(network, address string, timeout time.Duration) (net.Conn, error) {
 	ctx, cancel := context.WithTimeout(context.TODO(), timeout)
 	defer cancel()
-	return proxy.Dial(ctx, network, address)
+
+	dialer, err := d.dialer()
+	if err != nil {
+		return nil, err
+	}
+
+	if xd, ok := dialer.(proxy.ContextDialer); ok {
+		return xd.DialContext(ctx, network, address)
+	} else {
+		return nil, fmt.Errorf("unexpected protocol error")
+	}
+}
+
+func (d proxyDriver) dialer() (proxy.Dialer, error) {
+	proxyURL := d.proxyURL
+	if proxyURL == "" {
+		proxyURL = os.Getenv("PGPROXY")
+	}
+	if proxyURL == "" {
+		return proxy.FromEnvironment(), nil
+	}
+
+	u, err := url.Parse(proxyURL)
+	if err != nil {
+		return nil, err
+	}
+
+	dialer, err := proxy.FromURL(u, proxy.Direct)
+	if err != nil {
+		return nil, err
+	}
+
+	noProxy := ""
+	if v := os.Getenv("NO_PROXY"); v != "" {
+		noProxy = v
+	}
+	if v := os.Getenv("no_proxy"); noProxy == "" && v != "" {
+		noProxy = v
+	}
+	if noProxy != "" {
+		perHost := proxy.NewPerHost(dialer, proxy.Direct)
+		perHost.AddFromString(noProxy)
+
+		dialer = perHost
+	}
+
+	return dialer, nil
+}
+
+type proxyConnector struct {
+	dsn      string
+	proxyURL string
+}
+
+var _ driver.Connector = (*proxyConnector)(nil)
+
+func (c proxyConnector) Connect(ctx context.Context) (driver.Conn, error) {
+	return c.Driver().Open(c.dsn)
+}
+
+func (c proxyConnector) Driver() driver.Driver {
+	return proxyDriver{c.proxyURL}
 }
 
 func init() {

postgresql/resource_postgresql_grant_role_test.go

@@ -141,7 +141,7 @@ func TestAccPostgresqlGrantRole(t *testing.T) {
 
 func checkGrantRole(t *testing.T, dsn, role string, grantRole string, withAdmin bool) resource.TestCheckFunc {
 	return func(s *terraform.State) error {
-		db, err := sql.Open("postgres", dsn)
+		db, err := sql.Open(proxyDriverName, dsn)
 		if err != nil {
 			t.Fatalf("could to create connection pool: %v", err)
 		}

postgresql/resource_postgresql_grant_test.go

@@ -1442,7 +1442,7 @@ func TestAccPostgresqlGrantOwnerPG15(t *testing.T) {
 			// Change the owner to the new pg_database_owner role
 			func() {
 				config := getTestConfig(t)
-				db, err := sql.Open("postgres", config.connStr(dbName))
+				db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 				if err != nil {
 					t.Fatalf("could not connect to database %s: %v", dbName, err)
 				}

postgresql/resource_postgresql_role_test.go

@@ -309,7 +309,7 @@ func testAccCheckRoleCanLogin(t *testing.T, role, password string) resource.Test
 		config := getTestConfig(t)
 		config.Username = role
 		config.Password = password
-		db, err := sql.Open("postgres", config.connStr("postgres"))
+		db, err := sql.Open(proxyDriverName, config.connStr("postgres"))
 		if err != nil {
 			return fmt.Errorf("could not open SQL connection: %v", err)
 		}

postgresql/utils_test.go

@@ -77,7 +77,7 @@ func skipIfNotSuperuser(t *testing.T) {
 
 // dbExecute is a test helper to create a pool, execute one query then close the pool
 func dbExecute(t *testing.T, dsn, query string, args ...interface{}) {
-	db, err := sql.Open("postgres", dsn)
+	db, err := sql.Open(proxyDriverName, dsn)
 	if err != nil {
 		t.Fatalf("could to create connection pool: %v", err)
 	}
@@ -147,7 +147,7 @@ func createTestTables(t *testing.T, dbSuffix string, tables []string, owner stri
 	dbName, _ := getTestDBNames(dbSuffix)
 	adminUser := config.getDatabaseUsername()
 
-	db, err := sql.Open("postgres", config.connStr(dbName))
+	db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 	if err != nil {
 		t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 	}
@@ -182,7 +182,7 @@ func createTestTables(t *testing.T, dbSuffix string, tables []string, owner stri
 
 	// In this case we need to drop table after each test.
 	return func() {
-		db, err := sql.Open("postgres", config.connStr(dbName))
+		db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 		if err != nil {
 			t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 		}
@@ -213,7 +213,7 @@ func createTestSchemas(t *testing.T, dbSuffix string, schemas []string, owner st
 	dbName, _ := getTestDBNames(dbSuffix)
 	adminUser := config.getDatabaseUsername()
 
-	db, err := sql.Open("postgres", config.connStr(dbName))
+	db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 	if err != nil {
 		t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 	}
@@ -248,7 +248,7 @@ func createTestSchemas(t *testing.T, dbSuffix string, schemas []string, owner st
 
 	// In this case we need to drop schema after each test.
 	return func() {
-		db, err := sql.Open("postgres", config.connStr(dbName))
+		db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 		if err != nil {
 			t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 		}
@@ -278,7 +278,7 @@ func createTestSequences(t *testing.T, dbSuffix string, sequences []string, owne
 	dbName, _ := getTestDBNames(dbSuffix)
 	adminUser := config.getDatabaseUsername()
 
-	db, err := sql.Open("postgres", config.connStr(dbName))
+	db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 	if err != nil {
 		t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 	}
@@ -312,7 +312,7 @@ func createTestSequences(t *testing.T, dbSuffix string, sequences []string, owne
 	}
 
 	return func() {
-		db, err := sql.Open("postgres", config.connStr(dbName))
+		db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 		if err != nil {
 			t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 		}
@@ -362,7 +362,7 @@ func connectAsTestRole(t *testing.T, role, dbName string) *sql.DB {
 	config.Username = role
 	config.Password = testRolePassword
 
-	db, err := sql.Open("postgres", config.connStr(dbName))
+	db, err := sql.Open(proxyDriverName, config.connStr(dbName))
 	if err != nil {
 		t.Fatalf("could not open connection pool for db %s: %v", dbName, err)
 	}

tests/docker-compose.yml

@@ -18,9 +18,14 @@ services:
         environment:
             POSTGRES_PASSWORD: ${PGPASSWORD}
         ports:
-            - 25432:5432
+            - "25432:5432"
         healthcheck:
           test: [ "CMD-SHELL", "pg_isready" ]
           interval: 10s
           timeout: 5s
           retries: 5
+
+    proxy:
+        image: ghcr.io/httptoolkit/docker-socks-tunnel
+        ports:
+            - "11080:1080"

tests/switch_proxy.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+export TF_ACC=true
+export PGHOST=postgres
+export PGPORT=5432
+export PGUSER=postgres
+export PGPASSWORD=postgres
+export PGSSLMODE=disable
+export PGSUPERUSER=true
+export PGPROXY=socks5://127.0.0.1:11080

tests/switch_rds.sh

@@ -30,3 +30,4 @@ export PGUSER=rds
 export PGPASSWORD=rds
 export PGSSLMODE=disable
 export PGSUPERUSER=false
+export PGPROXY=""