gopkg/pgx
2
0
Fork 0
Code Activity

Unwatch and re-watch tls

Browse Source
This commit is contained in:
Blake Embrey
2021-12-23 08:53:33 -08:00
committed by Jack Christensen
parent c0a0be876d
commit 024de4c8f3
1 changed files with 21 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pgconn.go
+21 -15
View File
@@ -230,7 +230,7 @@ func connect(ctx context.Context, config *Config, fallbackConfig *FallbackConfig
var err error var err error
network, address := NetworkAddress(fallbackConfig.Host, fallbackConfig.Port) network, address := NetworkAddress(fallbackConfig.Host, fallbackConfig.Port)
pgConn.conn, err = config.DialFunc(ctx, network, address) netConn, err := config.DialFunc(ctx, network, address)
if err != nil { if err != nil {
var netErr net.Error var netErr net.Error
if errors.As(err, &netErr) && netErr.Timeout() { if errors.As(err, &netErr) && netErr.Timeout() {
@@ -239,26 +239,28 @@ func connect(ctx context.Context, config *Config, fallbackConfig *FallbackConfig
return nil, &connectError{config: config, msg: "dial error", err: err} return nil, &connectError{config: config, msg: "dial error", err: err}
} }
pgConn.parameterStatuses = make(map[string]string) pgConn.contextWatcher = contextWatcher(netConn)
pgConn.status = connStatusConnecting
pgConn.contextWatcher = ctxwatch.NewContextWatcher(
func() { pgConn.conn.SetDeadline(time.Date(1, 1, 1, 1, 1, 1, 1, time.UTC)) },
func() { pgConn.conn.SetDeadline(time.Time{}) },
)
pgConn.contextWatcher.Watch(ctx) pgConn.contextWatcher.Watch(ctx)
defer pgConn.contextWatcher.Unwatch() defer pgConn.contextWatcher.Unwatch()
pgConn.status = connStatusConnecting
pgConn.conn = netConn
if fallbackConfig.TLSConfig != nil { if fallbackConfig.TLSConfig != nil {
tlsConn, err := startTLS(pgConn.conn, fallbackConfig.TLSConfig) tlsConn, err := startTLS(netConn, fallbackConfig.TLSConfig)
if err != nil { if err != nil {
pgConn.conn.Close() netConn.Close()
return nil, &connectError{config: config, msg: "tls error", err: err} return nil, &connectError{config: config, msg: "tls error", err: err}
} }
pgConn.contextWatcher.Unwatch()
pgConn.contextWatcher = contextWatcher(tlsConn)
pgConn.contextWatcher.Watch(ctx)
pgConn.conn = tlsConn pgConn.conn = tlsConn
} }
pgConn.parameterStatuses = make(map[string]string)
pgConn.frontend = config.BuildFrontend(pgConn.conn, pgConn.conn) pgConn.frontend = config.BuildFrontend(pgConn.conn, pgConn.conn)
startupMsg := pgproto3.StartupMessage{ startupMsg := pgproto3.StartupMessage{
@@ -346,6 +348,13 @@ func connect(ctx context.Context, config *Config, fallbackConfig *FallbackConfig
} }
} }
func contextWatcher(conn net.Conn) *ctxwatch.ContextWatcher {
return ctxwatch.NewContextWatcher(
func() { conn.SetDeadline(time.Date(1, 1, 1, 1, 1, 1, 1, time.UTC)) },
func() { conn.SetDeadline(time.Time{}) },
)
}
func startTLS(conn net.Conn, tlsConfig *tls.Config) (net.Conn, error) { func startTLS(conn net.Conn, tlsConfig *tls.Config) (net.Conn, error) {
err := binary.Write(conn, binary.BigEndian, []int32{8, 80877103}) err := binary.Write(conn, binary.BigEndian, []int32{8, 80877103})
if err != nil { if err != nil {
@@ -1709,10 +1718,7 @@ func Construct(hc *HijackedConn) (*PgConn, error) {
cleanupDone: make(chan struct{}), cleanupDone: make(chan struct{}),
} }
pgConn.contextWatcher = ctxwatch.NewContextWatcher( pgConn.contextWatcher = contextWatcher(pgConn.conn)
func() { pgConn.conn.SetDeadline(time.Date(1, 1, 1, 1, 1, 1, 1, time.UTC)) },
func() { pgConn.conn.SetDeadline(time.Time{}) },
)
return pgConn, nil return pgConn, nil
} }