database/sql: make RawBytes safely usable with contexts

sql.RawBytes was added the very first Go release, Go 1. Its docs
say:

> RawBytes is a byte slice that holds a reference to memory owned by
> the database itself. After a Scan into a RawBytes, the slice is only
> valid until the next call to Next, Scan, or Close.

That "only valid until the next call" bit was true at the time,
until contexts were added to database/sql in Go 1.8.

In the past ~dozen releases it's been unsafe to use QueryContext with
a context that might become Done to get an *sql.Rows that's scanning
into a RawBytes. The Scan can succeed, but then while the caller's
reading the memory, a database/sql-managed goroutine can see the
context becoming done and call Close on the database/sql/driver and
make the caller's view of the RawBytes memory no longer valid,
introducing races, crashes, or database corruption. See #60304
and #53970 for details.

This change does the minimal surgery on database/sql to make it safe
again: Rows.Scan was already acquiring a mutex to check whether the
rows had been closed, so this change make Rows.Scan notice whether
*RawBytes was used and, if so, doesn't release the mutex on exit
before returning. That mean it's still locked while the user code
operates on the RawBytes memory and the concurrent context-watching
goroutine to close the database still runs, but if it fires, it then
gets blocked on the mutex until the next call to a Rows method (Next,
NextResultSet, Err, Close).

Updates #60304
Updates #53970 (earlier one I'd missed)

Change-Id: Ie41c0c6f32c24887b2f53ec3686c2aab73a1bfff
Reviewed-on: https://go-review.googlesource.com/c/go/+/497675
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@google.com>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
Auto-Submit: Ian Lance Taylor <iant@google.com>
Reviewed-by: Russ Cox <rsc@golang.org>

1 files changed, 58 insertions, 0 deletions

diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 2b3d76f513..29a6709f23 100644
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -4385,6 +4385,64 @@ func TestRowsScanProperlyWrapsErrors(t *testing.T) {
 	}
 }
 
+// From go.dev/issue/60304
+func TestContextCancelDuringRawBytesScan(t *testing.T) {
+	db := newTestDB(t, "people")
+	defer closeDB(t, db)
+
+	if _, err := db.Exec("USE_RAWBYTES"); err != nil {
+		t.Fatal(err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	r, err := db.QueryContext(ctx, "SELECT|people|name|")
+	if err != nil {
+		t.Fatal(err)
+	}
+	numRows := 0
+	var sink byte
+	for r.Next() {
+		numRows++
+		var s RawBytes
+		err = r.Scan(&s)
+		if !r.closemuScanHold {
+			t.Errorf("expected closemu to be held")
+		}
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Logf("read %q", s)
+		if numRows == 2 {
+			cancel() // invalidate the context, which used to call close asynchronously
+		}
+		for _, b := range s { // some operation reading from the raw memory
+			sink += b
+		}
+	}
+	if r.closemuScanHold {
+		t.Errorf("closemu held; should not be")
+	}
+
+	// There are 3 rows. We canceled after reading 2 so we expect either
+	// 2 or 3 depending on how the awaitDone goroutine schedules.
+	switch numRows {
+	case 0, 1:
+		t.Errorf("got %d rows; want 2+", numRows)
+	case 2:
+		if err := r.Err(); err != context.Canceled {
+			t.Errorf("unexpected error: %v (%T)", err, err)
+		}
+	default:
+		// Made it to the end. This is rare, but fine. Permit it.
+	}
+
+	if err := r.Close(); err != nil {
+		t.Fatal(err)
+	}
+}
+
 // badConn implements a bad driver.Conn, for TestBadDriver.
 // The Exec method panics.
 type badConn struct{}