Merge pull request #177 from dackroyd/invalidate-sql-rawbytes-memory

Invalidate memory scanned into sql.RawBytes

Author: l3pp4rd

rows.go

@@ -1,13 +1,16 @@
 package sqlmock
 
 import (
+	"bytes"
 	"database/sql/driver"
 	"encoding/csv"
 	"fmt"
 	"io"
 	"strings"
 )
 
+const invalidate = "☠☠☠ MEMORY OVERWRITTEN ☠☠☠ "
+
 // CSVColumnParser is a function which converts trimmed csv
 // column string to a []byte representation. Currently
 // transforms NULL to nil
@@ -23,13 +26,15 @@ type rowSets struct {
 	sets []*Rows
 	pos  int
 	ex   *ExpectedQuery
+	raw  [][]byte
 }
 
 func (rs *rowSets) Columns() []string {
 	return rs.sets[rs.pos].cols
 }
 
 func (rs *rowSets) Close() error {
+	rs.invalidateRaw()
 	rs.ex.rowsWereClosed = true
 	return rs.sets[rs.pos].closeErr
 }
@@ -38,11 +43,17 @@ func (rs *rowSets) Close() error {
 func (rs *rowSets) Next(dest []driver.Value) error {
 	r := rs.sets[rs.pos]
 	r.pos++
+	rs.invalidateRaw()
 	if r.pos > len(r.rows) {
 		return io.EOF // per interface spec
 	}
 
 	for i, col := range r.rows[r.pos-1] {
+		if b, ok := rawBytes(col); ok {
+			rs.raw = append(rs.raw, b)
+			dest[i] = b
+			continue
+		}
 		dest[i] = col
 	}
 
@@ -80,6 +91,30 @@ func (rs *rowSets) empty() bool {
 	return true
 }
 
+func rawBytes(col driver.Value) (_ []byte, ok bool) {
+	val, ok := col.([]byte)
+	if !ok || len(val) == 0 {
+		return nil, false
+	}
+	// Copy the bytes from the mocked row into a shared raw buffer, which we'll replace the content of later
+	// This allows scanning into sql.RawBytes to correctly become invalid on subsequent calls to Next(), Scan() or Close()
+	b := make([]byte, len(val))
+	copy(b, val)
+	return b, true
+}
+
+// Bytes that could have been scanned as sql.RawBytes are only valid until the next call to Next, Scan or Close.
+// If those occur, we must replace their content to simulate the shared memory to expose misuse of sql.RawBytes
+func (rs *rowSets) invalidateRaw() {
+	// Replace the content of slices previously returned
+	b := []byte(invalidate)
+	for _, r := range rs.raw {
+		copy(r, bytes.Repeat(b, len(r)/len(b)+1))
+	}
+	// Start with new slices for the next scan
+	rs.raw = nil
+}
+
 // Rows is a mocked collection of rows to
 // return for Query result
 type Rows struct {

rows_go13_test.go

@@ -0,0 +1,31 @@
+// +build go1.3
+
+package sqlmock
+
+import (
+	"database/sql"
+	"testing"
+)
+
+func TestQueryRowBytesNotInvalidatedByNext_stringIntoRawBytes(t *testing.T) {
+	t.Parallel()
+	rows := NewRows([]string{"raw"}).
+		AddRow(`one binary value with some text!`).
+		AddRow(`two binary value with even more text than the first one`)
+	scan := func(rs *sql.Rows) ([]byte, error) {
+		var raw sql.RawBytes
+		return raw, rs.Scan(&raw)
+	}
+	want := [][]byte{[]byte(`one binary value with some text!`), []byte(`two binary value with even more text than the first one`)}
+	queryRowBytesNotInvalidatedByNext(t, rows, scan, want)
+}
+
+func TestQueryRowBytesNotInvalidatedByClose_stringIntoRawBytes(t *testing.T) {
+	t.Parallel()
+	rows := NewRows([]string{"raw"}).AddRow(`one binary value with some text!`)
+	scan := func(rs *sql.Rows) ([]byte, error) {
+		var raw sql.RawBytes
+		return raw, rs.Scan(&raw)
+	}
+	queryRowBytesNotInvalidatedByClose(t, rows, scan, []byte(`one binary value with some text!`))
+}

rows_go18_test.go

@@ -3,6 +3,8 @@
 package sqlmock
 
 import (
+	"database/sql"
+	"encoding/json"
 	"fmt"
 	"testing"
 )
@@ -90,3 +92,114 @@ func TestQueryMultiRows(t *testing.T) {
 		t.Errorf("there were unfulfilled expectations: %s", err)
 	}
 }
+
+func TestQueryRowBytesInvalidatedByNext_jsonRawMessageIntoRawBytes(t *testing.T) {
+	t.Parallel()
+	replace := []byte(invalid)
+	rows := NewRows([]string{"raw"}).
+		AddRow(json.RawMessage(`{"thing": "one", "thing2": "two"}`)).
+		AddRow(json.RawMessage(`{"that": "foo", "this": "bar"}`))
+	scan := func(rs *sql.Rows) ([]byte, error) {
+		var raw sql.RawBytes
+		return raw, rs.Scan(&raw)
+	}
+	want := []struct {
+		Initial  []byte
+		Replaced []byte
+	}{
+		{Initial: []byte(`{"thing": "one", "thing2": "two"}`), Replaced: replace[:len(replace)-6]},
+		{Initial: []byte(`{"that": "foo", "this": "bar"}`), Replaced: replace[:len(replace)-9]},
+	}
+	queryRowBytesInvalidatedByNext(t, rows, scan, want)
+}
+
+func TestQueryRowBytesNotInvalidatedByNext_jsonRawMessageIntoBytes(t *testing.T) {
+	t.Parallel()
+	rows := NewRows([]string{"raw"}).
+		AddRow(json.RawMessage(`{"thing": "one", "thing2": "two"}`)).
+		AddRow(json.RawMessage(`{"that": "foo", "this": "bar"}`))
+	scan := func(rs *sql.Rows) ([]byte, error) {
+		var b []byte
+		return b, rs.Scan(&b)
+	}
+	want := [][]byte{[]byte(`{"thing": "one", "thing2": "two"}`), []byte(`{"that": "foo", "this": "bar"}`)}
+	queryRowBytesNotInvalidatedByNext(t, rows, scan, want)
+}
+
+func TestQueryRowBytesNotInvalidatedByNext_bytesIntoCustomBytes(t *testing.T) {
+	t.Parallel()
+	rows := NewRows([]string{"raw"}).
+		AddRow([]byte(`one binary value with some text!`)).
+		AddRow([]byte(`two binary value with even more text than the first one`))
+	scan := func(rs *sql.Rows) ([]byte, error) {
+		type customBytes []byte
+		var b customBytes
+		return b, rs.Scan(&b)
+	}
+	want := [][]byte{[]byte(`one binary value with some text!`), []byte(`two binary value with even more text than the first one`)}
+	queryRowBytesNotInvalidatedByNext(t, rows, scan, want)
+}
+
+func TestQueryRowBytesNotInvalidatedByNext_jsonRawMessageIntoCustomBytes(t *testing.T) {
+	t.Parallel()
+	rows := NewRows([]string{"raw"}).
+		AddRow(json.RawMessage(`{"thing": "one", "thing2": "two"}`)).
+		AddRow(json.RawMessage(`{"that": "foo", "this": "bar"}`))
+	scan := func(rs *sql.Rows) ([]byte, error) {
+		type customBytes []byte
+		var b customBytes
+		return b, rs.Scan(&b)
+	}
+	want := [][]byte{[]byte(`{"thing": "one", "thing2": "two"}`), []byte(`{"that": "foo", "this": "bar"}`)}
+	queryRowBytesNotInvalidatedByNext(t, rows, scan, want)
+}
+
+func TestQueryRowBytesNotInvalidatedByClose_bytesIntoCustomBytes(t *testing.T) {
+	t.Parallel()
+	rows := NewRows([]string{"raw"}).AddRow([]byte(`one binary value with some text!`))
+	scan := func(rs *sql.Rows) ([]byte, error) {
+		type customBytes []byte
+		var b customBytes
+		return b, rs.Scan(&b)
+	}
+	queryRowBytesNotInvalidatedByClose(t, rows, scan, []byte(`one binary value with some text!`))
+}
+
+func TestQueryRowBytesInvalidatedByClose_jsonRawMessageIntoRawBytes(t *testing.T) {
+	t.Parallel()
+	replace := []byte(invalid)
+	rows := NewRows([]string{"raw"}).AddRow(json.RawMessage(`{"thing": "one", "thing2": "two"}`))
+	scan := func(rs *sql.Rows) ([]byte, error) {
+		var raw sql.RawBytes
+		return raw, rs.Scan(&raw)
+	}
+	want := struct {
+		Initial  []byte
+		Replaced []byte
+	}{
+		Initial:  []byte(`{"thing": "one", "thing2": "two"}`),
+		Replaced: replace[:len(replace)-6],
+	}
+	queryRowBytesInvalidatedByClose(t, rows, scan, want)
+}
+
+func TestQueryRowBytesNotInvalidatedByClose_jsonRawMessageIntoBytes(t *testing.T) {
+	t.Parallel()
+	rows := NewRows([]string{"raw"}).AddRow(json.RawMessage(`{"thing": "one", "thing2": "two"}`))
+	scan := func(rs *sql.Rows) ([]byte, error) {
+		var b []byte
+		return b, rs.Scan(&b)
+	}
+	queryRowBytesNotInvalidatedByClose(t, rows, scan, []byte(`{"thing": "one", "thing2": "two"}`))
+}
+
+func TestQueryRowBytesNotInvalidatedByClose_jsonRawMessageIntoCustomBytes(t *testing.T) {
+	t.Parallel()
+	rows := NewRows([]string{"raw"}).AddRow(json.RawMessage(`{"thing": "one", "thing2": "two"}`))
+	scan := func(rs *sql.Rows) ([]byte, error) {
+		type customBytes []byte
+		var b customBytes
+		return b, rs.Scan(&b)
+	}
+	queryRowBytesNotInvalidatedByClose(t, rows, scan, []byte(`{"thing": "one", "thing2": "two"}`))
+}