Repository: jackc/pgx
Pull request #1849 (2023-12-22 21:05:41)

update description cache after exec prepare
I updated the `TestPrepareIdempotency` test but it didn't actually fail without this. I'm not sure there is any errors because of this other than just inefficiency and preparing more often than we need to.


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..d5a67bea 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 {
@@ -2159,6 +2177,7 @@ func (p *Pipeline) Close() error {
 	if p.closed {
 		return p.err
 	}
+
 	p.closed = true
 
 	if p.pendingSync {
@@ -2171,7 +2190,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 {


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)
+}