Repository: jackc/pgx
Pull request #1846 (2023-12-22 20:07:04)

Add snake_case support to RowToStructByName
RowToStructureByName doesn't support the conversion of snake_case to PascalCase, which is used in Go.


Patch:
diff --git a/conn.go b/conn.go
index d85cf127..64ae48ca 100644
--- a/conn.go
+++ b/conn.go
@@ -513,6 +513,7 @@ optionLoop:
 			if err != nil {
 				return pgconn.CommandTag{}, err
 			}
+			c.descriptionCache.Put(sd)
 		}
 
 		return c.execParams(ctx, sd, arguments)
diff --git a/pgconn/pgconn.go b/pgconn/pgconn.go
index 803b41d1..32bc0f47 100644
--- a/pgconn/pgconn.go
+++ b/pgconn/pgconn.go
@@ -2053,6 +2053,13 @@ func (p *Pipeline) Flush() error {
 
 // Sync establishes a synchronization point and flushes the queued requests.
 func (p *Pipeline) Sync() error {
+	if p.closed {
+		if p.err != nil {
+			return p.err
+		}
+		return errors.New("pipeline closed")
+	}
+
 	p.conn.frontend.SendSync(&pgproto3.Sync{})
 	err := p.Flush()
 	if err != nil {
@@ -2069,10 +2076,21 @@ func (p *Pipeline) Sync() error {
 // *PipelineSync. If an ErrorResponse is received from the server, results will be nil and err will be a *PgError. If no
 // results are available, results and err will both be nil.
 func (p *Pipeline) GetResults() (results any, err error) {
+	if p.closed {
+		if p.err != nil {
+			return nil, p.err
+		}
+		return nil, errors.New("pipeline closed")
+	}
+
 	if p.expectedReadyForQueryCount == 0 {
 		return nil, nil
 	}
 
+	return p.getResults()
+}
+
+func (p *Pipeline) getResults() (results any, err error) {
 	for {
 		msg, err := p.conn.receiveMessage()
 		if err != nil {
@@ -2099,7 +2117,8 @@ func (p *Pipeline) GetResults() (results any, err error) {
 		case *pgproto3.ParseComplete:
 			peekedMsg, err := p.conn.peekMessage()
 			if err != nil {
-				return nil, err
+				p.conn.asyncClose()
+				return nil, normalizeTimeoutError(p.ctx, err)
 			}
 			if _, ok := peekedMsg.(*pgproto3.ParameterDescription); ok {
 				return p.getResultsPrepare()
@@ -2159,6 +2178,7 @@ func (p *Pipeline) Close() error {
 	if p.closed {
 		return p.err
 	}
+
 	p.closed = true
 
 	if p.pendingSync {
@@ -2171,7 +2191,7 @@ func (p *Pipeline) Close() error {
 	}
 
 	for p.expectedReadyForQueryCount > 0 {
-		_, err := p.GetResults()
+		_, err := p.getResults()
 		if err != nil {
 			p.err = err
 			var pgErr *PgError
diff --git a/pgproto3/backend.go b/pgproto3/backend.go
index 6db77e4a..efa909c3 100644
--- a/pgproto3/backend.go
+++ b/pgproto3/backend.go
@@ -38,6 +38,7 @@ type Backend struct {
 	terminate      Terminate
 
 	bodyLen    int
+	maxBodyLen int // maxBodyLen is the maximum length of a message body in octets. If a message body exceeds this length, Receive will return an error.
 	msgType    byte
 	partialMsg bool
 	authType   uint32
@@ -158,6 +159,9 @@ func (b *Backend) Receive() (FrontendMessage, error) {
 
 		b.msgType = header[0]
 		b.bodyLen = int(binary.BigEndian.Uint32(header[1:])) - 4
+		if b.maxBodyLen > 0 && b.bodyLen > b.maxBodyLen {
+			return nil, &ExceededMaxBodyLenErr{b.maxBodyLen, b.bodyLen}
+		}
 		b.partialMsg = true
 	}
 
@@ -260,3 +264,12 @@ func (b *Backend) SetAuthType(authType uint32) error {
 
 	return nil
 }
+
+// SetMaxBodyLen sets the maximum length of a message body in octets. If a message body exceeds this length, Receive will return
+// an error. This is useful for protecting against malicious clients that send large messages with the intent of
+// causing memory exhaustion.
+// The default value is 0.
+// If maxBodyLen is 0, then no maximum is enforced.
+func (b *Backend) SetMaxBodyLen(maxBodyLen int) {
+	b.maxBodyLen = maxBodyLen
+}
diff --git a/pgproto3/pgproto3.go b/pgproto3/pgproto3.go
index ef5a5489..04be291c 100644
--- a/pgproto3/pgproto3.go
+++ b/pgproto3/pgproto3.go
@@ -70,6 +70,15 @@ func (e *writeError) Unwrap() error {
 	return e.err
 }
 
+type ExceededMaxBodyLenErr struct {
+	maxExpectedBodyLen int
+	actualBodyLen      int
+}
+
+func (e *ExceededMaxBodyLenErr) Error() string {
+	return fmt.Sprintf("invalid body length: expected at most %d, but got %d", e.maxExpectedBodyLen, e.actualBodyLen)
+}
+
 // getValueFromJSON gets the value from a protocol message representation in JSON.
 func getValueFromJSON(v map[string]string) ([]byte, error) {
 	if v == nil {
diff --git a/rows.go b/rows.go
index 73efffa0..1ad91765 100644
--- a/rows.go
+++ b/rows.go
@@ -667,7 +667,12 @@ const structTagKey = "db"
 func fieldPosByName(fldDescs []pgconn.FieldDescription, field string) (i int) {
 	i = -1
 	for i, desc := range fldDescs {
-		if strings.EqualFold(desc.Name, field) {
+
+		// Snake case support.
+		field = strings.ReplaceAll(field, "_", "")
+		descName := strings.ReplaceAll(desc.Name, "_", "")
+
+		if strings.EqualFold(descName, field) {
 			return i
 		}
 	}


Test patch:
diff --git a/conn_test.go b/conn_test.go
index 98924975..a7f7f2f8 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -482,31 +482,33 @@ func TestPrepareBadSQLFailure(t *testing.T) {
 func TestPrepareIdempotency(t *testing.T) {
 	t.Parallel()
 
-	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
-	defer closeConn(t, conn)
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
 
-	for i := 0; i < 2; i++ {
-		_, err := conn.Prepare(context.Background(), "test", "select 42::integer")
-		if err != nil {
-			t.Fatalf("%d. Unable to prepare statement: %v", i, err)
-		}
+	pgxtest.RunWithQueryExecModes(ctx, t, defaultConnTestRunner, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
+		for i := 0; i < 2; i++ {
+			_, err := conn.Prepare(context.Background(), "test", "select 42::integer")
+			if err != nil {
+				t.Fatalf("%d. Unable to prepare statement: %v", i, err)
+			}
 
-		var n int32
-		err = conn.QueryRow(context.Background(), "test").Scan(&n)
-		if err != nil {
-			t.Errorf("%d. Executing prepared statement failed: %v", i, err)
-		}
+			var n int32
+			err = conn.QueryRow(context.Background(), "test").Scan(&n)
+			if err != nil {
+				t.Errorf("%d. Executing prepared statement failed: %v", i, err)
+			}
 
-		if n != int32(42) {
-			t.Errorf("%d. Prepared statement did not return expected value: %v", i, n)
+			if n != int32(42) {
+				t.Errorf("%d. Prepared statement did not return expected value: %v", i, n)
+			}
 		}
-	}
 
-	_, err := conn.Prepare(context.Background(), "test", "select 'fail'::varchar")
-	if err == nil {
-		t.Fatalf("Prepare statement with same name but different SQL should have failed but it didn't")
-		return
-	}
+		_, err := conn.Prepare(context.Background(), "test", "select 'fail'::varchar")
+		if err == nil {
+			t.Fatalf("Prepare statement with same name but different SQL should have failed but it didn't")
+			return
+		}
+	})
 }
 
 func TestPrepareStatementCacheModes(t *testing.T) {
diff --git a/pgproto3/backend_test.go b/pgproto3/backend_test.go
index 596245dd..5655122a 100644
--- a/pgproto3/backend_test.go
+++ b/pgproto3/backend_test.go
@@ -120,3 +120,21 @@ func TestStartupMessage(t *testing.T) {
 		}
 	})
 }
+
+func TestBackendReceiveExceededMaxBodyLen(t *testing.T) {
+	t.Parallel()
+
+	server := &interruptReader{}
+	server.push([]byte{'Q', 0, 0, 10, 10})
+
+	backend := pgproto3.NewBackend(server, nil)
+
+	// Set max body len to 5
+	backend.SetMaxBodyLen(5)
+
+	// Receive regular msg
+	msg, err := backend.Receive()
+	assert.Nil(t, msg)
+	var invalidBodyLenErr *pgproto3.ExceededMaxBodyLenErr
+	assert.ErrorAs(t, err, &invalidBodyLenErr)
+}
diff --git a/rows_test.go b/rows_test.go
index be221ea4..31bd8c83 100644
--- a/rows_test.go
+++ b/rows_test.go
@@ -8,11 +8,12 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/jackc/pgx/v5/pgxtest"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
 type testRowScanner struct {
@@ -620,13 +621,14 @@ func TestRowToAddrOfStructPos(t *testing.T) {
 
 func TestRowToStructByName(t *testing.T) {
 	type person struct {
-		Last  string
-		First string
-		Age   int32
+		Last      string
+		First     string
+		Age       int32
+		AccountID string
 	}
 
 	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
-		rows, _ := conn.Query(ctx, `select 'John' as first, 'Smith' as last, n as age from generate_series(0, 9) n`)
+		rows, _ := conn.Query(ctx, `select 'John' as first, 'Smith' as last, n as age, 'd5e49d3f' as account_id from generate_series(0, 9) n`)
 		slice, err := pgx.CollectRows(rows, pgx.RowToStructByName[person])
 		assert.NoError(t, err)
 
@@ -635,6 +637,7 @@ func TestRowToStructByName(t *testing.T) {
 			assert.Equal(t, "Smith", slice[i].Last)
 			assert.Equal(t, "John", slice[i].First)
 			assert.EqualValues(t, i, slice[i].Age)
+			assert.Equal(t, "d5e49d3f", slice[i].AccountID)
 		}
 
 		// check missing fields in a returned row
@@ -643,7 +646,7 @@ func TestRowToStructByName(t *testing.T) {
 		assert.ErrorContains(t, err, "cannot find field First in returned row")
 
 		// check missing field in a destination struct
-		rows, _ = conn.Query(ctx, `select 'John' as first, 'Smith' as last, n as age, null as ignore from generate_series(0, 9) n`)
+		rows, _ = conn.Query(ctx, `select 'John' as first, 'Smith' as last, n as age, 'd5e49d3f' as account_id, null as ignore from generate_series(0, 9) n`)
 		_, err = pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[person])
 		assert.ErrorContains(t, err, "struct doesn't have corresponding row field ignore")
 	})