add message body size limits in frontend and backend

zene · zene · commit 8f8470edafaa · 2024-09-27T15:17:47.000+03:00
diff --git a/pgproto3/backend.go b/pgproto3/backend.go
@@ -5,8 +5,27 @@ import (
 	"encoding/binary"
 	"fmt"
 	"io"
+	"sync/atomic"
 )
 
+var (
+	// When using the PostgreSQL driver, it is impossible to set a limit using
+	// the structure method, however, sometimes it is necessary to set a limit
+	// for the safety of the application. A similar functionality
+	// has been made for client messages.
+	commonMaxBackendBodyLen atomic.Uint32
+)
+
+// SetCommonMaxBackendBodyLen 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 value is 0, then no maximum is enforced.
+func SetCommonMaxBackendBodyLen(value uint32) {
+	commonMaxBackendBodyLen.Store(value)
+}
+
 // Backend acts as a server for the PostgreSQL wire protocol version 3.
 type Backend struct {
 	cr *chunkReader
@@ -39,7 +58,9 @@ 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.
+	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
@@ -175,10 +196,20 @@ func (b *Backend) Receive() (FrontendMessage, error) {
 		}
 
 		b.msgType = header[0]
-		b.bodyLen = int(binary.BigEndian.Uint32(header[1:])) - 4
+
+		msgLength := int(binary.BigEndian.Uint32(header[1:]))
+		if msgLength < 4 {
+			return nil, fmt.Errorf("invalid message length: %d", msgLength)
+		}
+
+		b.bodyLen = msgLength - 4
 		if b.maxBodyLen > 0 && b.bodyLen > b.maxBodyLen {
 			return nil, &ExceededMaxBodyLenErr{b.maxBodyLen, b.bodyLen}
 		}
+		commonMaxBodyLen := int(commonMaxBackendBodyLen.Load())
+		if commonMaxBodyLen > 0 && b.bodyLen > commonMaxBodyLen {
+			return nil, &ExceededMaxBodyLenErr{commonMaxBodyLen, b.bodyLen}
+		}
 		b.partialMsg = true
 	}
 
@@ -282,9 +313,10 @@ 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.
+// 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) {
diff --git a/pgproto3/frontend.go b/pgproto3/frontend.go
@@ -6,8 +6,27 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"sync/atomic"
 )
 
+var (
+	// When using the PostgreSQL driver, it is impossible to set a limit using
+	// the structure method, however, sometimes it is necessary to set a limit
+	// for the safety of the application. A similar functionality
+	// has been made for client messages.
+	commonMaxFrontendBodyLen atomic.Uint32
+)
+
+// SetCommonMaxFrontendBodyLen 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 a corrupted server that sends
+// messages with incorrect length, which can cause memory exhaustion.
+// The default value is 0.
+// If value is 0, then no maximum is enforced.
+func SetCommonMaxFrontendBodyLen(value uint32) {
+	commonMaxFrontendBodyLen.Store(value)
+}
+
 // Frontend acts as a client for the PostgreSQL wire protocol version 3.
 type Frontend struct {
 	cr *chunkReader
@@ -54,6 +73,9 @@ type Frontend struct {
 	portalSuspended                 PortalSuspended
 
 	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
@@ -317,6 +339,13 @@ func (f *Frontend) Receive() (BackendMessage, error) {
 		}
 
 		f.bodyLen = msgLength - 4
+		if f.maxBodyLen > 0 && f.bodyLen > f.maxBodyLen {
+			return nil, &ExceededMaxBodyLenErr{f.maxBodyLen, f.bodyLen}
+		}
+		commonMaxBodyLen := int(commonMaxFrontendBodyLen.Load())
+		if commonMaxBodyLen > 0 && f.bodyLen > commonMaxBodyLen {
+			return nil, &ExceededMaxBodyLenErr{commonMaxBodyLen, f.bodyLen}
+		}
 		f.partialMsg = true
 	}
 
@@ -452,3 +481,13 @@ func (f *Frontend) GetAuthType() uint32 {
 func (f *Frontend) ReadBufferLen() int {
 	return f.cr.wp - f.cr.rp
 }
+
+// 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 a corrupted server that sends
+// messages with incorrect length, which can cause memory exhaustion.
+// The default value is 0.
+// If maxBodyLen is 0, then no maximum is enforced.
+func (f *Frontend) SetMaxBodyLen(maxBodyLen int) {
+	f.maxBodyLen = maxBodyLen
+}
diff --git a/pgproto3/frontend_test.go b/pgproto3/frontend_test.go
@@ -115,3 +115,21 @@ func TestErrorResponse(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, want, got)
 }
+
+func TestFrontendReceiveExceededMaxBodyLen(t *testing.T) {
+	t.Parallel()
+
+	client := &interruptReader{}
+	client.push([]byte{'D', 0, 0, 10, 10})
+
+	frontend := pgproto3.NewFrontend(client, nil)
+
+	// Set max body len to 5
+	frontend.SetMaxBodyLen(5)
+
+	// Receive regular msg
+	msg, err := frontend.Receive()
+	assert.Nil(t, msg)
+	var invalidBodyLenErr *pgproto3.ExceededMaxBodyLenErr
+	assert.ErrorAs(t, err, &invalidBodyLenErr)
+}
