From 478f111b248a7a63442b12d0cd184ccf919ad961 Mon Sep 17 00:00:00 2001 From: sachhg Date: Tue, 21 Jul 2026 20:50:05 -0700 Subject: [PATCH 1/6] feat(scram): implement SCRAM-SHA-256 for the client and server roles Pooling requires pgpilot to own authentication on both sides, so it needs SCRAM-SHA-256 as a client (to authenticate its own backend connections) and as a server (to verify connecting clients). This implements RFC 5802/7677 without channel binding, matching PostgreSQL's plain SCRAM-SHA-256: PBKDF2-salted keys, the client/server proof and signature exchange, and constant-time verification. Promotes golang.org/x/crypto to a direct dependency for pbkdf2. Unit tests cover a full exchange, wrong-password rejection, tampering, and malformed inputs. --- go.mod | 2 + internal/scram/scram.go | 267 +++++++++++++++++++++++++++++++++++ internal/scram/scram_test.go | 129 +++++++++++++++++ 3 files changed, 398 insertions(+) create mode 100644 internal/scram/scram.go create mode 100644 internal/scram/scram_test.go diff --git a/go.mod b/go.mod index a5fdbd6..515855c 100644 --- a/go.mod +++ b/go.mod @@ -3,3 +3,5 @@ module github.com/sachhg/pgpilot go 1.22 require github.com/jackc/pgx/v5 v5.7.1 + +require golang.org/x/crypto v0.27.0 diff --git a/internal/scram/scram.go b/internal/scram/scram.go new file mode 100644 index 0000000..0399456 --- /dev/null +++ b/internal/scram/scram.go @@ -0,0 +1,267 @@ +// Package scram implements SCRAM-SHA-256 (RFC 5802 / RFC 7677) for both roles: +// the client role, used to authenticate pgpilot's own connections to a backend, +// and the server role, used to verify clients connecting to pgpilot. Channel +// binding is not supported, matching PostgreSQL's plain SCRAM-SHA-256 mechanism. +package scram + +import ( + "crypto/hmac" + "crypto/rand" + "crypto/sha256" + "crypto/subtle" + "encoding/base64" + "errors" + "fmt" + "strconv" + "strings" + + "golang.org/x/crypto/pbkdf2" +) + +// ErrAuthFailed is returned by the server when a client's proof does not verify. +var ErrAuthFailed = errors.New("scram: authentication failed") + +const ( + defaultIterations = 4096 + saltLen = 16 + nonceLen = 18 // random bytes, then base64-encoded into the message + gs2Header = "n,," + // channelBinding is base64("n,,"), the c= value for a no-channel-binding + // client-final message. + channelBinding = "biws" +) + +// Client runs the SCRAM-SHA-256 client exchange. +type Client struct { + password string + clientNonce string + clientFirstBare string + serverSignature []byte +} + +// NewClient creates a Client for the given password with a fresh client nonce. +func NewClient(password string) (*Client, error) { + nonce, err := newNonce() + if err != nil { + return nil, err + } + return &Client{password: password, clientNonce: nonce}, nil +} + +// FirstMessage returns the client-first message. PostgreSQL takes the username +// from the startup packet, so the SCRAM message carries an empty username. +func (c *Client) FirstMessage() string { + c.clientFirstBare = "n=,r=" + c.clientNonce + return gs2Header + c.clientFirstBare +} + +// HandleServerFirst processes the server-first message and returns the +// client-final message. +func (c *Client) HandleServerFirst(serverFirst string) (string, error) { + combinedNonce, salt, iterations, err := parseServerFirst(serverFirst) + if err != nil { + return "", err + } + if !strings.HasPrefix(combinedNonce, c.clientNonce) { + return "", errors.New("scram: server nonce does not extend client nonce") + } + + salted := saltedPassword(c.password, salt, iterations) + clientKey := hmacSum(salted, []byte("Client Key")) + storedKey := sha256Sum(clientKey) + + clientFinalWithoutProof := "c=" + channelBinding + ",r=" + combinedNonce + authMessage := c.clientFirstBare + "," + serverFirst + "," + clientFinalWithoutProof + + clientSignature := hmacSum(storedKey, []byte(authMessage)) + clientProof := xorBytes(clientKey, clientSignature) + + serverKey := hmacSum(salted, []byte("Server Key")) + c.serverSignature = hmacSum(serverKey, []byte(authMessage)) + + return clientFinalWithoutProof + ",p=" + base64.StdEncoding.EncodeToString(clientProof), nil +} + +// HandleServerFinal verifies the server-final message's server signature. +func (c *Client) HandleServerFinal(serverFinal string) error { + if e, ok := field(serverFinal, "e="); ok { + return fmt.Errorf("scram: server error: %s", e) + } + v, ok := field(serverFinal, "v=") + if !ok { + return errors.New("scram: malformed server-final message") + } + sig, err := base64.StdEncoding.DecodeString(v) + if err != nil { + return fmt.Errorf("scram: bad server signature encoding: %w", err) + } + if subtle.ConstantTimeCompare(sig, c.serverSignature) != 1 { + return errors.New("scram: server signature mismatch") + } + return nil +} + +// Server runs the SCRAM-SHA-256 server exchange. +type Server struct { + password string + clientNonce string + serverNonce string + salt []byte + iterations int + clientFirstBare string + serverFirst string + storedKey []byte + serverKey []byte +} + +// NewServer creates a Server that verifies clients against the given password, +// with a fresh nonce and salt. +func NewServer(password string) (*Server, error) { + nonce, err := newNonce() + if err != nil { + return nil, err + } + salt := make([]byte, saltLen) + if _, err := rand.Read(salt); err != nil { + return nil, err + } + return &Server{password: password, serverNonce: nonce, salt: salt, iterations: defaultIterations}, nil +} + +// HandleClientFirst parses the client-first message and returns the server-first +// message. +func (s *Server) HandleClientFirst(clientFirst string) (string, error) { + if strings.HasPrefix(clientFirst, "p=") { + return "", errors.New("scram: channel binding is not supported") + } + // The gs2 header ends after the second comma; the rest is client-first-bare. + c1 := strings.IndexByte(clientFirst, ',') + if c1 < 0 { + return "", errors.New("scram: malformed client-first message") + } + c2 := strings.IndexByte(clientFirst[c1+1:], ',') + if c2 < 0 { + return "", errors.New("scram: malformed client-first message") + } + s.clientFirstBare = clientFirst[c1+1+c2+1:] + + nonce, ok := field(s.clientFirstBare, "r=") + if !ok { + return "", errors.New("scram: client-first missing nonce") + } + s.clientNonce = nonce + + combined := s.clientNonce + s.serverNonce + s.serverFirst = "r=" + combined + + ",s=" + base64.StdEncoding.EncodeToString(s.salt) + + ",i=" + strconv.Itoa(s.iterations) + + salted := saltedPassword(s.password, s.salt, s.iterations) + s.storedKey = sha256Sum(hmacSum(salted, []byte("Client Key"))) + s.serverKey = hmacSum(salted, []byte("Server Key")) + + return s.serverFirst, nil +} + +// HandleClientFinal verifies the client-final message and returns the +// server-final message. It returns ErrAuthFailed if the proof does not verify. +func (s *Server) HandleClientFinal(clientFinal string) (string, error) { + if cb, ok := field(clientFinal, "c="); !ok || cb != channelBinding { + return "", errors.New("scram: unexpected channel binding") + } + if r, ok := field(clientFinal, "r="); !ok || r != s.clientNonce+s.serverNonce { + return "", errors.New("scram: nonce mismatch") + } + pEnc, ok := field(clientFinal, "p=") + if !ok { + return "", errors.New("scram: client-final missing proof") + } + clientProof, err := base64.StdEncoding.DecodeString(pEnc) + if err != nil { + return "", fmt.Errorf("scram: bad client proof encoding: %w", err) + } + if len(clientProof) != sha256.Size { + return "", errors.New("scram: bad client proof length") + } + + idx := strings.Index(clientFinal, ",p=") + if idx < 0 { + return "", errors.New("scram: malformed client-final message") + } + clientFinalWithoutProof := clientFinal[:idx] + + authMessage := s.clientFirstBare + "," + s.serverFirst + "," + clientFinalWithoutProof + clientSignature := hmacSum(s.storedKey, []byte(authMessage)) + clientKey := xorBytes(clientProof, clientSignature) + if subtle.ConstantTimeCompare(sha256Sum(clientKey), s.storedKey) != 1 { + return "", ErrAuthFailed + } + serverSignature := hmacSum(s.serverKey, []byte(authMessage)) + return "v=" + base64.StdEncoding.EncodeToString(serverSignature), nil +} + +func parseServerFirst(s string) (nonce string, salt []byte, iterations int, err error) { + nonce, ok := field(s, "r=") + if !ok { + return "", nil, 0, errors.New("scram: server-first missing nonce") + } + saltEnc, ok := field(s, "s=") + if !ok { + return "", nil, 0, errors.New("scram: server-first missing salt") + } + iterStr, ok := field(s, "i=") + if !ok { + return "", nil, 0, errors.New("scram: server-first missing iteration count") + } + salt, err = base64.StdEncoding.DecodeString(saltEnc) + if err != nil { + return "", nil, 0, fmt.Errorf("scram: bad salt encoding: %w", err) + } + iterations, err = strconv.Atoi(iterStr) + if err != nil || iterations < 1 { + return "", nil, 0, fmt.Errorf("scram: bad iteration count %q", iterStr) + } + return nonce, salt, iterations, nil +} + +// field returns the value of the first comma-separated attribute in msg that has +// the given "key=" prefix. +func field(msg, prefix string) (string, bool) { + for _, part := range strings.Split(msg, ",") { + if strings.HasPrefix(part, prefix) { + return part[len(prefix):], true + } + } + return "", false +} + +func newNonce() (string, error) { + b := make([]byte, nonceLen) + if _, err := rand.Read(b); err != nil { + return "", err + } + return base64.RawStdEncoding.EncodeToString(b), nil +} + +func saltedPassword(password string, salt []byte, iterations int) []byte { + return pbkdf2.Key([]byte(password), salt, iterations, sha256.Size, sha256.New) +} + +func hmacSum(key, msg []byte) []byte { + m := hmac.New(sha256.New, key) + m.Write(msg) + return m.Sum(nil) +} + +func sha256Sum(b []byte) []byte { + h := sha256.Sum256(b) + return h[:] +} + +func xorBytes(a, b []byte) []byte { + out := make([]byte, len(a)) + for i := range a { + out[i] = a[i] ^ b[i] + } + return out +} diff --git a/internal/scram/scram_test.go b/internal/scram/scram_test.go new file mode 100644 index 0000000..083fb55 --- /dev/null +++ b/internal/scram/scram_test.go @@ -0,0 +1,129 @@ +package scram + +import ( + "errors" + "strings" + "testing" +) + +// runExchange drives a full client/server SCRAM exchange and returns any error. +func runExchange(t *testing.T, client *Client, server *Server) error { + t.Helper() + clientFirst := client.FirstMessage() + serverFirst, err := server.HandleClientFirst(clientFirst) + if err != nil { + return err + } + clientFinal, err := client.HandleServerFirst(serverFirst) + if err != nil { + return err + } + serverFinal, err := server.HandleClientFinal(clientFinal) + if err != nil { + return err + } + return client.HandleServerFinal(serverFinal) +} + +func TestExchange_Succeeds(t *testing.T) { + client, err := NewClient("pencil") + if err != nil { + t.Fatal(err) + } + server, err := NewServer("pencil") + if err != nil { + t.Fatal(err) + } + if err := runExchange(t, client, server); err != nil { + t.Fatalf("valid exchange failed: %v", err) + } +} + +func TestExchange_Deterministic(t *testing.T) { + // Fixed nonces and salt make the exchange reproducible; it must still verify + // end to end. This locks the message construction and crypto wiring. + client := &Client{password: "hunter2", clientNonce: "clientNONCEfixed01"} + server := &Server{ + password: "hunter2", + serverNonce: "serverNONCEfixed02", + salt: []byte("0123456789abcdef"), + iterations: 4096, + } + if err := runExchange(t, client, server); err != nil { + t.Fatalf("deterministic exchange failed: %v", err) + } +} + +func TestExchange_WrongPassword(t *testing.T) { + client, _ := NewClient("wrong-password") + server, _ := NewServer("real-password") + err := runExchange(t, client, server) + if !errors.Is(err, ErrAuthFailed) { + t.Fatalf("err = %v, want ErrAuthFailed", err) + } +} + +func TestClient_RejectsTamperedServerSignature(t *testing.T) { + client, _ := NewClient("pencil") + server, _ := NewServer("pencil") + + cf := client.FirstMessage() + sf, _ := server.HandleClientFirst(cf) + cfin, _ := client.HandleServerFirst(sf) + sfin, _ := server.HandleClientFinal(cfin) + + tampered := sfin[:len(sfin)-1] + string(flipLastChar(sfin)) + if err := client.HandleServerFinal(tampered); err == nil { + t.Fatal("client accepted a tampered server signature") + } +} + +func flipLastChar(s string) byte { + c := s[len(s)-1] + if c == 'A' { + return 'B' + } + return 'A' +} + +func TestClient_ServerError(t *testing.T) { + client, _ := NewClient("pencil") + if err := client.HandleServerFinal("e=invalid-proof"); err == nil || !strings.Contains(err.Error(), "invalid-proof") { + t.Fatalf("err = %v, want a server error carrying the reason", err) + } +} + +func TestServer_RejectsChannelBinding(t *testing.T) { + server, _ := NewServer("pencil") + if _, err := server.HandleClientFirst("p=tls-server-end-point,,n=,r=abc"); err == nil { + t.Fatal("server accepted a channel-binding request it cannot satisfy") + } +} + +func TestServer_RejectsMalformedClientFirst(t *testing.T) { + server, _ := NewServer("pencil") + for _, bad := range []string{"", "n", "n,", "n,,", "n,,x=1"} { + if _, err := server.HandleClientFirst(bad); err == nil { + t.Errorf("server accepted malformed client-first %q", bad) + } + } +} + +func TestClient_RejectsMalformedServerFirst(t *testing.T) { + client, _ := NewClient("pencil") + _ = client.FirstMessage() + for _, bad := range []string{"", "r=only", "s=abc,i=1", "r=x,s=notbase64!!,i=1"} { + if _, err := client.HandleServerFirst(bad); err == nil { + t.Errorf("client accepted malformed server-first %q", bad) + } + } +} + +func TestClient_RejectsServerNonceNotExtendingClientNonce(t *testing.T) { + client := &Client{password: "pencil", clientNonce: "AAAA"} + _ = client.FirstMessage() + // server nonce that does not start with the client nonce + if _, err := client.HandleServerFirst("r=ZZZZ,s=MDEyMzQ1Njc4OWFiY2RlZg==,i=4096"); err == nil { + t.Fatal("client accepted a server nonce that does not extend its own") + } +} From deb91059471f398fd84db649205181853ed34382 Mon Sep 17 00:00:00 2001 From: sachhg Date: Tue, 21 Jul 2026 20:50:05 -0700 Subject: [PATCH 2/6] feat(config): load pgpilot JSON configuration Introduce a JSON config with the listen address, the backend to route to, the users pgpilot authenticates (name and password, used to verify clients and to open pooled backend connections), and pool sizing. Durations parse from Go duration strings; unknown fields are rejected; defaults and validation are applied on load. Config hot-reload arrives in Phase 6. --- internal/config/config.go | 130 +++++++++++++++++++++++++++++++++ internal/config/config_test.go | 94 ++++++++++++++++++++++++ 2 files changed, 224 insertions(+) create mode 100644 internal/config/config.go create mode 100644 internal/config/config_test.go diff --git a/internal/config/config.go b/internal/config/config.go new file mode 100644 index 0000000..0082533 --- /dev/null +++ b/internal/config/config.go @@ -0,0 +1,130 @@ +// Package config loads pgpilot's JSON configuration: the listen address, the +// backend to route to, the users pgpilot authenticates (used both to verify +// clients and to authenticate its own pooled backend connections), and pool +// sizing. Hot-reload arrives in Phase 6. +package config + +import ( + "bytes" + "encoding/json" + "fmt" + "os" + "time" +) + +// Config is the top-level pgpilot configuration. +type Config struct { + Listen string `json:"listen"` + Primary string `json:"primary"` + Users []User `json:"users"` + Pool Pool `json:"pool"` +} + +// User is a role pgpilot authenticates. The password is used both to verify a +// connecting client (SCRAM server side) and to open pooled backend connections +// (SCRAM client side). +type User struct { + Name string `json:"name"` + Password string `json:"password"` +} + +// Pool holds connection-pool sizing, shared by every (user, database) pool. +type Pool struct { + MaxSize int `json:"max_size"` + MaxWaiters int `json:"max_waiters"` + AcquireTimeout Duration `json:"acquire_timeout"` + IdleTimeout Duration `json:"idle_timeout"` +} + +// Duration is a time.Duration that marshals to and from a Go duration string +// (e.g. "5s", "5m") in JSON. +type Duration time.Duration + +// Std returns the value as a time.Duration. +func (d Duration) Std() time.Duration { return time.Duration(d) } + +// UnmarshalJSON parses a duration from a JSON string. +func (d *Duration) UnmarshalJSON(b []byte) error { + var s string + if err := json.Unmarshal(b, &s); err != nil { + return fmt.Errorf("config: duration must be a string: %w", err) + } + parsed, err := time.ParseDuration(s) + if err != nil { + return fmt.Errorf("config: invalid duration %q: %w", s, err) + } + *d = Duration(parsed) + return nil +} + +// MarshalJSON renders the duration as a Go duration string. +func (d Duration) MarshalJSON() ([]byte, error) { + return json.Marshal(time.Duration(d).String()) +} + +// Load reads and validates the config file at path, applying defaults. +func Load(path string) (*Config, error) { + raw, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("config: read %s: %w", path, err) + } + var c Config + dec := json.NewDecoder(bytes.NewReader(raw)) + dec.DisallowUnknownFields() + if err := dec.Decode(&c); err != nil { + return nil, fmt.Errorf("config: parse %s: %w", path, err) + } + c.applyDefaults() + if err := c.validate(); err != nil { + return nil, err + } + return &c, nil +} + +func (c *Config) applyDefaults() { + if c.Pool.MaxSize == 0 { + c.Pool.MaxSize = 10 + } + if c.Pool.AcquireTimeout == 0 { + c.Pool.AcquireTimeout = Duration(5 * time.Second) + } + if c.Pool.IdleTimeout == 0 { + c.Pool.IdleTimeout = Duration(5 * time.Minute) + } +} + +func (c *Config) validate() error { + if c.Listen == "" { + return fmt.Errorf("config: listen is required") + } + if c.Primary == "" { + return fmt.Errorf("config: primary is required") + } + if len(c.Users) == 0 { + return fmt.Errorf("config: at least one user is required") + } + seen := make(map[string]struct{}, len(c.Users)) + for _, u := range c.Users { + if u.Name == "" { + return fmt.Errorf("config: a user is missing a name") + } + if _, dup := seen[u.Name]; dup { + return fmt.Errorf("config: duplicate user %q", u.Name) + } + seen[u.Name] = struct{}{} + } + if c.Pool.MaxSize < 1 { + return fmt.Errorf("config: pool.max_size must be >= 1") + } + return nil +} + +// User returns the named user's configuration and whether it exists. +func (c *Config) User(name string) (User, bool) { + for _, u := range c.Users { + if u.Name == name { + return u, true + } + } + return User{}, false +} diff --git a/internal/config/config_test.go b/internal/config/config_test.go new file mode 100644 index 0000000..8285cdc --- /dev/null +++ b/internal/config/config_test.go @@ -0,0 +1,94 @@ +package config_test + +import ( + "os" + "path/filepath" + "testing" + "time" + + "github.com/sachhg/pgpilot/internal/config" +) + +func writeConfig(t *testing.T, body string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "pgpilot.json") + if err := os.WriteFile(path, []byte(body), 0o600); err != nil { + t.Fatal(err) + } + return path +} + +func TestLoad_Valid(t *testing.T) { + path := writeConfig(t, `{ + "listen": "127.0.0.1:6432", + "primary": "127.0.0.1:55432", + "users": [{"name": "pgpilot", "password": "pw"}], + "pool": {"max_size": 8, "max_waiters": 50, "acquire_timeout": "3s", "idle_timeout": "2m"} + }`) + c, err := config.Load(path) + if err != nil { + t.Fatalf("Load: %v", err) + } + if c.Listen != "127.0.0.1:6432" || c.Primary != "127.0.0.1:55432" { + t.Errorf("addresses = %q / %q", c.Listen, c.Primary) + } + if c.Pool.MaxSize != 8 || c.Pool.MaxWaiters != 50 { + t.Errorf("pool sizes = %+v", c.Pool) + } + if c.Pool.AcquireTimeout.Std() != 3*time.Second || c.Pool.IdleTimeout.Std() != 2*time.Minute { + t.Errorf("durations = %v / %v", c.Pool.AcquireTimeout.Std(), c.Pool.IdleTimeout.Std()) + } + if u, ok := c.User("pgpilot"); !ok || u.Password != "pw" { + t.Errorf("User lookup = %+v, %v", u, ok) + } + if _, ok := c.User("nobody"); ok { + t.Error("unexpected user found") + } +} + +func TestLoad_AppliesDefaults(t *testing.T) { + path := writeConfig(t, `{ + "listen": "127.0.0.1:6432", + "primary": "127.0.0.1:55432", + "users": [{"name": "u", "password": "p"}] + }`) + c, err := config.Load(path) + if err != nil { + t.Fatalf("Load: %v", err) + } + if c.Pool.MaxSize != 10 { + t.Errorf("default MaxSize = %d, want 10", c.Pool.MaxSize) + } + if c.Pool.AcquireTimeout.Std() != 5*time.Second { + t.Errorf("default AcquireTimeout = %v, want 5s", c.Pool.AcquireTimeout.Std()) + } + if c.Pool.IdleTimeout.Std() != 5*time.Minute { + t.Errorf("default IdleTimeout = %v, want 5m", c.Pool.IdleTimeout.Std()) + } +} + +func TestLoad_Errors(t *testing.T) { + cases := map[string]string{ + "missing listen": `{"primary":"x","users":[{"name":"u","password":"p"}]}`, + "missing primary": `{"listen":"x","users":[{"name":"u","password":"p"}]}`, + "no users": `{"listen":"x","primary":"y","users":[]}`, + "duplicate user": `{"listen":"x","primary":"y","users":[{"name":"u","password":"a"},{"name":"u","password":"b"}]}`, + "bad duration": `{"listen":"x","primary":"y","users":[{"name":"u","password":"p"}],"pool":{"acquire_timeout":"nope"}}`, + "unknown field": `{"listen":"x","primary":"y","users":[{"name":"u","password":"p"}],"bogus":1}`, + "bad json": `{`, + } + for name, body := range cases { + t.Run(name, func(t *testing.T) { + path := writeConfig(t, body) + if _, err := config.Load(path); err == nil { + t.Errorf("Load(%s) succeeded, want an error", name) + } + }) + } +} + +func TestLoad_MissingFile(t *testing.T) { + if _, err := config.Load(filepath.Join(t.TempDir(), "nope.json")); err == nil { + t.Error("Load of a missing file should fail") + } +} From 0a19f9159078f41f49fc79c550a3349be618d2ce Mon Sep 17 00:00:00 2001 From: sachhg Date: Tue, 21 Jul 2026 20:50:05 -0700 Subject: [PATCH 3/6] feat(backend): authenticate pooled backend connections via SCRAM Dial opens a TCP connection to the primary and runs the SCRAM-SHA-256 client handshake to reach ReadyForQuery, capturing the backend's ParameterStatus and BackendKeyData. Reset returns a connection to a clean state for reuse (ROLLBACK then DISCARD ALL, sent as separate simple queries so DISCARD ALL is not blocked by an implicit transaction), and Ping is a lightweight health check. An integration test authenticates against the real primary, proving the SCRAM client correct against PostgreSQL's reference implementation, and confirms a wrong password is rejected. --- internal/backend/backend.go | 191 +++++++++++++++++++ internal/backend/backend_integration_test.go | 41 ++++ 2 files changed, 232 insertions(+) create mode 100644 internal/backend/backend.go create mode 100644 internal/backend/backend_integration_test.go diff --git a/internal/backend/backend.go b/internal/backend/backend.go new file mode 100644 index 0000000..7aed454 --- /dev/null +++ b/internal/backend/backend.go @@ -0,0 +1,191 @@ +// Package backend opens and authenticates pgpilot's own connections to a +// PostgreSQL backend using SCRAM-SHA-256, leaving them ready to relay client +// traffic. The connections are what the pool manages. +package backend + +import ( + "context" + "fmt" + "net" + "slices" + "time" + + "github.com/jackc/pgx/v5/pgproto3" + + "github.com/sachhg/pgpilot/internal/scram" +) + +// protocolVersion3 is the PostgreSQL 3.0 protocol version. +const protocolVersion3 = 196608 + +const scramMechanism = "SCRAM-SHA-256" + +// Conn is an authenticated backend connection, quiescent at ReadyForQuery. It +// implements pool.Conn. +type Conn struct { + raw net.Conn + params map[string]string + pid uint32 + key uint32 +} + +// Close closes the underlying connection. +func (c *Conn) Close() error { return c.raw.Close() } + +// NetConn returns the underlying connection for byte-level relaying. The backend +// is quiescent at ReadyForQuery when Dial returns, so no protocol bytes are +// buffered ahead. +func (c *Conn) NetConn() net.Conn { return c.raw } + +// Params returns the ParameterStatus values the backend reported during startup. +func (c *Conn) Params() map[string]string { return c.params } + +// Dial opens a TCP connection to addr and authenticates as user/password for the +// given database with SCRAM-SHA-256, returning once the backend reaches +// ReadyForQuery. +func Dial(ctx context.Context, addr, user, password, database string) (*Conn, error) { + var d net.Dialer + raw, err := d.DialContext(ctx, "tcp", addr) + if err != nil { + return nil, fmt.Errorf("backend: dial %s: %w", addr, err) + } + if deadline, ok := ctx.Deadline(); ok { + _ = raw.SetDeadline(deadline) + } + c := &Conn{raw: raw, params: make(map[string]string)} + if err := c.authenticate(user, password, database); err != nil { + _ = raw.Close() + return nil, err + } + _ = raw.SetDeadline(time.Time{}) // clear the handshake deadline before relaying + return c, nil +} + +func (c *Conn) authenticate(user, password, database string) error { + fe := pgproto3.NewFrontend(c.raw, c.raw) + fe.Send(&pgproto3.StartupMessage{ + ProtocolVersion: protocolVersion3, + Parameters: map[string]string{"user": user, "database": database}, + }) + if err := fe.Flush(); err != nil { + return fmt.Errorf("backend: send startup: %w", err) + } + + var sc *scram.Client + for { + msg, err := fe.Receive() + if err != nil { + return fmt.Errorf("backend: receive during startup: %w", err) + } + switch m := msg.(type) { + case *pgproto3.AuthenticationSASL: + if !slices.Contains(m.AuthMechanisms, scramMechanism) { + return fmt.Errorf("backend: server offered unsupported SASL mechanisms %v", m.AuthMechanisms) + } + sc, err = scram.NewClient(password) + if err != nil { + return err + } + fe.Send(&pgproto3.SASLInitialResponse{AuthMechanism: scramMechanism, Data: []byte(sc.FirstMessage())}) + if err := fe.Flush(); err != nil { + return err + } + case *pgproto3.AuthenticationSASLContinue: + if sc == nil { + return fmt.Errorf("backend: unexpected SASL continue") + } + final, err := sc.HandleServerFirst(string(m.Data)) + if err != nil { + return err + } + fe.Send(&pgproto3.SASLResponse{Data: []byte(final)}) + if err := fe.Flush(); err != nil { + return err + } + case *pgproto3.AuthenticationSASLFinal: + if sc == nil { + return fmt.Errorf("backend: unexpected SASL final") + } + if err := sc.HandleServerFinal(string(m.Data)); err != nil { + return err + } + case *pgproto3.AuthenticationOk: + // Authentication done; keep reading to ReadyForQuery. + case *pgproto3.ParameterStatus: + c.params[m.Name] = m.Value + case *pgproto3.BackendKeyData: + c.pid, c.key = m.ProcessID, m.SecretKey + case *pgproto3.ReadyForQuery: + return nil + case *pgproto3.ErrorResponse: + return fmt.Errorf("backend: server rejected connection: %s %s", m.Code, m.Message) + case *pgproto3.NoticeResponse: + // ignore + default: + return fmt.Errorf("backend: unexpected message %T during startup", msg) + } + } +} + +// Reset returns the connection to a clean state for reuse by another client: it +// rolls back any open transaction and discards session state (prepared +// statements, temp tables, LISTEN registrations, session GUCs), reading to +// ReadyForQuery. Callers should reuse the connection only when Reset succeeds. +// +// ROLLBACK and DISCARD ALL are sent as separate simple queries: combined into +// one multi-statement query they would run inside an implicit transaction block, +// which DISCARD ALL forbids. +func (c *Conn) Reset(ctx context.Context) error { + fe, restore := c.frontend(ctx) + defer restore() + if err := runSimple(fe, "ROLLBACK"); err != nil { + return fmt.Errorf("backend: reset rollback: %w", err) + } + if err := runSimple(fe, "DISCARD ALL"); err != nil { + return fmt.Errorf("backend: reset discard: %w", err) + } + return nil +} + +// Ping checks the connection is alive and at a clean ReadyForQuery state by +// running a no-op query. It is suitable as a pool health check. +func (c *Conn) Ping(ctx context.Context) error { + fe, restore := c.frontend(ctx) + defer restore() + if err := runSimple(fe, ";"); err != nil { + return fmt.Errorf("backend: ping: %w", err) + } + return nil +} + +// frontend returns a pgproto3 Frontend over the connection, applying ctx's +// deadline for the duration; restore clears it. +func (c *Conn) frontend(ctx context.Context) (fe *pgproto3.Frontend, restore func()) { + fe = pgproto3.NewFrontend(c.raw, c.raw) + if deadline, ok := ctx.Deadline(); ok { + _ = c.raw.SetDeadline(deadline) + return fe, func() { _ = c.raw.SetDeadline(time.Time{}) } + } + return fe, func() {} +} + +// runSimple sends one simple query and reads until ReadyForQuery, returning an +// error only on ErrorResponse (notices and command tags are ignored). +func runSimple(fe *pgproto3.Frontend, sql string) error { + fe.Send(&pgproto3.Query{String: sql}) + if err := fe.Flush(); err != nil { + return err + } + for { + msg, err := fe.Receive() + if err != nil { + return err + } + switch m := msg.(type) { + case *pgproto3.ReadyForQuery: + return nil + case *pgproto3.ErrorResponse: + return fmt.Errorf("%s %s", m.Code, m.Message) + } + } +} diff --git a/internal/backend/backend_integration_test.go b/internal/backend/backend_integration_test.go new file mode 100644 index 0000000..9db4b5b --- /dev/null +++ b/internal/backend/backend_integration_test.go @@ -0,0 +1,41 @@ +//go:build integration + +package backend + +import ( + "context" + "testing" + "time" +) + +const primaryAddr = "127.0.0.1:55432" + +// TestDial_AuthenticatesAgainstPrimary proves the SCRAM-SHA-256 client is +// correct against PostgreSQL's reference implementation: a successful Dial means +// the real server accepted our proof. Requires `make up`. +func TestDial_AuthenticatesAgainstPrimary(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + c, err := Dial(ctx, primaryAddr, "pgpilot", "pgpilot", "pgpilot") + if err != nil { + t.Fatalf("Dial: %v (is `make up` running?)", err) + } + defer func() { _ = c.Close() }() + + if c.Params()["server_version"] == "" { + t.Error("no server_version ParameterStatus captured during startup") + } + if err := c.Ping(ctx); err != nil { + t.Errorf("Ping: %v", err) + } +} + +func TestDial_WrongPasswordRejected(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + if _, err := Dial(ctx, primaryAddr, "pgpilot", "not-the-password", "pgpilot"); err == nil { + t.Fatal("Dial succeeded with a wrong password") + } +} From 69f15a6446c0ab2c36977b360b4848534b7a012a Mon Sep 17 00:00:00 2001 From: sachhg Date: Tue, 21 Jul 2026 20:50:05 -0700 Subject: [PATCH 4/6] feat(backend): manage per-(user,database) connection pools Manager keeps one pool (internal/pool) per (user, database), each opening SCRAM-authenticated connections to the primary and health-checking them on reuse. A connection authenticated as one user against one database cannot serve another, so keying pools by both keeps identities isolated. Integration tests confirm acquire/reset/release reuses the same connection, an unknown user is rejected, and the size cap holds against the live cluster. --- internal/backend/manager.go | 129 +++++++++++++++++++ internal/backend/manager_integration_test.go | 71 ++++++++++ 2 files changed, 200 insertions(+) create mode 100644 internal/backend/manager.go create mode 100644 internal/backend/manager_integration_test.go diff --git a/internal/backend/manager.go b/internal/backend/manager.go new file mode 100644 index 0000000..e29c008 --- /dev/null +++ b/internal/backend/manager.go @@ -0,0 +1,129 @@ +package backend + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/sachhg/pgpilot/internal/pool" +) + +// PoolConfig sizes each per-(user, database) pool. +type PoolConfig struct { + MaxSize int + MaxWaiters int + AcquireTimeout time.Duration + IdleTimeout time.Duration +} + +// Manager owns one connection pool per (user, database), each holding +// SCRAM-authenticated connections to the primary. It is safe for concurrent use. +type Manager struct { + primary string + users map[string]string // user name -> password + poolCfg PoolConfig + + mu sync.Mutex + pools map[poolKey]*pool.Pool +} + +type poolKey struct { + user string + database string +} + +// NewManager creates a Manager that routes to primary and authenticates with the +// given user passwords (keyed by user name). +func NewManager(primary string, users map[string]string, cfg PoolConfig) *Manager { + return &Manager{ + primary: primary, + users: users, + poolCfg: cfg, + pools: make(map[poolKey]*pool.Pool), + } +} + +// Acquire returns a pooled, authenticated backend connection for (user, +// database). The user must be known to the manager. +func (m *Manager) Acquire(ctx context.Context, user, database string) (*Conn, error) { + p, err := m.poolFor(user, database) + if err != nil { + return nil, err + } + c, err := p.Acquire(ctx) + if err != nil { + return nil, err + } + return c.(*Conn), nil +} + +// Release returns a connection to its (user, database) pool for reuse. +func (m *Manager) Release(user, database string, c *Conn) { + if p := m.lookup(user, database); p != nil { + p.Release(c) + return + } + _ = c.Close() +} + +// Discard closes a connection and frees its slot instead of reusing it. +func (m *Manager) Discard(user, database string, c *Conn) { + if p := m.lookup(user, database); p != nil { + p.Discard(c) + return + } + _ = c.Close() +} + +// Close closes every pool. +func (m *Manager) Close() { + m.mu.Lock() + pools := make([]*pool.Pool, 0, len(m.pools)) + for _, p := range m.pools { + pools = append(pools, p) + } + m.pools = make(map[poolKey]*pool.Pool) + m.mu.Unlock() + for _, p := range pools { + _ = p.Close() + } +} + +func (m *Manager) poolFor(user, database string) (*pool.Pool, error) { + password, ok := m.users[user] + if !ok { + return nil, fmt.Errorf("backend: unknown user %q", user) + } + key := poolKey{user: user, database: database} + + m.mu.Lock() + defer m.mu.Unlock() + if p, ok := m.pools[key]; ok { + return p, nil + } + primary := m.primary + p, err := pool.New(pool.Config{ + MaxSize: m.poolCfg.MaxSize, + MaxWaiters: m.poolCfg.MaxWaiters, + AcquireTimeout: m.poolCfg.AcquireTimeout, + IdleTimeout: m.poolCfg.IdleTimeout, + New: func(ctx context.Context) (pool.Conn, error) { + return Dial(ctx, primary, user, password, database) + }, + Ping: func(ctx context.Context, c pool.Conn) error { + return c.(*Conn).Ping(ctx) + }, + }) + if err != nil { + return nil, err + } + m.pools[key] = p + return p, nil +} + +func (m *Manager) lookup(user, database string) *pool.Pool { + m.mu.Lock() + defer m.mu.Unlock() + return m.pools[poolKey{user: user, database: database}] +} diff --git a/internal/backend/manager_integration_test.go b/internal/backend/manager_integration_test.go new file mode 100644 index 0000000..d655868 --- /dev/null +++ b/internal/backend/manager_integration_test.go @@ -0,0 +1,71 @@ +//go:build integration + +package backend + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/sachhg/pgpilot/internal/pool" +) + +func testManager(maxSize int) *Manager { + return NewManager(primaryAddr, + map[string]string{"pgpilot": "pgpilot"}, + PoolConfig{MaxSize: maxSize, AcquireTimeout: 3 * time.Second, IdleTimeout: time.Minute}, + ) +} + +func TestManager_AcquireResetReleaseReuses(t *testing.T) { + m := testManager(2) + defer m.Close() + ctx := context.Background() + + c1, err := m.Acquire(ctx, "pgpilot", "pgpilot") + if err != nil { + t.Fatalf("acquire: %v (is `make up` running?)", err) + } + if err := c1.Reset(ctx); err != nil { + t.Fatalf("reset: %v", err) + } + m.Release("pgpilot", "pgpilot", c1) + + c2, err := m.Acquire(ctx, "pgpilot", "pgpilot") + if err != nil { + t.Fatalf("re-acquire: %v", err) + } + if c1 != c2 { + t.Error("expected the released backend connection to be reused") + } + m.Release("pgpilot", "pgpilot", c2) +} + +func TestManager_UnknownUser(t *testing.T) { + m := testManager(1) + defer m.Close() + if _, err := m.Acquire(context.Background(), "ghost", "pgpilot"); err == nil { + t.Fatal("acquire for an unknown user should fail") + } +} + +func TestManager_RespectsMaxSize(t *testing.T) { + m := testManager(2) + defer m.Close() + ctx := context.Background() + + a, err := m.Acquire(ctx, "pgpilot", "pgpilot") + if err != nil { + t.Fatalf("acquire a: %v", err) + } + b, err := m.Acquire(ctx, "pgpilot", "pgpilot") + if err != nil { + t.Fatalf("acquire b: %v", err) + } + if _, err := m.Acquire(ctx, "pgpilot", "pgpilot"); !errors.Is(err, pool.ErrTimeout) { + t.Fatalf("third acquire err = %v, want ErrTimeout", err) + } + m.Release("pgpilot", "pgpilot", a) + m.Release("pgpilot", "pgpilot", b) +} From a2406ce95caab0e6f1b383dfe7de22cd6c8cfd18 Mon Sep 17 00:00:00 2001 From: sachhg Date: Tue, 21 Jul 2026 20:50:05 -0700 Subject: [PATCH 5/6] docs(adr): record the auth-termination and pooling architecture Document why pooling forces pgpilot to terminate authentication on both sides (superseding ADR 0002's pass-through), the choice of a static config over auth_query or cleartext, (user, database) pool keying, and DISCARD ALL reset for reuse, each with its rejected alternatives and consequences. --- docs/adr/0004-auth-termination-and-pooling.md | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 docs/adr/0004-auth-termination-and-pooling.md diff --git a/docs/adr/0004-auth-termination-and-pooling.md b/docs/adr/0004-auth-termination-and-pooling.md new file mode 100644 index 0000000..f704729 --- /dev/null +++ b/docs/adr/0004-auth-termination-and-pooling.md @@ -0,0 +1,67 @@ +# 4. Terminate authentication to enable connection pooling + +Date: 2026-07-21 + +## Status + +Accepted. Supersedes the connection-handling decision in +[ADR 0002](0002-transparent-proxy-ssl-refusal.md), which passed authentication +through untouched. + +## Context + +Session and transaction pooling reuse a backend connection across *different* +clients. A backend connection can only be reused if it is already authenticated, +which means the pooler — not the client — has to own the backend's +authentication, and therefore has to authenticate the client itself. ADR 0002's +pass-through model, where the client authenticates end-to-end with the backend, +makes reuse impossible: a pooled connection is already past its startup and +cannot replay a new client's handshake. Pooling forces this to change. + +## Decision + +**pgpilot terminates authentication on both sides, pgbouncer-style.** + +- **Backend side (client role):** pgpilot opens its own connections to the + primary and authenticates them with **SCRAM-SHA-256** (`internal/scram` + client, driven by `internal/backend`). These connections are what the pool + holds. +- **Client side (server role):** pgpilot verifies a connecting client with + **SCRAM-SHA-256** (`internal/scram` server) before handing it a pooled + backend. (This half is wired into the proxy in the next step of the phase.) + +**Credentials come from a static config file** (`internal/config`, JSON): a list +of users with passwords, used both to verify clients and to open backend +connections. This is pgbouncer's "userlist" model. We rejected two alternatives: + +- **`auth_query`** (fetch each user's SCRAM verifier from the backend at + runtime) is more production-faithful but adds a bootstrap connection, verifier + caching, and failure modes; overkill for this project's scope. +- **Cleartext client auth** would avoid implementing the SCRAM *server* but is a + weak posture and would not exercise the interesting half of the protocol. + +**Pools are keyed by `(user, database)`.** A connection authenticated as one user +against one database cannot serve another; separate pools keep them isolated, +matching how a client's identity maps to a backend session. + +**Connections are reset before reuse** with `ROLLBACK` followed by `DISCARD ALL`, +sent as *separate* simple queries — combined into one multi-statement query they +would run inside an implicit transaction block, which `DISCARD ALL` forbids. +DISCARD ALL clears prepared statements, temp tables, LISTEN registrations, and +session GUCs, so a reused connection looks fresh. A connection that fails to +reset is discarded rather than reused. + +## Consequences + +- The SCRAM-SHA-256 implementation is validated against real PostgreSQL: the + backend connector authenticates to the primary, and wrong passwords are + rejected by the server. +- pgpilot now needs a config file (`-config`); the `-primary` flag is replaced. + Config hot-reload is a Phase 6 concern. +- Because pgpilot authenticates clients itself, `sslmode=require` still fails + (TLS termination remains future work, per ADR 0002), and cancel-request + handling is deferred — a client's CancelRequest cannot yet be mapped to a + pooled backend. +- Reset-based reuse assumes a well-behaved client that waits for its final + ReadyForQuery before disconnecting; an abrupt mid-response disconnect causes + the connection to be discarded rather than reused. From 5f5ca186252da5cbce477896b89c266dd67339e5 Mon Sep 17 00:00:00 2001 From: sachhg Date: Tue, 21 Jul 2026 20:50:05 -0700 Subject: [PATCH 6/6] docs(changelog): note the SCRAM auth and pooling foundation Record internal/scram, internal/backend (auth + reset + per-(user,database) pools), and internal/config under Unreleased, plus the x/crypto direct dependency. --- CHANGELOG.md | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5258faf..2e1f58f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,21 +29,29 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 proxy behaves identically to psql direct. Rationale recorded in `docs/adr/0002-transparent-proxy-ssl-refusal.md`. - Protocol codec (Phase 3): `internal/protocol` decodes the wire messages - pgpilot routes on (Query, Parse, Bind, Describe, Execute, Sync, Terminate, - ReadyForQuery, CommandComplete, ErrorResponse, RowDescription, DataRow) via - `jackc/pgx/v5/pgproto3`, with round-trip tests and a panic-safe, fuzzed - decoder. The proxy now relays messages frame-by-frame instead of as opaque - bytes and tracks each session's transaction status from ReadyForQuery - (`I`/`T`/`E`). `cmd/pgpilot` gains a `-log-level` flag. Rationale recorded in + pgpilot routes on via `jackc/pgx/v5/pgproto3`, with round-trip tests and a + panic-safe, fuzzed decoder. The proxy now relays messages frame-by-frame and + tracks each session's transaction status from ReadyForQuery. `cmd/pgpilot` + gains a `-log-level` flag. Rationale in `docs/adr/0003-message-aware-relay.md`. - Connection pool (Phase 4a): `internal/pool` is a bounded, health-checked pool of backend connections — configurable max size, acquire timeout, idle timeout with a background reaper, and per-connection health checks — that applies - backpressure (a bounded wait, or an immediate refusal once a waiter limit is - reached) instead of queueing acquirers without bound. It is the primitive the - session- and transaction-pooling layers will build on. + backpressure instead of queueing acquirers without bound. +- SCRAM authentication and pooling foundation (Phase 4b): `internal/scram` + implements SCRAM-SHA-256 for both the client and server roles; + `internal/backend` opens and authenticates pgpilot's own connections to the + primary with SCRAM, resets them for reuse (`ROLLBACK`, then `DISCARD ALL`), + and manages one connection pool per `(user, database)`; `internal/config` + loads pgpilot's JSON configuration (backend, users, pool sizing). The SCRAM + client is validated against real PostgreSQL. Wiring the client-facing SCRAM + authenticator and session pooling into the proxy is the remaining half of this + phase. Rationale recorded in + `docs/adr/0004-auth-termination-and-pooling.md`. ### Dependencies - Added `github.com/jackc/pgx/v5` v5.7.1 (pinned to keep the module's `go 1.22` floor) for its `pgproto3` wire-protocol codec. +- Promoted `golang.org/x/crypto` to a direct dependency for `pbkdf2`, used by + the SCRAM implementation.