database/sql: avoid clobbering driver-owned memory in RawBytes

Depending on the query, a RawBytes can contain memory owned by the
driver or by database/sql:

If the driver provides the column as a []byte,
RawBytes aliases that []byte.

If the driver provides the column as any other type,
RawBytes contains memory allocated by database/sql.
Prior to this CL, Rows.Scan will reuse existing capacity in a
RawBytes to permit a single allocation to be reused across rows.

When a RawBytes is reused across queries, this can result
in database/sql writing to driver-owned memory.

Add a buffer to Rows to store RawBytes data, and reuse this
buffer across calls to Rows.Scan.

Fixes #65201

Change-Id: Iac640174c7afa97eeb39496f47dec202501b2483
Reviewed-on: https://go-review.googlesource.com/c/go/+/557917
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Reviewed-by: Roland Shoemaker <roland@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

4 files changed, 94 insertions, 7 deletions

diff --git a/src/database/sql/convert.go b/src/database/sql/convert.go
index dac3f246ae..8f71d5b867 100644
--- a/src/database/sql/convert.go
+++ b/src/database/sql/convert.go
@@ -237,7 +237,7 @@ func convertAssignRows(dest, src any, rows *Rows) error {
 			if d == nil {
 				return errNilPtr
 			}
-			*d = append((*d)[:0], s...)
+			*d = rows.setrawbuf(append(rows.rawbuf(), s...))
 			return nil
 		}
 	case []byte:
@@ -285,7 +285,7 @@ func convertAssignRows(dest, src any, rows *Rows) error {
 			if d == nil {
 				return errNilPtr
 			}
-			*d = s.AppendFormat((*d)[:0], time.RFC3339Nano)
+			*d = rows.setrawbuf(s.AppendFormat(rows.rawbuf(), time.RFC3339Nano))
 			return nil
 		}
 	case decimalDecompose:
@@ -366,8 +366,8 @@ func convertAssignRows(dest, src any, rows *Rows) error {
 		}
 	case *RawBytes:
 		sv = reflect.ValueOf(src)
-		if b, ok := asBytes([]byte(*d)[:0], sv); ok {
-			*d = RawBytes(b)
+		if b, ok := asBytes(rows.rawbuf(), sv); ok {
+			*d = rows.setrawbuf(b)
 			return nil
 		}
 	case *bool:
diff --git a/src/database/sql/convert_test.go b/src/database/sql/convert_test.go
index 6d09fa1eae..f94db8e5f8 100644
--- a/src/database/sql/convert_test.go
+++ b/src/database/sql/convert_test.go
@@ -354,9 +354,10 @@ func TestRawBytesAllocs(t *testing.T) {
 		{"time", time.Unix(2, 5).UTC(), "1970-01-01T00:00:02.000000005Z"},
 	}
 
-	buf := make(RawBytes, 10)
+	var buf RawBytes
+	rows := &Rows{}
 	test := func(name string, in any, want string) {
-		if err := convertAssign(&buf, in); err != nil {
+		if err := convertAssignRows(&buf, in, rows); err != nil {
 			t.Fatalf("%s: convertAssign = %v", name, err)
 		}
 		match := len(buf) == len(want)
@@ -375,6 +376,7 @@ func TestRawBytesAllocs(t *testing.T) {
 
 	n := testing.AllocsPerRun(100, func() {
 		for _, tt := range tests {
+			rows.raw = rows.raw[:0]
 			test(tt.name, tt.in, tt.want)
 		}
 	})
@@ -383,7 +385,11 @@ func TestRawBytesAllocs(t *testing.T) {
 	// and gc. With 32-bit words there are more convT2E allocs, and
 	// with gccgo, only pointers currently go in interface data.
 	// So only care on amd64 gc for now.
-	measureAllocs := runtime.GOARCH == "amd64" && runtime.Compiler == "gc"
+	measureAllocs := false
+	switch runtime.GOARCH {
+	case "amd64", "arm64":
+		measureAllocs = runtime.Compiler == "gc"
+	}
 
 	if n > 0.5 && measureAllocs {
 		t.Fatalf("allocs = %v; want 0", n)
diff --git a/src/database/sql/sql.go b/src/database/sql/sql.go
index 5b4a3f5409..fdbe4b2172 100644
--- a/src/database/sql/sql.go
+++ b/src/database/sql/sql.go
@@ -2930,6 +2930,13 @@ type Rows struct {
 	// not to be called concurrently.
 	lastcols []driver.Value
 
+	// raw is a buffer for RawBytes that persists between Scan calls.
+	// This is used when the driver returns a mismatched type that requires
+	// a cloning allocation. For example, if the driver returns a *string and
+	// the user is scanning into a *RawBytes, we need to copy the string.
+	// The raw buffer here lets us reuse the memory for that copy across Scan calls.
+	raw []byte
+
 	// closemuScanHold is whether the previous call to Scan kept closemu RLock'ed
 	// without unlocking it. It does that when the user passes a *RawBytes scan
 	// target. In that case, we need to prevent awaitDone from closing the Rows
@@ -3124,6 +3131,32 @@ func (rs *Rows) Err() error {
 	return rs.lasterrOrErrLocked(nil)
 }
 
+// rawbuf returns the buffer to append RawBytes values to.
+// This buffer is reused across calls to Rows.Scan.
+//
+// Usage:
+//
+//	rawBytes = rows.setrawbuf(append(rows.rawbuf(), value...))
+func (rs *Rows) rawbuf() []byte {
+	if rs == nil {
+		// convertAssignRows can take a nil *Rows; for simplicity handle it here
+		return nil
+	}
+	return rs.raw
+}
+
+// setrawbuf updates the RawBytes buffer with the result of appending a new value to it.
+// It returns the new value.
+func (rs *Rows) setrawbuf(b []byte) RawBytes {
+	if rs == nil {
+		// convertAssignRows can take a nil *Rows; for simplicity handle it here
+		return RawBytes(b)
+	}
+	off := len(rs.raw)
+	rs.raw = b
+	return RawBytes(rs.raw[off:])
+}
+
 var errRowsClosed = errors.New("sql: Rows are closed")
 var errNoRows = errors.New("sql: no Rows available")
 
@@ -3331,6 +3364,7 @@ func (rs *Rows) Scan(dest ...any) error {
 
 	if scanArgsContainRawBytes(dest) {
 		rs.closemuScanHold = true
+		rs.raw = rs.raw[:0]
 	} else {
 		rs.closemu.RUnlock()
 	}
diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 25ca5ff0ad..7dfc6434e0 100644
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -4566,6 +4566,53 @@ func TestNilErrorAfterClose(t *testing.T) {
 	}
 }
 
+// Issue #65201.
+//
+// If a RawBytes is reused across multiple queries,
+// subsequent queries shouldn't overwrite driver-owned memory from previous queries.
+func TestRawBytesReuse(t *testing.T) {
+	db := newTestDB(t, "people")
+	defer closeDB(t, db)
+
+	if _, err := db.Exec("USE_RAWBYTES"); err != nil {
+		t.Fatal(err)
+	}
+
+	var raw RawBytes
+
+	// The RawBytes in this query aliases driver-owned memory.
+	rows, err := db.Query("SELECT|people|name|")
+	if err != nil {
+		t.Fatal(err)
+	}
+	rows.Next()
+	rows.Scan(&raw) // now raw is pointing to driver-owned memory
+	name1 := string(raw)
+	rows.Close()
+
+	// The RawBytes in this query does not alias driver-owned memory.
+	rows, err = db.Query("SELECT|people|age|")
+	if err != nil {
+		t.Fatal(err)
+	}
+	rows.Next()
+	rows.Scan(&raw) // this must not write to the driver-owned memory in raw
+	rows.Close()
+
+	// Repeat the first query. Nothing should have changed.
+	rows, err = db.Query("SELECT|people|name|")
+	if err != nil {
+		t.Fatal(err)
+	}
+	rows.Next()
+	rows.Scan(&raw) // raw points to driver-owned memory again
+	name2 := string(raw)
+	rows.Close()
+	if name1 != name2 {
+		t.Fatalf("Scan read name %q, want %q", name2, name1)
+	}
+}
+
 // badConn implements a bad driver.Conn, for TestBadDriver.
 // The Exec method panics.
 type badConn struct{}