fix(postgres): prevent uint32 underflow & OOM in message parsing (#10099)

* fix(postgres): prevent uint32 underflow & OOM in message parsing

* postgres: drop redundant startup guard, use maxStartupMessageSize const

The msgTotalLen < 8 check already guarantees msgLength >= 4, so the extra
msgLength < 4 guard before reading the protocol version was unreachable.
Point the startup size limit at maxStartupMessageSize instead of a literal.

* postgres: trim query terminator safely, cap pre-auth payloads

Use strings.TrimSuffix for the simple-query null terminator so a
non-null-terminated body isn't silently shortened, matching the auth
handlers. Bound password/MD5 reads with a dedicated maxAuthMessageSize
(10 KiB) instead of the 100 MiB maxMessageSize, since these payloads are
read before authentication.

---------

Co-authored-by: shangshuhan <shangshuhan@cmict.chinamobile.com>
Co-authored-by: Chris Lu <chris.lu@gmail.com>

weed/server/postgres/protocol.go

@@ -16,6 +16,22 @@ import (
 	"github.com/seaweedfs/seaweedfs/weed/util/version"
 )
 
+const (
+	// maxStartupMessageSize limits the startup handshake body to 10 KiB.
+	// Startup parameters (username, database, etc.) are always small.
+	maxStartupMessageSize = 10000
+
+	// maxMessageSize limits runtime message bodies to 100 MiB to prevent OOM.
+	// PostgreSQL protocol allows up to 1 GiB per message in theory, but real
+	// queries are rarely larger than a few megabytes.
+	maxMessageSize = 100 * 1024 * 1024 // 100 MiB
+
+	// maxAuthMessageSize caps password/MD5 payloads, which arrive before the
+	// client is authenticated. Credentials are tiny, so the pre-auth allocation
+	// stays far below maxMessageSize.
+	maxAuthMessageSize = 10 * 1024 // 10 KiB
+)
+
 // mapErrorToPostgreSQLCode maps SeaweedFS SQL engine errors to appropriate PostgreSQL error codes
 func mapErrorToPostgreSQLCode(err error) string {
 	if err == nil {
@@ -102,7 +118,14 @@ func (s *PostgreSQLServer) handleMessage(session *PostgreSQLSession) error {
 		return err
 	}
 
-	msgLength := binary.BigEndian.Uint32(length) - 4
+	rawLength := binary.BigEndian.Uint32(length)
+	if rawLength < 4 {
+		return fmt.Errorf("invalid message length: %d (must be >= 4)", rawLength)
+	}
+	msgLength := rawLength - 4
+	if msgLength > maxMessageSize {
+		return fmt.Errorf("message too large: %d bytes", msgLength)
+	}
 	msgBody := make([]byte, msgLength)
 	if msgLength > 0 {
 		_, err = io.ReadFull(session.reader, msgBody)
@@ -114,7 +137,8 @@ func (s *PostgreSQLServer) handleMessage(session *PostgreSQLSession) error {
 	// Process message based on type
 	switch msgType[0] {
 	case PG_MSG_QUERY:
-		return s.handleSimpleQuery(session, string(msgBody[:len(msgBody)-1])) // Remove null terminator
+		query := strings.TrimSuffix(string(msgBody), "\x00") // Remove null terminator if present
+		return s.handleSimpleQuery(session, query)
 	case PG_MSG_PARSE:
 		return s.handleParse(session, msgBody)
 	case PG_MSG_BIND:

weed/server/postgres/server.go

@@ -294,7 +294,11 @@ func (s *PostgreSQLServer) handleConnection(conn net.Conn) {
 
 	// Generate unique connection ID
 	connID := s.generateConnectionID()
-	secretKey := s.generateSecretKey()
+	secretKey, err := s.generateSecretKey()
+	if err != nil {
+		glog.Errorf("handleConnection: %v", err)
+		return
+	}
 
 	// Create session
 	session := &PostgreSQLSession{
@@ -328,7 +332,7 @@ func (s *PostgreSQLServer) handleConnection(conn net.Conn) {
 	glog.V(2).Infof("New PostgreSQL connection from %s (ID: %d)", conn.RemoteAddr(), connID)
 
 	// Handle startup
-	err := s.handleStartup(session)
+	err = s.handleStartup(session)
 	if err != nil {
 		// Handle common disconnection scenarios more gracefully
 		if strings.Contains(err.Error(), "client disconnected") {
@@ -395,7 +399,7 @@ func (s *PostgreSQLServer) handleStartup(session *PostgreSQLSession) error {
 		}
 
 		msgLength := msgTotalLen - 4
-		if msgLength > 10000 { // Reasonable limit for startup messages
+		if msgLength > maxStartupMessageSize {
 			return fmt.Errorf("startup message too large: %d bytes", msgLength)
 		}
 
@@ -412,7 +416,7 @@ func (s *PostgreSQLServer) handleStartup(session *PostgreSQLSession) error {
 			return fmt.Errorf("failed to read startup message: %v", err)
 		}
 
-		// Parse protocol version
+		// Parse protocol version (msgTotalLen >= 8 above guarantees msgLength >= 4)
 		protocolVersion := binary.BigEndian.Uint32(msg[0:4])
 
 		switch protocolVersion {
@@ -563,16 +567,24 @@ func (s *PostgreSQLServer) handlePasswordAuth(session *PostgreSQLSession) error
 		return err
 	}
 
-	msgLength := binary.BigEndian.Uint32(length) - 4
+	rawLength := binary.BigEndian.Uint32(length)
+	if rawLength < 4 {
+		return fmt.Errorf("password message too short: %d bytes", rawLength)
+	}
+	msgLength := rawLength - 4
+	if msgLength > maxAuthMessageSize {
+		return fmt.Errorf("password message too large: %d bytes", msgLength)
+	}
 	password := make([]byte, msgLength)
 	_, err = io.ReadFull(session.reader, password)
 	if err != nil {
 		return err
 	}
 
-	// Verify password
+	// Verify password (safe null-terminator removal)
 	expectedPassword, exists := s.config.Users[session.username]
-	if !exists || string(password[:len(password)-1]) != expectedPassword { // Remove null terminator
+	passwordStr := strings.TrimSuffix(string(password), "\x00")
+	if !exists || passwordStr != expectedPassword {
 		return s.sendError(session, "28P01", "authentication failed for user \""+session.username+"\"")
 	}
 
@@ -621,7 +633,14 @@ func (s *PostgreSQLServer) handleMD5Auth(session *PostgreSQLSession) error {
 		return err
 	}
 
-	msgLength := binary.BigEndian.Uint32(length) - 4
+	rawLength := binary.BigEndian.Uint32(length)
+	if rawLength < 4 {
+		return fmt.Errorf("MD5 response too short: %d bytes", rawLength)
+	}
+	msgLength := rawLength - 4
+	if msgLength > maxAuthMessageSize {
+		return fmt.Errorf("MD5 response too large: %d bytes", msgLength)
+	}
 	response := make([]byte, msgLength)
 	_, err = io.ReadFull(session.reader, response)
 	if err != nil {
@@ -638,7 +657,7 @@ func (s *PostgreSQLServer) handleMD5Auth(session *PostgreSQLSession) error {
 	inner := md5.Sum([]byte(expectedPassword + session.username))
 	expected := fmt.Sprintf("md5%x", md5.Sum(append([]byte(fmt.Sprintf("%x", inner)), salt...)))
 
-	if string(response[:len(response)-1]) != expected { // Remove null terminator
+	if strings.TrimSuffix(string(response), "\x00") != expected {
 		return s.sendError(session, "28P01", "authentication failed for user \""+session.username+"\"")
 	}
 
@@ -655,10 +674,12 @@ func (s *PostgreSQLServer) generateConnectionID() uint32 {
 }
 
 // generateSecretKey generates a secret key for the connection
-func (s *PostgreSQLServer) generateSecretKey() uint32 {
+func (s *PostgreSQLServer) generateSecretKey() (uint32, error) {
 	key := make([]byte, 4)
-	rand.Read(key)
-	return binary.BigEndian.Uint32(key)
+	if _, err := rand.Read(key); err != nil {
+		return 0, fmt.Errorf("generate secret key: %w", err)
+	}
+	return binary.BigEndian.Uint32(key), nil
 }
 
 // close marks the session as closed