database/sql: avoid closing Rows while scan is in progress

A database/sql/driver.Rows can return database-owned data
from Rows.Next. The driver.Rows documentation doesn't explicitly
document the lifetime guarantees for this data, but a reasonable
expectation is that the caller of Next should only access it
until the next call to Rows.Close or Rows.Next.

Avoid violating that constraint when a query is cancelled while
a call to database/sql.Rows.Scan (note the difference between
the two different Rows types!) is in progress. We previously
took care to avoid closing a driver.Rows while the user has
access to driver-owned memory via a RawData, but we could still
close a driver.Rows while a Scan call was in the process of
reading previously-returned driver-owned data.

Update the fake DB used in database/sql tests to invalidate
returned data to help catch other places we might be
incorrectly retaining it.

Fixes #74831.

Change-Id: Ice45b5fad51b679c38e3e1d21ef39156b56d6037
Reviewed-on: https://go-internal-review.googlesource.com/c/go/+/2540
Reviewed-by: Roland Shoemaker <bracewell@google.com>
Reviewed-by: Neal Patel <nealpatel@google.com>
Reviewed-on: https://go-review.googlesource.com/c/go/+/693735
Auto-Submit: Dmitri Shuralyov <dmitshur@google.com>
Reviewed-by: Roland Shoemaker <roland@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Dmitri Shuralyov <dmitshur@google.com>

1 files changed, 22 insertions, 25 deletions

diff --git a/src/database/sql/fakedb_test.go b/src/database/sql/fakedb_test.go
index 3dfcd447b5..003e6c6298 100644
--- a/src/database/sql/fakedb_test.go
+++ b/src/database/sql/fakedb_test.go
@@ -5,6 +5,7 @@
 package sql
 
 import (
+	"bytes"
 	"context"
 	"database/sql/driver"
 	"errors"
@@ -15,7 +16,6 @@ import (
 	"strconv"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"testing"
 	"time"
 )
@@ -91,8 +91,6 @@ func (cc *fakeDriverCtx) OpenConnector(name string) (driver.Connector, error) {
 type fakeDB struct {
 	name string
 
-	useRawBytes atomic.Bool
-
 	mu       sync.Mutex
 	tables   map[string]*table
 	badConn  bool
@@ -684,8 +682,6 @@ func (c *fakeConn) PrepareContext(ctx context.Context, query string) (driver.Stm
 		switch cmd {
 		case "WIPE":
 			// Nothing
-		case "USE_RAWBYTES":
-			c.db.useRawBytes.Store(true)
 		case "SELECT":
 			stmt, err = c.prepareSelect(stmt, parts)
 		case "CREATE":
@@ -789,9 +785,6 @@ func (s *fakeStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (d
 	case "WIPE":
 		db.wipe()
 		return driver.ResultNoRows, nil
-	case "USE_RAWBYTES":
-		s.c.db.useRawBytes.Store(true)
-		return driver.ResultNoRows, nil
 	case "CREATE":
 		if err := db.createTable(s.table, s.colName, s.colType); err != nil {
 			return nil, err
@@ -1076,10 +1069,9 @@ type rowsCursor struct {
 	errPos int
 	err    error
 
-	// a clone of slices to give out to clients, indexed by the
-	// original slice's first byte address.  we clone them
-	// just so we're able to corrupt them on close.
-	bytesClone map[*byte][]byte
+	// Data returned to clients.
+	// We clone and stash it here so it can be invalidated by Close and Next.
+	driverOwnedMemory [][]byte
 
 	// Every operation writes to line to enable the race detector
 	// check for data races.
@@ -1096,9 +1088,19 @@ func (rc *rowsCursor) touchMem() {
 	rc.line++
 }
 
+func (rc *rowsCursor) invalidateDriverOwnedMemory() {
+	for _, buf := range rc.driverOwnedMemory {
+		for i := range buf {
+			buf[i] = 'x'
+		}
+	}
+	rc.driverOwnedMemory = nil
+}
+
 func (rc *rowsCursor) Close() error {
 	rc.touchMem()
 	rc.parentMem.touchMem()
+	rc.invalidateDriverOwnedMemory()
 	rc.closed = true
 	return rc.closeErr
 }
@@ -1129,6 +1131,8 @@ func (rc *rowsCursor) Next(dest []driver.Value) error {
 	if rc.posRow >= len(rc.rows[rc.posSet]) {
 		return io.EOF // per interface spec
 	}
+	// Corrupt any previously returned bytes.
+	rc.invalidateDriverOwnedMemory()
 	for i, v := range rc.rows[rc.posSet][rc.posRow].cols {
 		// TODO(bradfitz): convert to subset types? naah, I
 		// think the subset types should only be input to
@@ -1136,20 +1140,13 @@ func (rc *rowsCursor) Next(dest []driver.Value) error {
 		// a wider range of types coming out of drivers. all
 		// for ease of drivers, and to prevent drivers from
 		// messing up conversions or doing them differently.
-		dest[i] = v
-
-		if bs, ok := v.([]byte); ok && !rc.db.useRawBytes.Load() {
-			if rc.bytesClone == nil {
-				rc.bytesClone = make(map[*byte][]byte)
-			}
-			clone, ok := rc.bytesClone[&bs[0]]
-			if !ok {
-				clone = make([]byte, len(bs))
-				copy(clone, bs)
-				rc.bytesClone[&bs[0]] = clone
-			}
-			dest[i] = clone
+		if bs, ok := v.([]byte); ok {
+			// Clone []bytes and stash for later invalidation.
+			bs = bytes.Clone(bs)
+			rc.driverOwnedMemory = append(rc.driverOwnedMemory, bs)
+			v = bs
 		}
+		dest[i] = v
 	}
 	return nil
 }