From dd268d6f71112a4c374c254440b2b3f8e8636e1c Mon Sep 17 00:00:00 2001
From: Rafi Shamim <rafi@cockroachlabs.com>
Date: Thu, 18 May 2023 18:03:34 -0400
Subject: [PATCH] pgwire: adds bounds check for FormatCodes

Now we check the bounds when looking for the format code to use.

Getting an out of bounds index is not expected, so we use an assertion
error. This is less disruptive than a panic that crashes the node.

Release note: None
---
 pkg/sql/pgwire/conn.go                               | 6 ++++++
 pkg/sql/pgwire/pgwirebase/BUILD.bazel                | 1 +
 pkg/sql/pgwire/pgwirebase/encoding.go                | 6 ++++++
 pkg/testutils/lint/passes/redactcheck/redactcheck.go | 3 +++
 4 files changed, 16 insertions(+)

diff --git a/pkg/sql/pgwire/conn.go b/pkg/sql/pgwire/conn.go
index 4b2bb0279b7c..47cc575b8682 100644
--- a/pkg/sql/pgwire/conn.go
+++ b/pkg/sql/pgwire/conn.go
@@ -1370,6 +1370,9 @@ func (c *conn) bufferRow(ctx context.Context, row tree.Datums, r *commandResult)
 	for i, col := range row {
 		fmtCode := pgwirebase.FormatText
 		if r.formatCodes != nil {
+			if i >= len(r.formatCodes) {
+				return errors.AssertionFailedf("could not find format code for column %d in %v", i, r.formatCodes)
+			}
 			fmtCode = r.formatCodes[i]
 		}
 		switch fmtCode {
@@ -1412,6 +1415,9 @@ func (c *conn) bufferBatch(ctx context.Context, batch coldata.Batch, r *commandR
 			for vecIdx := 0; vecIdx < len(c.vecsScratch.Vecs); vecIdx++ {
 				fmtCode := pgwirebase.FormatText
 				if r.formatCodes != nil {
+					if vecIdx >= len(r.formatCodes) {
+						return errors.AssertionFailedf("could not find format code for column %d in %v", vecIdx, r.formatCodes)
+					}
 					fmtCode = r.formatCodes[vecIdx]
 				}
 				switch fmtCode {
diff --git a/pkg/sql/pgwire/pgwirebase/BUILD.bazel b/pkg/sql/pgwire/pgwirebase/BUILD.bazel
index 489e4f28f33c..9ad52b9b1ad2 100644
--- a/pkg/sql/pgwire/pgwirebase/BUILD.bazel
+++ b/pkg/sql/pgwire/pgwirebase/BUILD.bazel
@@ -39,6 +39,7 @@ go_library(
         "//pkg/util/tsearch",
         "//pkg/util/uint128",
         "@com_github_cockroachdb_errors//:errors",
+        "@com_github_cockroachdb_redact//:redact",
         "@com_github_dustin_go_humanize//:go-humanize",
         "@com_github_lib_pq//oid",
     ],
diff --git a/pkg/sql/pgwire/pgwirebase/encoding.go b/pkg/sql/pgwire/pgwirebase/encoding.go
index 52dae906d0f9..21aa603309e9 100644
--- a/pkg/sql/pgwire/pgwirebase/encoding.go
+++ b/pkg/sql/pgwire/pgwirebase/encoding.go
@@ -42,6 +42,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/tsearch"
 	"github.com/cockroachdb/cockroach/pkg/util/uint128"
 	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/redact"
 	"github.com/dustin/go-humanize"
 	"github.com/lib/pq/oid"
 )
@@ -73,6 +74,11 @@ var ReadBufferMaxMessageSizeClusterSetting = settings.RegisterByteSizeSetting(
 //go:generate stringer -type=FormatCode
 type FormatCode uint16
 
+var _ redact.SafeValue = FormatCode(0)
+
+// SafeValue implements the redact.SafeValue interface.
+func (i FormatCode) SafeValue() {}
+
 const (
 	// FormatText is the default, text format.
 	FormatText FormatCode = 0
diff --git a/pkg/testutils/lint/passes/redactcheck/redactcheck.go b/pkg/testutils/lint/passes/redactcheck/redactcheck.go
index 97641fead33b..99758a845119 100644
--- a/pkg/testutils/lint/passes/redactcheck/redactcheck.go
+++ b/pkg/testutils/lint/passes/redactcheck/redactcheck.go
@@ -130,6 +130,9 @@ func runAnalyzer(pass *analysis.Pass) (interface{}, error) {
 						"IndexDescriptorVersion":       {},
 						"MutationID":                   {},
 					},
+					"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgwirebase": {
+						"FormatCode": {},
+					},
 					"github.com/cockroachdb/cockroach/pkg/sql/sem/catconstants": {
 						"ConstraintType": {},
 					},