proxy.go

package pgmux

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"runtime/debug"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/jackc/pgproto3/v2"
)

type (
	// TLSConfig holds TLS configuration for the proxy server
	TLSConfig struct {
		// Enable TLS support
		Enabled bool
		// Path to certificate file
		CertFile string
		// Path to key file
		KeyFile string
		// Required rejects connections that are not under TLS. No effect
		// unless Enabled is also set.
		Required bool
		// Config supplies advanced settings such as MinVersion. It is cloned,
		// not adopted. Certificates are loaded from CertFile/KeyFile unless it
		// already carries them, and reloaded when those files change.
		Config *tls.Config
	}

	// Limits configures runtime resource limits on the proxy.
	// Zero or negative values fall back to defaults.
	Limits struct {
		MaxConnections int // default 1024
		// MaxConnectionsPerIP caps concurrent connections per source
		// address; 0 (default) is unlimited. IPv6 is keyed by /64.
		// Requires the peer address to be the real client: behind a
		// shared egress (L4 balancer, NAT) one key covers many clients.
		MaxConnectionsPerIP int
		MaxMessageSize      int           // bytes, default 16 MiB; applies to client-facing reads only
		ClientIdleTimeout   time.Duration // kills conn if client sends nothing for this long; 0 disables (default)
	}

	// ProxyServer is a PostgreSQL proxy server that routes connections based on username
	ProxyServer struct {
		listenAddr     string
		router         Router
		mu             sync.RWMutex
		listener       net.Listener
		closing        chan struct{}
		conns          map[net.Conn]struct{}
		backendConns   map[net.Conn]struct{}
		wg             sync.WaitGroup
		sem            chan struct{}
		ipMu           sync.Mutex
		perIP          map[string]int
		cancelMu       sync.Mutex
		cancelKeys     map[cancelKey]cancelTarget
		accepted       atomic.Uint64
		rejected       atomic.Uint64
		rejectedPerIP  atomic.Uint64
		tlsConfig      *TLSConfig
		limits         *Limits
		logger         *slog.Logger
		startupRewrite StartupRewriteFunc
		resolvedTLS    *tls.Config
	}

	// StartupRewriteFunc may mutate the startup parameters sent to the backend.
	// It runs after routing and the user/database rewrites, immediately before
	// the StartupMessage is encoded. pgmux ships no policy of its own here.
	StartupRewriteFunc func(clientAddr net.Addr, params map[string]string)

	// Stats is a point-in-time view of the proxy's connection accounting.
	Stats struct {
		Current  int    // sessions being handled right now
		Accepted uint64 // connections admitted since Start
		Rejected uint64 // connections refused at either cap
		// RejectedPerIP is the subset refused by MaxConnectionsPerIP
		RejectedPerIP uint64
		Max           int // configured connection cap
	}
)

// Client-visible failure messages are fixed strings: a wrapped Go error would
// leak the backend host, port and topology. Detail goes to the log.
const (
	msgBackendUnavailable = "backend unavailable"
	msgNotAuthorized      = "no backend available for this connection"
	msgSSLRequired        = "SSL connection is required"
)

const (
	defaultMaxConnections = 1024
	defaultMaxMessageSize = 16 * 1024 * 1024
	tlsHandshakeTimeout   = 10 * time.Second
	// startupTimeout bounds how long a connection may occupy a slot without
	// sending its StartupMessage.
	startupTimeout = 10 * time.Second
	// forceCloseGrace bounds how long Shutdown waits for sessions to unwind
	// after their connections have been closed.
	forceCloseGrace = 2 * time.Second
)

// cappedChunkReader rejects Next(n) when n exceeds max, before pgproto3
// allocates buffer sized from an attacker-controlled wire length prefix.
type cappedChunkReader struct {
	inner pgproto3.ChunkReader
	max   int
}

func (c *cappedChunkReader) Next(n int) ([]byte, error) {
	if n > c.max {
		return nil, fmt.Errorf("pgmux: client message size %d exceeds max %d", n, c.max)
	}
	return c.inner.Next(n)
}

func (ps *ProxyServer) newClientChunkReader(r io.Reader) pgproto3.ChunkReader {
	return &cappedChunkReader{inner: pgproto3.NewChunkReader(r), max: ps.maxMessageSize()}
}

// NewProxyServer creates a new ProxyServer with the given listen address and router
func NewProxyServer(listenAddr string, router Router) *ProxyServer {
	return &ProxyServer{
		listenAddr:   listenAddr,
		router:       router,
		closing:      make(chan struct{}),
		conns:        make(map[net.Conn]struct{}),
		backendConns: make(map[net.Conn]struct{}),
		perIP:        make(map[string]int),
		cancelKeys:   make(map[cancelKey]cancelTarget),
	}
}

// WithTLS configures TLS support for the proxy server.
//
// All With… setters must be called before Start: they write fields the accept
// loop reads without synchronisation.
func (ps *ProxyServer) WithTLS(config *TLSConfig) *ProxyServer {
	ps.tlsConfig = config
	return ps
}

// WithLimits configures resource limits for the proxy server.
func (ps *ProxyServer) WithLimits(limits *Limits) *ProxyServer {
	ps.limits = limits
	return ps
}

// WithLogger sets the structured logger. Defaults to slog.Default().
// Per-connection detail (startup parameters, protocol chatter) is logged at
// debug level; enable it with a handler whose level is slog.LevelDebug.
func (ps *ProxyServer) WithLogger(logger *slog.Logger) *ProxyServer {
	ps.logger = logger
	return ps
}

// WithStartupRewrite installs a hook that may mutate the startup parameters
// forwarded to the backend. Nil (the default) is a no-op.
func (ps *ProxyServer) WithStartupRewrite(fn StartupRewriteFunc) *ProxyServer {
	ps.startupRewrite = fn
	return ps
}

func (ps *ProxyServer) log() *slog.Logger {
	if ps.logger != nil {
		return ps.logger
	}
	return slog.Default()
}

func (ps *ProxyServer) tlsRequired() bool {
	return ps.tlsConfig != nil && ps.tlsConfig.Required
}

func (ps *ProxyServer) maxConnections() int {
	if ps.limits != nil && ps.limits.MaxConnections > 0 {
		return ps.limits.MaxConnections
	}
	return defaultMaxConnections
}

func (ps *ProxyServer) maxConnectionsPerIP() int {
	if ps.limits != nil && ps.limits.MaxConnectionsPerIP > 0 {
		return ps.limits.MaxConnectionsPerIP
	}
	return 0
}

func (ps *ProxyServer) maxMessageSize() int {
	if ps.limits != nil && ps.limits.MaxMessageSize > 0 {
		return ps.limits.MaxMessageSize
	}
	return defaultMaxMessageSize
}

func (ps *ProxyServer) clientIdleTimeout() time.Duration {
	if ps.limits != nil && ps.limits.ClientIdleTimeout > 0 {
		return ps.limits.ClientIdleTimeout
	}
	return 0
}

// Start starts the proxy server and listens for connections
func (ps *ProxyServer) Start(ctx context.Context) error {
	// Resolve TLS once, at startup, so a bad path fails here rather than on
	// the first client connection.
	if ps.tlsConfig != nil && ps.tlsConfig.Enabled {
		resolved, err := resolveServerTLS(ps.tlsConfig, ps.log())
		if err != nil {
			return err
		}
		ps.resolvedTLS = resolved
	}
	// Requiring TLS without the means to serve it must not start.
	if ps.tlsConfig != nil && ps.tlsConfig.Required && ps.resolvedTLS == nil {
		return errors.New("TLS required but not enabled")
	}

	// Always start with a plain TCP listener
	// TLS upgrade happens after SSL negotiation
	listener, err := net.Listen("tcp", ps.listenAddr)
	if err != nil {
		return fmt.Errorf("failed to listen: %w", err)
	}
	defer listener.Close()

	ps.mu.Lock()
	ps.listener = listener
	select {
	case <-ps.closing:
		// Shutdown already ran; serving now would accept-and-reject forever.
		ps.mu.Unlock()
		return nil
	default:
	}
	ps.mu.Unlock()
	defer func() {
		ps.mu.Lock()
		ps.listener = nil
		ps.mu.Unlock()
	}()

	sem := make(chan struct{}, ps.maxConnections())
	ps.mu.Lock()
	ps.sem = sem
	ps.mu.Unlock()

	ps.log().Info("proxy listening",
		"addr", ps.listenAddr,
		"tls", ps.resolvedTLS != nil,
		"max_connections", cap(sem))

	// Also select on closing, or a Shutdown that never cancels ctx leaks
	// this goroutine.
	go func() {
		select {
		case <-ctx.Done():
			listener.Close()
		case <-ps.closing:
		}
	}()

	for {
		conn, err := listener.Accept()
		if err != nil {
			// A closed listener is terminal: retrying it spins the loop at
			// full tilt, one log line per iteration, until the process dies.
			// Check this before ctx, because the listener can be closed
			// without the context being cancelled.
			if errors.Is(err, net.ErrClosed) {
				return nil
			}
			select {
			case <-ctx.Done():
				return nil
			default:
				ps.log().Error("failed to accept connection", "error", err)
				continue
			}
		}

		// Checked before the global semaphore so one host cannot take every slot.
		ip := remoteIP(conn)
		if !ps.reserveIP(ip) {
			ps.rejected.Add(1)
			ps.rejectedPerIP.Add(1)
			go func() {
				defer ps.recoverConn(conn.RemoteAddr())
				ps.rejectOverCapacity(conn, "at max connections for this address")
			}()
			continue
		}

		select {
		case sem <- struct{}{}:
			// Must not start a session once Shutdown has stopped waiting.
			if !ps.trackConn(conn) {
				<-sem
				ps.releaseIP(ip)
				conn.Close()
				continue
			}
			// Logged only on the accepted path, so counting these lines gives
			// the connection count even when the cap is being hit.
			ps.accepted.Add(1)
			ps.log().Info("connection accepted", "client", conn.RemoteAddr().String())
			go func() {
				defer func() { <-sem }()
				defer ps.releaseIP(ip)
				defer ps.untrackConn(conn)
				defer ps.recoverConn(conn.RemoteAddr())
				ps.handleConnection(ctx, conn)
			}()
		default:
			ps.releaseIP(ip)
			ps.rejected.Add(1)
			// In a goroutine so a client that stops reading cannot stall the
			// accept loop.
			go func() {
				defer ps.recoverConn(conn.RemoteAddr())
				ps.rejectOverCapacity(conn, "at max connections")
			}()
		}
	}
}

// ipv6ClientPrefix is the smallest standard end-site allocation; keying on
// the full /128 would let anyone with a /64 hold 2^64 keys
const ipv6ClientPrefix = 64

// remoteIP returns the per-IP limit's key: the source address without port.
func remoteIP(conn net.Conn) string {
	addr, ok := conn.RemoteAddr().(*net.TCPAddr)
	if !ok {
		host, _, err := net.SplitHostPort(addrString(conn.RemoteAddr()))
		if err != nil {
			return addrString(conn.RemoteAddr())
		}
		return host
	}

	// To4 also folds IPv4-mapped IPv6 (::ffff:1.2.3.4) into one key
	if v4 := addr.IP.To4(); v4 != nil {
		return v4.String()
	}
	return addr.IP.Mask(net.CIDRMask(ipv6ClientPrefix, 128)).String()
}

// reserveIP takes one of ip's slots, reporting false when none are left.
// Every true must be paired with a releaseIP
func (ps *ProxyServer) reserveIP(ip string) bool {
	limit := ps.maxConnectionsPerIP()
	if limit <= 0 {
		return true
	}

	ps.ipMu.Lock()
	defer ps.ipMu.Unlock()
	if ps.perIP[ip] >= limit {
		return false
	}
	ps.perIP[ip]++
	return true
}

// releaseIP returns a slot, using the count rather than re-reading the limit.
func (ps *ProxyServer) releaseIP(ip string) {
	ps.ipMu.Lock()
	defer ps.ipMu.Unlock()
	// Delete at zero so the map, keyed by client addresses, stays bounded.
	if ps.perIP[ip] <= 1 {
		delete(ps.perIP, ip)
		return
	}
	ps.perIP[ip]--
}

// Stats reports the proxy's current connection accounting. Safe to call
// before Start and after Shutdown.
func (ps *ProxyServer) Stats() Stats {
	ps.mu.RLock()
	sem := ps.sem
	ps.mu.RUnlock()

	// Current reads the semaphore so the live count cannot drift from the cap.
	return Stats{
		Current:       len(sem),
		Accepted:      ps.accepted.Load(),
		Rejected:      ps.rejected.Load(),
		RejectedPerIP: ps.rejectedPerIP.Load(),
		Max:           ps.maxConnections(),
	}
}

// Addr returns the address the proxy is listening on, or nil when it is not
// listening. Use it to recover the real port after listening on :0.
func (ps *ProxyServer) Addr() net.Addr {
	ps.mu.RLock()
	defer ps.mu.RUnlock()

	if ps.listener == nil {
		return nil
	}
	return ps.listener.Addr()
}

// trackConn registers an accepted connection so Shutdown can wait for it and
// close it past the drain deadline. It reports false once shutdown has begun:
// wg.Add racing a zero-counter wg.Wait panics.
//
// The raw TCP conn is tracked, not the *tls.Conn wrapping it later — closing
// the underlying connection is what unblocks a reader parked in the TLS layer.
func (ps *ProxyServer) trackConn(conn net.Conn) bool {
	ps.mu.Lock()
	defer ps.mu.Unlock()

	select {
	case <-ps.closing:
		return false
	default:
	}

	ps.conns[conn] = struct{}{}
	ps.wg.Add(1)
	return true
}

func (ps *ProxyServer) untrackConn(conn net.Conn) {
	ps.mu.Lock()
	delete(ps.conns, conn)
	ps.mu.Unlock()
	ps.wg.Done()
}

// trackBackendConn registers a dialed backend connection so a forced Shutdown
// closes it too; it reports false once shutdown has begun, in which case the
// caller must abandon the session. Not counted in wg — the session's client
// connection already is.
func (ps *ProxyServer) trackBackendConn(conn net.Conn) bool {
	ps.mu.Lock()
	defer ps.mu.Unlock()

	select {
	case <-ps.closing:
		return false
	default:
	}

	ps.backendConns[conn] = struct{}{}
	return true
}

func (ps *ProxyServer) untrackBackendConn(conn net.Conn) {
	ps.mu.Lock()
	delete(ps.backendConns, conn)
	ps.mu.Unlock()
}

// forceCloseConns closes every live client and backend connection and reports
// how many client sessions were open. Sessions parked in Receive ignore
// context cancellation, so closing the connections under them is the only way
// to end the wait — and closing only the client side would leave a session
// parked in a backend write.
func (ps *ProxyServer) forceCloseConns() int {
	ps.mu.Lock()
	defer ps.mu.Unlock()

	n := len(ps.conns)
	for conn := range ps.conns {
		conn.Close()
	}
	for conn := range ps.backendConns {
		conn.Close()
	}
	return n
}

// Shutdown stops accepting connections and waits for the open ones to finish,
// up to ctx's deadline. Past the deadline the remaining connections are closed
// under their sessions and an error naming how many were cut is returned; a
// clean drain returns nil.
//
// Safe to call before Start, concurrently, or twice. Shutdown is terminal.
func (ps *ProxyServer) Shutdown(ctx context.Context) error {
	ps.mu.Lock()
	select {
	case <-ps.closing:
		// Another caller is already draining; fall through and wait with them.
	default:
		close(ps.closing)
	}
	if ps.listener != nil {
		// Ends the accept loop, which returns nil on net.ErrClosed.
		ps.listener.Close()
	}
	ps.mu.Unlock()

	drained := make(chan struct{})
	go func() {
		ps.wg.Wait()
		close(drained)
	}()

	select {
	case <-drained:
		return nil
	case <-ctx.Done():
	}

	n := ps.forceCloseConns()
	select {
	case <-drained:
	case <-time.After(forceCloseGrace):
	}
	return fmt.Errorf("pgmux: drain deadline passed with %d connection(s) open: %w", n, ctx.Err())
}

// recoverConn keeps one malformed connection from taking the process down.
// Every per-connection goroutine defers this.
func (ps *ProxyServer) recoverConn(addr net.Addr) {
	if r := recover(); r != nil {
		ps.log().Error("panic handling connection",
			"client", addrString(addr),
			"panic", fmt.Sprint(r),
			"stack", string(debug.Stack()))
	}
}

func addrString(addr net.Addr) string {
	if addr == nil {
		return ""
	}
	return addr.String()
}

// resolveServerTLS turns a TLSConfig into the *tls.Config used for every client
// handshake, failing if it cannot produce a usable one.
func resolveServerTLS(cfg *TLSConfig, logger *slog.Logger) (*tls.Config, error) {
	// Clone rather than adopt: the caller's config is theirs, and a *tls.Config
	// handed to two servers must not pick up one server's certificates.
	out := &tls.Config{}
	if cfg.Config != nil {
		out = cfg.Config.Clone()
	}

	// Certificates come from Config when it carries them and from the files
	// otherwise, so settings like MinVersion no longer cost the caller the
	// CertFile/KeyFile convenience — previously Config was returned verbatim
	// and the two were mutually exclusive.
	if len(out.Certificates) > 0 && cfg.CertFile != "" && logger != nil {
		// Certificates in Config win, so the files are never consulted again
		// and a renewal would go unnoticed until the certificate expired.
		logger.Warn("TLS certificates supplied in Config: CertFile and KeyFile are ignored and will not be reloaded",
			"cert_file", cfg.CertFile)
	}

	if len(out.Certificates) == 0 && out.GetCertificate == nil {
		if cfg.CertFile == "" || cfg.KeyFile == "" {
			return nil, errors.New("TLS enabled but no certificates provided")
		}
		reloader, err := newCertReloader(cfg.CertFile, cfg.KeyFile, logger)
		if err != nil {
			return nil, err
		}
		// GetCertificate only, never alongside Certificates: crypto/tls
		// consults it only when Certificates is empty or the client sent SNI,
		// so populating both would skip reloads for connections by IP.
		out.GetCertificate = reloader.GetCertificate
	}
	return out, nil
}

// sanitizeAuthError strips a backend ErrorResponse down to the fields a client
// legitimately needs during authentication: PostgreSQL populates Detail, Hint,
// Where, File, Line and Routine with server internals. Applies pre-auth only;
// once a session is established, query errors are relayed intact.
func sanitizeAuthError(msg *pgproto3.ErrorResponse) *pgproto3.ErrorResponse {
	return &pgproto3.ErrorResponse{
		Severity: msg.Severity,
		Code:     msg.Code,
		Message:  msg.Message,
	}
}

// sendFatal writes a FATAL ErrorResponse to the client. The message must be a
// fixed string, never a wrapped error: see msgBackendUnavailable.
func (ps *ProxyServer) sendFatal(conn net.Conn, code, message string) {
	errorMsg := &pgproto3.ErrorResponse{
		Severity: "FATAL",
		Code:     code,
		Message:  message,
	}
	buf, _ := errorMsg.Encode(nil)
	conn.SetWriteDeadline(time.Now().Add(500 * time.Millisecond))
	_, _ = conn.Write(buf)
	conn.SetWriteDeadline(time.Time{})
}

func (ps *ProxyServer) rejectOverCapacity(conn net.Conn, reason string) {
	defer conn.Close()
	// Debug, not Warn: an unauthenticated client must not set log volume by
	// reconnecting once a cap is reached. Stats() counts these
	ps.log().Debug("rejecting connection", "client", addrString(conn.RemoteAddr()), "reason", reason)
	errorMsg := &pgproto3.ErrorResponse{
		Severity: "FATAL",
		Code:     "53300",
		Message:  "too many connections for proxy",
	}
	buf, _ := errorMsg.Encode(nil)
	conn.SetWriteDeadline(time.Now().Add(500 * time.Millisecond))
	_, _ = conn.Write(buf)
}

func (ps *ProxyServer) handleConnection(ctx context.Context, clientConn net.Conn) {
	defer clientConn.Close()

	backend := pgproto3.NewBackend(ps.newClientChunkReader(clientConn), clientConn)

	clientConn.SetReadDeadline(time.Now().Add(startupTimeout))
	startupMsg, err := backend.ReceiveStartupMessage()
	if err != nil {
		ps.log().Debug("failed to receive startup message", "client", addrString(clientConn.RemoteAddr()), "error", err)
		return
	}

	ps.log().Debug("startup message type", "client", addrString(clientConn.RemoteAddr()), "type", fmt.Sprintf("%T", startupMsg))

	switch msg := startupMsg.(type) {
	case *pgproto3.SSLRequest:
		ps.handleSSLRequest(ctx, backend, clientConn)
	case *pgproto3.GSSEncRequest:
		// libpq sends GSSENCRequest before SSLRequest when gssencmode is prefer
		// (the default) and the client holds Kerberos credentials. Decline it
		// and read the client's next packet, same as the SSL path.
		if _, err := clientConn.Write([]byte{'N'}); err != nil {
			ps.log().Debug("failed to decline GSS encryption",
				"client", addrString(clientConn.RemoteAddr()), "error", err)
			return
		}
		clientConn.SetReadDeadline(time.Now().Add(startupTimeout))
		next, err := backend.ReceiveStartupMessage()
		if err != nil {
			ps.log().Debug("failed to receive startup message after GSS declined",
				"client", addrString(clientConn.RemoteAddr()), "error", err)
			return
		}
		clientConn.SetReadDeadline(time.Time{})
		if _, ok := next.(*pgproto3.SSLRequest); ok {
			ps.handleSSLRequest(ctx, backend, clientConn)
			return
		}
		ps.dispatchStartup(ctx, backend, next, clientConn)
	default:
		clientConn.SetReadDeadline(time.Time{})
		ps.dispatchStartup(ctx, backend, msg, clientConn)
	}
}

// dispatchStartup routes a startup-phase message once any negotiation is
// settled. Four paths reach it — a bare message, one after a declined
// SSLRequest, one over TLS, and one after a declined GSSEncRequest — so
// message handling lives here rather than being repeated, and forgotten, in
// each branch. CancelRequest was previously dropped on all four.
func (ps *ProxyServer) dispatchStartup(ctx context.Context, backend *pgproto3.Backend,
	msg pgproto3.FrontendMessage, clientConn net.Conn,
) {
	switch m := msg.(type) {
	case *pgproto3.StartupMessage:
		ps.log().Debug("protocol version", "major", m.ProtocolVersion>>16, "minor", m.ProtocolVersion&0xFFFF)
		ps.handleStartupMessage(ctx, backend, m, clientConn)
	case *pgproto3.CancelRequest:
		ps.forwardCancel(ctx, m, clientConn)
	default:
		ps.log().Debug("unexpected startup message type",
			"client", addrString(clientConn.RemoteAddr()), "type", fmt.Sprintf("%T", msg))
	}
}

// handleSSLRequest performs PostgreSQL SSL negotiation and then reads the real
// StartupMessage, over TLS if the upgrade happened.
func (ps *ProxyServer) handleSSLRequest(ctx context.Context, backend *pgproto3.Backend, clientConn net.Conn) {
	client := addrString(clientConn.RemoteAddr())

	if ps.resolvedTLS == nil {
		// TLS not configured: decline and continue in cleartext.
		if _, err := clientConn.Write([]byte{'N'}); err != nil {
			ps.log().Debug("failed to send SSL negotiation response", "client", client, "error", err)
			return
		}
		clientConn.SetReadDeadline(time.Now().Add(startupTimeout))
		startupMsg, err := backend.ReceiveStartupMessage()
		if err != nil {
			ps.log().Debug("failed to receive startup message after SSL negotiation", "client", client, "error", err)
			return
		}
		clientConn.SetReadDeadline(time.Time{})
		ps.dispatchStartup(ctx, backend, startupMsg, clientConn)
		return
	}

	if _, err := clientConn.Write([]byte{'S'}); err != nil {
		ps.log().Debug("failed to send SSL negotiation response", "client", client, "error", err)
		return
	}

	tlsConn := tls.Server(clientConn, ps.resolvedTLS)
	clientConn.SetDeadline(time.Now().Add(tlsHandshakeTimeout))
	if err := tlsConn.Handshake(); err != nil {
		ps.log().Debug("TLS handshake failed", "client", client, "error", err)
		return
	}
	clientConn.SetDeadline(time.Time{})
	ps.log().Debug("TLS connection established", "client", client)

	tlsBackend := pgproto3.NewBackend(ps.newClientChunkReader(tlsConn), tlsConn)
	tlsConn.SetReadDeadline(time.Now().Add(startupTimeout))
	startupMsg, err := tlsBackend.ReceiveStartupMessage()
	if err != nil {
		ps.log().Debug("failed to receive startup message after TLS", "client", client, "error", err)
		return
	}
	tlsConn.SetReadDeadline(time.Time{})
	ps.dispatchStartup(ctx, tlsBackend, startupMsg, tlsConn)
}

func (ps *ProxyServer) handleStartupMessage(ctx context.Context, clientBackend *pgproto3.Backend,
	startupMsg *pgproto3.StartupMessage, clientConn net.Conn,
) {
	originalUser := startupMsg.Parameters["user"]
	clientAddr := clientConn.RemoteAddr()

	// Check the connection's type, not the negotiation path: a
	// StartupMessage can also arrive with no SSLRequest at all.
	if ps.tlsRequired() {
		if _, ok := clientConn.(*tls.Conn); !ok {
			ps.log().Warn("rejecting plaintext connection",
				"client", addrString(clientAddr))
			ps.sendFatal(clientConn, "28000", msgSSLRequired)
			return
		}
	}

	// Startup parameters are attacker-supplied on a public endpoint; keep them
	// out of the default log stream.
	ps.log().Debug("startup message received", "client", addrString(clientAddr),
		"user", originalUser,
		"parameters", fmt.Sprintf("%+v", startupMsg.Parameters))

	// Route the user to get backend configuration
	backendConfig, err := ps.router.Route(ctx, originalUser)
	if err != nil {
		// Do not echo the username or routing error back: one confirms which
		// usernames exist, the other describes the topology.
		code, msg := "08001", msgBackendUnavailable
		if errors.Is(err, ErrUserNotFound) {
			code, msg = "28P01", msgNotAuthorized
		}
		ps.log().Warn("routing failed",
			"client", addrString(clientAddr), "user", originalUser, "error", err)
		ps.sendFatal(clientConn, code, msg)
		return
	}

	// Create new connection for authentication (with retries for port changes)
	addr := net.JoinHostPort(backendConfig.Host, strconv.Itoa(backendConfig.Port))
	ps.log().Debug("connecting to backend", "backend", addr, "backend_user", backendConfig.User)

	var backendConn net.Conn
	maxRetries := 3
	for attempt := range maxRetries {
		if attempt > 0 {
			// progressive wait
			time.Sleep(time.Duration(attempt) * 500 * time.Millisecond)
			backendConfig, err = ps.router.Route(ctx, originalUser)
			if err != nil {
				break
			}
			addr = net.JoinHostPort(backendConfig.Host, strconv.Itoa(backendConfig.Port))
			ps.log().Debug("retrying backend connection", "attempt", attempt+1, "backend", addr)
		}
		dialer := net.Dialer{Timeout: 5 * time.Second}
		backendConn, err = dialer.DialContext(ctx, "tcp", addr)
		if err == nil {
			break
		}
		ps.log().Warn("backend dial attempt failed", "attempt", attempt+1, "max_attempts", maxRetries, "error", err)
	}
	if err != nil {
		// The dial error names the backend host and port: log it, do not send it.
		ps.log().Error("backend dial failed",
			"client", addrString(clientAddr), "user", originalUser, "error", err)
		ps.sendFatal(clientConn, "08001", msgBackendUnavailable)
		return
	}
	defer backendConn.Close()

	// Registered after the dial loop, before TLS upgrade: the raw conn is the
	// key either way, and a forced Shutdown must be able to close it.
	if !ps.trackBackendConn(backendConn) {
		return
	}
	defer ps.untrackBackendConn(backendConn)

	if backendConfig.TLS != nil {
		upgraded, err := upgradeBackendToTLS(ctx, backendConn, backendConfig.TLS)
		if err != nil {
			ps.log().Error("backend TLS negotiation failed",
				"client", addrString(clientAddr), "error", err)
			ps.sendFatal(clientConn, "08001", msgBackendUnavailable)
			return
		}
		backendConn = upgraded
	}

	// Rewrite the identity parameters, keep all others.
	startupMsg.Parameters["user"] = backendConfig.User
	if backendConfig.Database != "" {
		startupMsg.Parameters["database"] = backendConfig.Database
	}

	if ps.startupRewrite != nil {
		ps.startupRewrite(clientAddr, startupMsg.Parameters)
	}

	serverFrontend := pgproto3.NewFrontend(pgproto3.NewChunkReader(backendConn), backendConn)

	buf, _ := startupMsg.Encode(nil)
	ps.log().Debug("forwarding startup message to backend",
		"client", addrString(clientAddr),
		"parameters", fmt.Sprintf("%+v", startupMsg.Parameters))
	_, err = backendConn.Write(buf)
	if err != nil {
		ps.log().Error("failed to send startup message to backend",
			"client", addrString(clientAddr), "error", err)
		ps.sendFatal(clientConn, "08001", msgBackendUnavailable)
		return
	}

	// Installed before the call, so a panic during authentication cannot
	// strand an entry: the registry is the one piece of per-session state
	// with no other bound.
	var issued cancelKey
	defer func() {
		if issued != (cancelKey{}) {
			ps.forgetCancelKey(issued)
		}
	}()

	if err := ps.handleAuthentication(clientBackend, serverFrontend, clientConn, backendConn, backendConfig, &issued); err != nil {
		ps.log().Warn("authentication failed", "client", addrString(clientAddr), "user", originalUser, "error", err)
		return
	}

	// Debug, not Info: "connection accepted" is the single per-connection
	// record at default level.
	ps.log().Debug("authentication successful",
		"client", addrString(clientAddr), "user", originalUser, "backend_user", backendConfig.User)
	ps.proxyMessages(ctx, clientBackend, serverFrontend, clientConn, backendConn)
}

// handleAuthentication relays the authentication exchange. It returns the
// cancel key issued to the client, which the caller must forget when the
// session ends — a non-zero key is returned even alongside an error, because
// BackendKeyData can arrive before the exchange fails.
func (ps *ProxyServer) handleAuthentication(clientBackend *pgproto3.Backend, serverFrontend *pgproto3.Frontend,
	clientConn, serverConn net.Conn, backendConfig *BackendConfig, issued *cancelKey,
) error {
	// Set a reasonable timeout for authentication
	serverConn.SetReadDeadline(time.Now().Add(30 * time.Second))
	clientConn.SetReadDeadline(time.Now().Add(30 * time.Second))
	defer func() {
		serverConn.SetReadDeadline(time.Time{})
		clientConn.SetReadDeadline(time.Time{})
	}()

authLoop:
	for {
		msg, err := serverFrontend.Receive()
		if err != nil {
			return fmt.Errorf("failed to receive from backend: %w", err)
		}

		ps.log().Debug("auth message from backend", "type", fmt.Sprintf("%T", msg))

		var buf []byte
		switch msg := msg.(type) {
		case *pgproto3.AuthenticationOk:
			ps.log().Debug("authentication OK received")
			buf, _ = msg.Encode(nil)
		case *pgproto3.AuthenticationCleartextPassword:
			buf, _ = msg.Encode(nil)
			_, err = clientConn.Write(buf)
			if err != nil {
				return fmt.Errorf("failed to send auth request to client: %w", err)
			}

			passMsg, err := clientBackend.Receive()
			if err != nil {
				return fmt.Errorf("failed to receive password: %w", err)
			}

			if pm, ok := passMsg.(*pgproto3.PasswordMessage); ok {
				buf, _ = pm.Encode(nil)
				_, err = serverConn.Write(buf)
				if err != nil {
					return fmt.Errorf("failed to send password to server: %w", err)
				}
			}
			continue
		case *pgproto3.AuthenticationMD5Password:
			buf, _ = msg.Encode(nil)
			_, err = clientConn.Write(buf)
			if err != nil {
				return fmt.Errorf("failed to send auth request to client: %w", err)
			}

			passMsg, err := clientBackend.Receive()
			if err != nil {
				return fmt.Errorf("failed to receive password: %w", err)
			}

			if pm, ok := passMsg.(*pgproto3.PasswordMessage); ok {
				buf, _ = pm.Encode(nil)
				_, err = serverConn.Write(buf)
				if err != nil {
					return fmt.Errorf("failed to send password to server: %w", err)
				}
			}
			continue
		case *pgproto3.AuthenticationSASL:
			// SASL authentication - forward server's mechanism list to client.
			buf, _ := msg.Encode(nil)
			if _, err := clientConn.Write(buf); err != nil {
				return fmt.Errorf("failed to send SASL auth to client: %w", err)
			}

			// Receive the client's SASLInitialResponse via the framed Backend
			// rather than a raw Read. A single Read assumes one TCP segment
			// equals one protocol message; that's not guaranteed and the
			// resulting mis-framing would corrupt the proof bytes.
			if err := clientBackend.SetAuthType(pgproto3.AuthTypeSASL); err != nil {
				return fmt.Errorf("failed to set SASL auth type: %w", err)
			}
			clientMsg, err := clientBackend.Receive()
			if err != nil {
				return fmt.Errorf("failed to receive SASL initial response from client: %w", err)
			}
			initResp, ok := clientMsg.(*pgproto3.SASLInitialResponse)
			if !ok {
				return fmt.Errorf("unexpected client message during SASL init: %T", clientMsg)
			}
			buf, _ = initResp.Encode(nil)
			if _, err := serverConn.Write(buf); err != nil {
				return fmt.Errorf("failed to forward client SASL response: %w", err)
			}

			// Handle the rest of the SASL handshake.
			for {
				serverMsg, err := serverFrontend.Receive()
				if err != nil {
					return fmt.Errorf("failed to receive from server during SASL: %w", err)
				}

				var buf []byte
				switch msg := serverMsg.(type) {
				case *pgproto3.AuthenticationSASLContinue:
					buf, _ = msg.Encode(nil)
				case *pgproto3.AuthenticationSASLFinal:
					buf, _ = msg.Encode(nil)
				case *pgproto3.AuthenticationOk:
					buf, _ = msg.Encode(nil)
					if _, err := clientConn.Write(buf); err != nil {
						return fmt.Errorf("failed to send auth ok to client: %w", err)
					}
					// Not the end of the exchange: ParameterStatus,
					// BackendKeyData and ReadyForQuery still follow, and
					// BackendKeyData is where the cancel key is issued.
					continue authLoop
				case *pgproto3.ErrorResponse:
					ps.log().Warn("backend rejected authentication",
						"code", msg.Code, "message", msg.Message,
						"detail", msg.Detail, "where", msg.Where)
					buf, _ = sanitizeAuthError(msg).Encode(nil)
					clientConn.Write(buf)
					return fmt.Errorf("server auth error: %s", msg.Message)
				default:
					if encoder, ok := msg.(interface{ Encode([]byte) ([]byte, error) }); ok {
						buf, _ = encoder.Encode(nil)
					}
				}

				if buf != nil {
					if _, err := clientConn.Write(buf); err != nil {
						return fmt.Errorf("failed to forward server message to client: %w", err)
					}
				}

				if _, ok := serverMsg.(*pgproto3.AuthenticationSASLContinue); ok {
					if err := clientBackend.SetAuthType(pgproto3.AuthTypeSASLContinue); err != nil {
						return fmt.Errorf("failed to set SASL continue auth type: %w", err)
					}
					contMsg, err := clientBackend.Receive()
					if err != nil {
						return fmt.Errorf("failed to read SASL response from client: %w", err)
					}
					saslResp, ok := contMsg.(*pgproto3.SASLResponse)
					if !ok {
						return fmt.Errorf("unexpected client message during SASL continue: %T", contMsg)
					}
					buf, _ := saslResp.Encode(nil)
					if _, err := serverConn.Write(buf); err != nil {
						return fmt.Errorf("failed to forward client SASL response: %w", err)
					}
				}
			}
		case *pgproto3.ParameterStatus:
			buf, _ = msg.Encode(nil)
		case *pgproto3.BackendKeyData:
			// Substituted, never relayed: the backend's key names its real
			// PID and is only usable against the backend's own address.
			key, err := ps.registerCancelKey(backendConfig, msg)
			if err != nil {
				return fmt.Errorf("failed to register cancel key: %w", err)
			}
			// A second BackendKeyData would otherwise strand the first entry
			// in the registry, pointing at a live backend PID for good.
			if *issued != (cancelKey{}) {
				ps.forgetCancelKey(*issued)
			}
			*issued = key
			buf, _ = (&pgproto3.BackendKeyData{ProcessID: key.pid, SecretKey: key.secret}).Encode(nil)
		case *pgproto3.ReadyForQuery:
			buf, _ = msg.Encode(nil)
			_, err = clientConn.Write(buf)
			if err != nil {
				return fmt.Errorf("failed to send ready to client: %w", err)
			}
			return nil
		case *pgproto3.ErrorResponse:
			ps.log().Warn("backend rejected authentication",
				"code", msg.Code, "message", msg.Message,
				"detail", msg.Detail, "where", msg.Where)
			buf, _ = sanitizeAuthError(msg).Encode(nil)
			_, err = clientConn.Write(buf)
			if err != nil {
				return fmt.Errorf("failed to send error to client: %w", err)
			}
			return fmt.Errorf("authentication error: %s", msg.Message)
		default:
			ps.log().Debug("unexpected auth message type", "type", fmt.Sprintf("%T", msg))
			continue
		}

		if buf != nil {
			_, err = clientConn.Write(buf)
			if err != nil {
				return fmt.Errorf("failed to forward auth message: %w", err)
			}
		}
	}
}

func (ps *ProxyServer) proxyMessages(ctx context.Context, clientBackend *pgproto3.Backend,
	serverFrontend *pgproto3.Frontend, clientConn, serverConn net.Conn,
) {
	errChan := make(chan error, 2)

	idle := ps.clientIdleTimeout()

	// Client to server
	go func() {
		// Always signal on exit: returning silently would leave the opposite
		// goroutine parked in Receive and the select below blocked forever,
		// leaking the connection slot.
		var err error
		defer func() { errChan <- err }()
		defer ps.recoverConn(clientConn.RemoteAddr())
		for {
			select {
			case <-ctx.Done():
				return
			default:
				if idle > 0 {
					clientConn.SetReadDeadline(time.Now().Add(idle))
				}
				var msg pgproto3.FrontendMessage
				msg, err = clientBackend.Receive()
				if err != nil {
					err = fmt.Errorf("client receive: %w", err)
					return
				}

				var buf []byte
				switch m := msg.(type) {
				case *pgproto3.Query:
					buf, _ = m.Encode(nil)
				case *pgproto3.Parse:
					buf, _ = m.Encode(nil)
				case *pgproto3.Bind:
					buf, _ = m.Encode(nil)
				case *pgproto3.Execute:
					buf, _ = m.Encode(nil)
				case *pgproto3.Describe:
					buf, _ = m.Encode(nil)
				case *pgproto3.Sync:
					buf, _ = m.Encode(nil)
				case *pgproto3.Close:
					buf, _ = m.Encode(nil)
				case *pgproto3.Terminate:
					buf, _ = m.Encode(nil)
				case *pgproto3.CopyData:
					buf, _ = m.Encode(nil)
				case *pgproto3.CopyDone:
					buf, _ = m.Encode(nil)
				case *pgproto3.CopyFail:
					buf, _ = m.Encode(nil)
				case *pgproto3.Flush:
					buf, _ = m.Encode(nil)
				default:
					ps.log().Debug("unknown client message type", "type", fmt.Sprintf("%T", m))
					continue
				}

				if buf != nil {
					if _, err = serverConn.Write(buf); err != nil {
						err = fmt.Errorf("server send: %w", err)
						return
					}
				}
			}
		}
	}()

	// Server to client
	go func() {
		var err error
		defer func() { errChan <- err }()
		defer ps.recoverConn(serverConn.RemoteAddr())
		for {
			select {
			case <-ctx.Done():
				return
			default:
				var msg pgproto3.BackendMessage
				msg, err = serverFrontend.Receive()
				if err != nil {
					err = fmt.Errorf("server receive: %w", err)
					return
				}

				var buf []byte
				switch m := msg.(type) {
				case *pgproto3.RowDescription:
					buf, _ = m.Encode(nil)
				case *pgproto3.DataRow:
					buf, _ = m.Encode(nil)
				case *pgproto3.CommandComplete:
					buf, _ = m.Encode(nil)
				case *pgproto3.ReadyForQuery:
					buf, _ = m.Encode(nil)
				case *pgproto3.ErrorResponse:
					buf, _ = m.Encode(nil)
				case *pgproto3.NoticeResponse:
					buf, _ = m.Encode(nil)
				case *pgproto3.ParameterStatus:
					buf, _ = m.Encode(nil)
				case *pgproto3.BackendKeyData:
					// Dropped, not relayed: BackendKeyData belongs to startup,
					// and forwarding one here would hand the client the real
					// backend PID that the substitution during auth withheld.
					ps.log().Debug("dropping unexpected BackendKeyData from backend")
					continue
				case *pgproto3.ParseComplete:
					buf, _ = m.Encode(nil)
				case *pgproto3.BindComplete:
					buf, _ = m.Encode(nil)
				case *pgproto3.NoData:
					buf, _ = m.Encode(nil)
				case *pgproto3.EmptyQueryResponse:
					buf, _ = m.Encode(nil)
				case *pgproto3.ParameterDescription:
					buf, _ = m.Encode(nil)
				case *pgproto3.CloseComplete:
					buf, _ = m.Encode(nil)
				case *pgproto3.NotificationResponse:
					buf, _ = m.Encode(nil)
				case *pgproto3.CopyInResponse:
					buf, _ = m.Encode(nil)
				case *pgproto3.CopyOutResponse:
					buf, _ = m.Encode(nil)
				case *pgproto3.CopyBothResponse:
					buf, _ = m.Encode(nil)
				case *pgproto3.CopyData:
					buf, _ = m.Encode(nil)
				case *pgproto3.CopyDone:
					buf, _ = m.Encode(nil)
				case *pgproto3.PortalSuspended:
					buf, _ = m.Encode(nil)
				default:
					ps.log().Debug("unknown server message type", "type", fmt.Sprintf("%T", m))
					continue
				}

				if buf != nil {
					if _, err = clientConn.Write(buf); err != nil {
						err = fmt.Errorf("client send: %w", err)
						return
					}
				}
			}
		}
	}()

	// Returning runs the caller's deferred Close on both connections, which
	// unblocks whichever direction is still in Receive.
	select {
	case err := <-errChan:
		switch {
		case err == nil:
			ps.log().Debug("proxy session ended")
		case isConnectionClosed(err) || errors.Is(err, io.ErrUnexpectedEOF):
			// Ordinary hang-ups: EOF, reset, or an idle-timeout deadline.
			ps.log().Debug("proxy session closed", "reason", err)
		default:
			ps.log().Warn("proxy session ended with error", "error", err)
		}
	case <-ctx.Done():
		ps.log().Debug("context cancelled, closing proxy connection")
	}
}

// upgradeBackendToTLS performs the PostgreSQL SSLRequest handshake on an
// existing backend connection and wraps it in a TLS client. Refuses to
// proceed if the backend declines SSL ('N') or returns an error ('E') —
// silent downgrade would defeat the point of using TLS here.
func upgradeBackendToTLS(ctx context.Context, conn net.Conn, cfg *tls.Config) (net.Conn, error) {
	// SSLRequest: int32 length (8) + int32 magic (80877103 = 0x04D2162F).
	sslRequest := []byte{0x00, 0x00, 0x00, 0x08, 0x04, 0xd2, 0x16, 0x2f}

	conn.SetDeadline(time.Now().Add(10 * time.Second))
	defer conn.SetDeadline(time.Time{})

	if _, err := conn.Write(sslRequest); err != nil {
		return nil, fmt.Errorf("send SSLRequest: %w", err)
	}

	resp := make([]byte, 1)
	if _, err := io.ReadFull(conn, resp); err != nil {
		return nil, fmt.Errorf("read SSLRequest response: %w", err)
	}
	switch resp[0] {
	case 'S':
		// Backend agrees to TLS; proceed.
	case 'N':
		return nil, fmt.Errorf("backend refused SSL (responded 'N'); refusing to send credentials over plaintext")
	case 'E':
		return nil, fmt.Errorf("backend returned error in response to SSLRequest")
	default:
		return nil, fmt.Errorf("unexpected SSLRequest response byte: %q", resp[0])
	}

	tlsConn := tls.Client(conn, cfg)
	if err := tlsConn.HandshakeContext(ctx); err != nil {
		return nil, fmt.Errorf("backend TLS handshake: %w", err)
	}
	return tlsConn, nil
}

func isConnectionClosed(err error) bool {
	if netErr, ok := err.(*net.OpError); ok {
		return netErr.Op == "read" || netErr.Op == "write"
	}
	// Errors reach us wrapped by fmt.Errorf, and tls.Conn wraps read failures
	// in its own permanentError, so the bare type assertion above is not enough.
	var opErr *net.OpError
	if errors.As(err, &opErr) {
		return opErr.Op == "read" || opErr.Op == "write"
	}
	return errors.Is(err, net.ErrClosed) || errors.Is(err, io.EOF)
}