LinkPort/internal/protocol/pgproxy/postgres_test.go

package pgproxy

import (
	"agentdb/internal/model"
	"context"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgproto3"
	"net"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"
)

func TestPostgresAuthenticationSimpleExtendedAndAudit(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	dial := func(_ context.Context, _, _ string) (net.Conn, error) {
		a, b := net.Pipe()
		go func() {
			defer b.Close()
			b.SetDeadline(time.Now().Add(5 * time.Second))
			backend := pgproto3.NewBackend(b, b)
			msg, err := backend.ReceiveStartupMessage()
			if err != nil {
				return
			}
			startup, ok := msg.(*pgproto3.StartupMessage)
			if !ok || startup.Parameters["user"] != "target" {
				return
			}
			backend.SetAuthType(pgproto3.AuthTypeCleartextPassword)
			backend.Send(&pgproto3.AuthenticationCleartextPassword{})
			backend.Flush()
			m, e := backend.Receive()
			if e != nil {
				return
			}
			pass, ok := m.(*pgproto3.PasswordMessage)
			if !ok || pass.Password != "target-secret" {
				return
			}
			backend.Send(&pgproto3.AuthenticationOk{})
			backend.Send(&pgproto3.ParameterStatus{Name: "server_version", Value: "16.0"})
			backend.Send(&pgproto3.ParameterStatus{Name: "client_encoding", Value: "UTF8"})
			backend.Send(&pgproto3.ReadyForQuery{TxStatus: 'I'})
			backend.Flush()
			row := func() {
				backend.Send(&pgproto3.RowDescription{Fields: []pgproto3.FieldDescription{{Name: []byte("n"), DataTypeOID: 23, DataTypeSize: 4, TypeModifier: -1}}})
			}
			for {
				m, e = backend.Receive()
				if e != nil {
					return
				}
				switch m.(type) {
				case *pgproto3.Query:
					row()
					backend.Send(&pgproto3.DataRow{Values: [][]byte{[]byte("42")}})
					backend.Send(&pgproto3.CommandComplete{CommandTag: []byte("SELECT 1")})
					backend.Send(&pgproto3.ReadyForQuery{TxStatus: 'I'})
				case *pgproto3.Parse:
					backend.Send(&pgproto3.ParseComplete{})
				case *pgproto3.Bind:
					backend.Send(&pgproto3.BindComplete{})
				case *pgproto3.Describe:
					row()
				case *pgproto3.Execute:
					backend.Send(&pgproto3.DataRow{Values: [][]byte{[]byte("42")}})
					backend.Send(&pgproto3.CommandComplete{CommandTag: []byte("SELECT 1")})
				case *pgproto3.Sync:
					backend.Send(&pgproto3.ReadyForQuery{TxStatus: 'I'})
				case *pgproto3.Terminate:
					return
				}
				if backend.Flush() != nil {
					return
				}
			}
		}()
		return a, nil
	}
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer listener.Close()
	p := model.Proxy{TargetHost: "pg.internal", TargetPort: 5432, TargetUser: "target", TargetPassword: "target-secret", GatewayUser: "gateway", GatewayPassword: "gateway-secret", AuthSource: "postgres"}
	var mu sync.Mutex
	var records []model.Command
	auditCtx := model.WithCommandAudit(ctx, func(c model.Command) error { mu.Lock(); records = append(records, c); mu.Unlock(); return nil })
	go func() {
		for {
			c, e := listener.Accept()
			if e != nil {
				return
			}
			go New().Serve(auditCtx, c, p, dial, nil, nil)
		}
	}()
	connect := func(password string) (*pgconn.PgConn, error) {
		cfg, _ := pgconn.ParseConfig("postgres://gateway@localhost/postgres?sslmode=disable")
		host, port, _ := net.SplitHostPort(listener.Addr().String())
		n, _ := strconv.Atoi(port)
		cfg.Host, cfg.Port, cfg.Password = host, uint16(n), password
		return pgconn.ConnectConfig(ctx, cfg)
	}
	c, err := connect("gateway-secret")
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close(ctx)
	results, err := c.Exec(ctx, "SELECT 42").ReadAll()
	if err != nil || len(results) != 1 || string(results[0].Rows[0][0]) != "42" {
		t.Fatal("simple query", err)
	}
	extended := c.ExecParams(ctx, "SELECT $1", [][]byte{[]byte("private-parameter")}, nil, nil, nil).Read()
	if extended.Err != nil {
		t.Fatal(extended.Err)
	}
	bad, e := connect("wrong")
	if e == nil {
		bad.Close(ctx)
		t.Fatal("wrong gateway password accepted")
	}
	testResult, err := New().Test(ctx, p, dial)
	if err != nil || testResult.Version != "16.0" {
		t.Fatal(testResult, err)
	}
	mu.Lock()
	defer mu.Unlock()
	query, prepare := false, false
	for _, r := range records {
		if r.Kind == "query" {
			query = true
		}
		if r.Kind == "prepare" {
			prepare = true
		}
		if strings.Contains(r.Text, "secret") || strings.Contains(r.Text, "private-parameter") {
			t.Fatal("secret in audit")
		}
	}
	if !query || !prepare {
		t.Fatal("missing SQL audit")
	}
}