bufio: reject UnreadByte or UnreadRune after a Discard or WriteTo

Discard is not really a read operation, and in theory it could
Seek the underlying Reader without actually reading anything,
so an UnreadByte following a Discard is disallowed.

Similarly, although WriteTo usually does end up calling Read on the
underlying buffer, if the underlying Reader implements io.WriterTo it
may instead terminate in a call to WriteTo, without ever buffering or
even seeing the last byte written. (It is conceptually read-like, but
not strictly “a read operation”.)

Fixes #48446

Change-Id: Ide6f2b157332b423486810399f66140c914144e5
Reviewed-on: https://go-review.googlesource.com/c/go/+/351810
Trust: Bryan C. Mills <bcmills@google.com>
Trust: Joe Tsai <joetsai@digital-static.net>
Reviewed-by: Joe Tsai <joetsai@digital-static.net>

2 files changed, 43 insertions, 2 deletions

diff --git a/src/bufio/bufio.go b/src/bufio/bufio.go
index 506b84f6ba..a58df25494 100644
--- a/src/bufio/bufio.go
+++ b/src/bufio/bufio.go
@@ -173,6 +173,10 @@ func (b *Reader) Discard(n int) (discarded int, err error) {
 	if n == 0 {
 		return
 	}
+
+	b.lastByte = -1
+	b.lastRuneSize = -1
+
 	remain := n
 	for {
 		skip := b.Buffered()
@@ -266,8 +270,8 @@ func (b *Reader) ReadByte() (byte, error) {
 // UnreadByte unreads the last byte. Only the most recently read byte can be unread.
 //
 // UnreadByte returns an error if the most recent method called on the
-// Reader was not a read operation. Notably, Peek is not considered a
-// read operation.
+// Reader was not a read operation. Notably, Peek, Discard, and WriteTo are not
+// considered read operations.
 func (b *Reader) UnreadByte() error {
 	if b.lastByte < 0 || b.r == 0 && b.w > 0 {
 		return ErrInvalidUnreadByte
@@ -502,6 +506,9 @@ func (b *Reader) ReadString(delim byte) (string, error) {
 // If the underlying reader supports the WriteTo method,
 // this calls the underlying WriteTo without buffering.
 func (b *Reader) WriteTo(w io.Writer) (n int64, err error) {
+	b.lastByte = -1
+	b.lastRuneSize = -1
+
 	n, err = b.writeBuf(w)
 	if err != nil {
 		return
diff --git a/src/bufio/bufio_test.go b/src/bufio/bufio_test.go
index 04a810c206..8e8a8a1778 100644
--- a/src/bufio/bufio_test.go
+++ b/src/bufio/bufio_test.go
@@ -304,6 +304,40 @@ func TestNoUnreadByteAfterPeek(t *testing.T) {
 	}
 }
 
+func TestNoUnreadRuneAfterDiscard(t *testing.T) {
+	br := NewReader(strings.NewReader("example"))
+	br.ReadRune()
+	br.Discard(1)
+	if err := br.UnreadRune(); err == nil {
+		t.Error("UnreadRune didn't fail after Discard")
+	}
+}
+
+func TestNoUnreadByteAfterDiscard(t *testing.T) {
+	br := NewReader(strings.NewReader("example"))
+	br.ReadByte()
+	br.Discard(1)
+	if err := br.UnreadByte(); err == nil {
+		t.Error("UnreadByte didn't fail after Discard")
+	}
+}
+
+func TestNoUnreadRuneAfterWriteTo(t *testing.T) {
+	br := NewReader(strings.NewReader("example"))
+	br.WriteTo(io.Discard)
+	if err := br.UnreadRune(); err == nil {
+		t.Error("UnreadRune didn't fail after WriteTo")
+	}
+}
+
+func TestNoUnreadByteAfterWriteTo(t *testing.T) {
+	br := NewReader(strings.NewReader("example"))
+	br.WriteTo(io.Discard)
+	if err := br.UnreadByte(); err == nil {
+		t.Error("UnreadByte didn't fail after WriteTo")
+	}
+}
+
 func TestUnreadByte(t *testing.T) {
 	segments := []string{"Hello, ", "world"}
 	r := NewReader(&StringReader{data: segments})