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
 }