From 8f8470edafaaf52c4a0fc94abb3c8bb28697515c Mon Sep 17 00:00:00 2001 From: zene Date: Fri, 27 Sep 2024 15:17:47 +0300 Subject: [PATCH 1/3] add message body size limits in frontend and backend --- pgproto3/backend.go | 42 ++++++++++++++++++++++++++++++++++----- pgproto3/frontend.go | 39 ++++++++++++++++++++++++++++++++++++ pgproto3/frontend_test.go | 18 +++++++++++++++++ 3 files changed, 94 insertions(+), 5 deletions(-) diff --git a/pgproto3/backend.go b/pgproto3/backend.go index d146c3384..c4b004398 100644 --- 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 index b41abbe10..8510fbb18 100644 --- 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 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) +} From 0290507ff2ecace4a53e7fa9a9db8e9d8c163373 Mon Sep 17 00:00:00 2001 From: zene Date: Fri, 4 Oct 2024 09:26:37 +0300 Subject: [PATCH 2/3] remove global atomics --- pgproto3/backend.go | 23 ----------------------- pgproto3/frontend.go | 23 ----------------------- 2 files changed, 46 deletions(-) diff --git a/pgproto3/backend.go b/pgproto3/backend.go index c4b004398..492c0f306 100644 --- a/pgproto3/backend.go +++ b/pgproto3/backend.go @@ -5,27 +5,8 @@ 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 @@ -206,10 +187,6 @@ func (b *Backend) Receive() (FrontendMessage, error) { 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 } diff --git a/pgproto3/frontend.go b/pgproto3/frontend.go index 8510fbb18..a2b9d366a 100644 --- a/pgproto3/frontend.go +++ b/pgproto3/frontend.go @@ -6,27 +6,8 @@ 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 @@ -342,10 +323,6 @@ func (f *Frontend) Receive() (BackendMessage, error) { 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 } From 10e11952bd54d733c288f9a087a9164e41ec0afe Mon Sep 17 00:00:00 2001 From: zene Date: Sat, 5 Oct 2024 19:54:02 +0300 Subject: [PATCH 3/3] changed style of two comments --- pgproto3/backend.go | 4 +--- pgproto3/frontend.go | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/pgproto3/backend.go b/pgproto3/backend.go index 492c0f306..28cff049a 100644 --- a/pgproto3/backend.go +++ b/pgproto3/backend.go @@ -39,9 +39,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. + 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 diff --git a/pgproto3/frontend.go b/pgproto3/frontend.go index a2b9d366a..056e547cd 100644 --- a/pgproto3/frontend.go +++ b/pgproto3/frontend.go @@ -54,9 +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. + 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