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>

diff --git a/src/bufio/bufio.go b/src/bufio/bufio.go
index 506b84f6baf94f9d4bd82c81d84cc492ae64d3c6..a58df254941166ae72c71888c15d3d14eebf84d9 100644 (file)
--- 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 04a810c206500d846174c291e2b25fc4af59848f..8e8a8a1778a690ade5bed0a4d8bef18ead5324e8 100644 (file)
--- 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})