From 4e23d157fda041d3901771f78989d1659d77f555 Mon Sep 17 00:00:00 2001
From: Paul Gerste <79814126+paul-gerste-sonarsource@users.noreply.github.com>
Date: Tue, 21 May 2024 12:01:03 +0200
Subject: [PATCH] Enforce size limits

---
 buf.go  | 7 +++++++
 conn.go | 4 ++++
 copy.go | 4 ++++
 3 files changed, 15 insertions(+)

diff --git a/buf.go b/buf.go
index 4b0a0a8f7..c381e050b 100644
--- a/buf.go
+++ b/buf.go
@@ -3,6 +3,7 @@ package pq
 import (
 	"bytes"
 	"encoding/binary"
+	"math"
 
 	"github.com/lib/pq/oid"
 )
@@ -79,12 +80,18 @@ func (b *writeBuf) bytes(v []byte) {
 
 func (b *writeBuf) wrap() []byte {
 	p := b.buf[b.pos:]
+	if len(p) > math.MaxUint32 {
+		panic("message too large")
+	}
 	binary.BigEndian.PutUint32(p, uint32(len(p)))
 	return b.buf
 }
 
 func (b *writeBuf) next(c byte) {
 	p := b.buf[b.pos:]
+	if len(p) > math.MaxUint32 {
+		panic("message too large")
+	}
 	binary.BigEndian.PutUint32(p, uint32(len(p)))
 	b.pos = len(b.buf) + 1
 	b.buf = append(b.buf, c, 0, 0, 0, 0)
diff --git a/conn.go b/conn.go
index bc0983608..cdc2d664c 100644
--- a/conn.go
+++ b/conn.go
@@ -12,6 +12,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 	"net"
 	"os"
 	"os/user"
@@ -820,6 +821,9 @@ func decideColumnFormats(
 		return colFmts, colFmtDataAllText
 	} else {
 		colFmtData = make([]byte, 2+len(colFmts)*2)
+		if len(colFmts) > math.MaxUint16 {
+			panic("too many columns")
+		}
 		binary.BigEndian.PutUint16(colFmtData, uint16(len(colFmts)))
 		for i, v := range colFmts {
 			binary.BigEndian.PutUint16(colFmtData[2+i*2:], uint16(v))
diff --git a/copy.go b/copy.go
index a8f16b2b2..473c953d7 100644
--- a/copy.go
+++ b/copy.go
@@ -7,6 +7,7 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"math"
 	"sync"
 )
 
@@ -140,6 +141,9 @@ awaitCopyInResponse:
 }
 
 func (ci *copyin) flush(buf []byte) {
+	if len(buf)-1 > math.MaxUint32 {
+		panic("too many columns")
+	}
 	// set message length (without message identifier)
 	binary.BigEndian.PutUint32(buf[1:], uint32(len(buf)-1))