diff --git a/pgproto3/backend.go b/pgproto3/backend.go index d146c3384..28cff049a 100644 --- a/pgproto3/backend.go +++ b/pgproto3/backend.go @@ -175,7 +175,13 @@ 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} } @@ -282,9 +288,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 index b41abbe10..056e547cd 100644 --- a/pgproto3/frontend.go +++ b/pgproto3/frontend.go @@ -54,6 +54,7 @@ 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 +318,9 @@ 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} + } f.partialMsg = true } @@ -452,3 +456,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 index e02457d68..c2872661d 100644 --- 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) +}