Add support for custom tls.Config

Allow registering a custom TLS configuration, for example to use
encrypted keys.

Based on #1066

Fixes #766
Fixes #789
Fixes #811
Fixes #849

Co-authored-by: Martin Tournoij <martin@arp242.net>

Author: daqingshu

CHANGELOG.md

@@ -28,6 +28,8 @@ newer. Previously PostgreSQL 8.4 and newer were supported.
            12 | );
                 ^
 
+- Allow using a custom `tls.Config`, for example for encrypted keys ([#1228]).
+
 - Add `PQGO_DEBUG=1` print the communication with PostgreSQL to stderr, to aid
   in debugging, testing, and bug reports ([#1223]).
 
@@ -88,6 +90,7 @@ newer. Previously PostgreSQL 8.4 and newer were supported.
 [#1223]: https://github.com/lib/pq/pull/1223
 [#1224]: https://github.com/lib/pq/pull/1224
 [#1226]: https://github.com/lib/pq/pull/1226
+[#1228]: https://github.com/lib/pq/pull/1228
 
 
 v1.10.9 (2023-04-26)

doc.go

@@ -66,6 +66,8 @@ Valid values for sslmode are:
   - verify-full - Always SSL (verify that the certification presented by
     the server was signed by a trusted CA and the server host name
     matches the one in the certificate)
+  - A custom TLS configuration registered with [RegisterTLSConfig]. These must
+    be prefixed with "pqgo-".
 
 See http://www.postgresql.org/docs/current/static/libpq-connect.html#LIBPQ-CONNSTRING
 for more information about connection string parameters.

example_test.go

@@ -1,9 +1,12 @@
 package pq_test
 
 import (
+	"crypto/tls"
+	"crypto/x509"
 	"database/sql"
 	"fmt"
 	"log"
+	"os"
 
 	"github.com/lib/pq"
 )
@@ -23,6 +26,7 @@ func ExampleNewConnector() {
 		log.Fatalf("could not start transaction: %v", err)
 	}
 	tx.Rollback()
+	// Output:
 }
 
 func ExampleConnectorWithNoticeHandler() {
@@ -44,7 +48,38 @@ func ExampleConnectorWithNoticeHandler() {
 	if _, err := db.Exec(sql); err != nil {
 		log.Fatal(err)
 	}
-
 	// Output:
 	// Notice sent: test notice
 }
+
+func ExampleRegisterTLSConfig() {
+	pem, err := os.ReadFile("testdata/init/root.crt")
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	root := x509.NewCertPool()
+	root.AppendCertsFromPEM(pem)
+
+	certs, err := tls.LoadX509KeyPair("testdata/init/postgresql.crt", "testdata/init/postgresql.key")
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	pq.RegisterTLSConfig("mytls", &tls.Config{
+		RootCAs:      root,
+		Certificates: []tls.Certificate{certs},
+		ServerName:   "postgres",
+	})
+
+	db, err := sql.Open("postgres", "host=postgres dbname=pqgo sslmode=pqgo-mytls")
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	err = db.Ping()
+	if err != nil {
+		log.Fatal(err)
+	}
+	// Output:
+}

ssl.go

@@ -10,19 +10,57 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
+	"sync"
 	"syscall"
 
 	"github.com/lib/pq/internal/pqutil"
 )
 
+// Registry for custom tls.Configs
+var (
+	tlsConfs   = make(map[string]*tls.Config)
+	tlsConfsMu sync.RWMutex
+)
+
+// RegisterTLSConfig registers a custom [tls.Config]. They are used by using
+// sslmode=pqgo-«key» in the connection string.
+//
+// Set the config to nil to remove a configuration.
+func RegisterTLSConfig(key string, config *tls.Config) error {
+	key = strings.TrimPrefix(key, "pqgo-")
+	if config == nil {
+		tlsConfsMu.Lock()
+		delete(tlsConfs, key)
+		tlsConfsMu.Unlock()
+		return nil
+	}
+
+	tlsConfsMu.Lock()
+	tlsConfs[key] = config
+	tlsConfsMu.Unlock()
+	return nil
+}
+
+func getTLSConfigClone(key string) *tls.Config {
+	tlsConfsMu.RLock()
+	if v, ok := tlsConfs[key]; ok {
+		return v.Clone()
+	}
+	tlsConfsMu.RUnlock()
+	return nil
+}
+
 // ssl generates a function to upgrade a net.Conn based on the "sslmode" and
 // related settings. The function is nil when no upgrade should take place.
 func ssl(o values) (func(net.Conn) (net.Conn, error), error) {
-	verifyCaOnly := false
-	tlsConf := tls.Config{}
-	switch mode := o["sslmode"]; mode {
+	var (
+		verifyCaOnly = false
+		tlsConf      = &tls.Config{}
+		mode         = o["sslmode"]
+	)
+	switch {
 	// "require" is the default.
-	case "", "require":
+	case mode == "" || mode == "require":
 		// We must skip TLS's own verification since it requires full
 		// verification since Go 1.3.
 		tlsConf.InsecureSkipVerify = true
@@ -42,15 +80,20 @@ func ssl(o values) (func(net.Conn) (net.Conn, error), error) {
 				delete(o, "sslrootcert")
 			}
 		}
-	case "verify-ca":
+	case mode == "verify-ca":
 		// We must skip TLS's own verification since it requires full
 		// verification since Go 1.3.
 		tlsConf.InsecureSkipVerify = true
 		verifyCaOnly = true
-	case "verify-full":
+	case mode == "verify-full":
 		tlsConf.ServerName = o["host"]
-	case "disable":
+	case mode == "disable":
 		return nil, nil
+	case strings.HasPrefix(mode, "pqgo-"):
+		tlsConf = getTLSConfigClone(mode[5:])
+		if tlsConf == nil {
+			return nil, fmt.Errorf(`pq: unknown custom sslmode %q`, mode)
+		}
 	default:
 		return nil, fmt.Errorf(
 			`pq: unsupported sslmode %q; only "require" (default), "verify-full", "verify-ca", and "disable" supported`,
@@ -67,11 +110,11 @@ func ssl(o values) (func(net.Conn) (net.Conn, error), error) {
 		tlsConf.ServerName = o["host"]
 	}
 
-	err := sslClientCertificates(&tlsConf, o)
+	err := sslClientCertificates(tlsConf, o)
 	if err != nil {
 		return nil, err
 	}
-	err = sslCertificateAuthority(&tlsConf, o)
+	err = sslCertificateAuthority(tlsConf, o)
 	if err != nil {
 		return nil, err
 	}
@@ -84,7 +127,7 @@ func ssl(o values) (func(net.Conn) (net.Conn, error), error) {
 	tlsConf.Renegotiation = tls.RenegotiateFreelyAsClient
 
 	return func(conn net.Conn) (net.Conn, error) {
-		client := tls.Client(conn, &tlsConf)
+		client := tls.Client(conn, tlsConf)
 		if verifyCaOnly {
 			err := client.Handshake()
 			if err != nil {