From c23579f031ecd09bf37c644723b33736dffa8b92 Mon Sep 17 00:00:00 2001 From: Damien Neil Date: Tue, 23 Jan 2024 15:59:47 -0800 Subject: [PATCH] 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 Reviewed-by: Roland Shoemaker LUCI-TryBot-Result: Go LUCI --- src/database/sql/convert.go | 8 +++--- src/database/sql/convert_test.go | 12 ++++++-- src/database/sql/sql.go | 34 +++++++++++++++++++++++ src/database/sql/sql_test.go | 47 ++++++++++++++++++++++++++++++++ 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{} -- 2.48.1