Have encapsulation.ReadData return an error when the buffer is short.

https://gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake/-/merge_requests/154#note_2919109

Still ignoring the io.ErrShortBuffer at the callers, which retains
current behavior.

4 files changed, 24 insertions, 16 deletions

diff --git a/client/lib/turbotunnel.go b/client/lib/turbotunnel.go
index d36eadd..642b0d4 100644
--- a/client/lib/turbotunnel.go
+++ b/client/lib/turbotunnel.go
@@ -38,10 +38,10 @@ func newEncapsulationPacketConn(
 // ReadFrom reads an encapsulated packet from the stream.
 func (c *encapsulationPacketConn) ReadFrom(p []byte) (int, net.Addr, error) {
 	n, err := encapsulation.ReadData(c.ReadWriteCloser, p)
-	if err != nil {
-		return n, c.remoteAddr, err
+	if err == io.ErrShortBuffer {
+		err = nil
 	}
-	return n, c.remoteAddr, nil
+	return n, c.remoteAddr, err
 }
 
 // WriteTo writes an encapsulated packet to the stream.
diff --git a/common/encapsulation/encapsulation.go b/common/encapsulation/encapsulation.go
index 1ccf6ab..a8ed539 100644
--- a/common/encapsulation/encapsulation.go
+++ b/common/encapsulation/encapsulation.go
@@ -53,11 +53,12 @@ var ErrTooLong = errors.New("length prefix is too long")
 
 // ReadData the next available data chunk, skipping over any padding chunks that
 // may come first, and copies the data into p. If p is shorter than the length
-// of the data chunk, only the first len(p) bytes are copied into p. The
-// returned error value is nil if and only if a data chunk was present and was
-// read in its entirety. The returned error is io.EOF only if r ended before the
-// first byte of a length prefix. If r ended in the middle of a length prefix or
-// data/padding, the returned error is io.ErrUnexpectedEOF.
+// of the data chunk, only the first len(p) bytes are copied into p, and the
+// error return is io.ErrShortBuffer. The returned error value is nil if and
+// only if a data chunk was present and was read in its entirety. The returned
+// error is io.EOF only if r ended before the first byte of a length prefix. If
+// r ended in the middle of a length prefix or data/padding, the returned error
+// is io.ErrUnexpectedEOF.
 func ReadData(r io.Reader, p []byte) (int, error) {
 	for {
 		var b [1]byte
@@ -89,9 +90,13 @@ func ReadData(r io.Reader, p []byte) (int, error) {
 			}
 			numData, err := io.ReadFull(r, p)
 			if err == nil && numData < n {
-				// Discard the rest of the data, if the caller's
-				// buffer was too short.
+				// If the caller's buffer was too short, discard
+				// the rest of the data and return
+				// io.ErrShortBuffer.
 				_, err = io.CopyN(ioutil.Discard, r, int64(n-numData))
+				if err == nil {
+					err = io.ErrShortBuffer
+				}
 			}
 			if err == io.EOF {
 				err = io.ErrUnexpectedEOF
diff --git a/common/encapsulation/encapsulation_test.go b/common/encapsulation/encapsulation_test.go
index 9fa8759..f499bfb 100644
--- a/common/encapsulation/encapsulation_test.go
+++ b/common/encapsulation/encapsulation_test.go
@@ -342,8 +342,8 @@ func TestReadDataTruncate(t *testing.T) {
 	var p [4]byte
 	// First ReadData should return truncated "1234".
 	n, err := ReadData(&enc, p[:])
-	if err != nil {
-		t.Fatalf("got error %v, expected %v", err, nil)
+	if err != io.ErrShortBuffer {
+		t.Fatalf("got error %v, expected %v", err, io.ErrShortBuffer)
 	}
 	if !bytes.Equal(p[:n], []byte("1234")) {
 		t.Fatalf("got <%x>, expected <%x>", p[:n], []byte("1234"))
@@ -351,8 +351,8 @@ func TestReadDataTruncate(t *testing.T) {
 	// Second ReadData should return truncated "abcd", not the rest of
 	// "12345678".
 	n, err = ReadData(&enc, p[:])
-	if err != nil {
-		t.Fatalf("got error %v, expected %v", err, nil)
+	if err != io.ErrShortBuffer {
+		t.Fatalf("got error %v, expected %v", err, io.ErrShortBuffer)
 	}
 	if !bytes.Equal(p[:n], []byte("abcd")) {
 		t.Fatalf("got <%x>, expected <%x>", p[:n], []byte("abcd"))
@@ -377,8 +377,8 @@ func TestReadDataTruncateFull(t *testing.T) {
 	}()
 	var p [8]byte
 	n, err := ReadData(pr, p[:])
-	if err != nil {
-		t.Fatalf("got error %v, expected %v", err, nil)
+	if err != io.ErrShortBuffer {
+		t.Fatalf("got error %v, expected %v", err, io.ErrShortBuffer)
 	}
 	// Should not stop after "hello".
 	if !bytes.Equal(p[:n], []byte("hellowor")) {
diff --git a/server/lib/http.go b/server/lib/http.go
index f6f2f37..4d33363 100644
--- a/server/lib/http.go
+++ b/server/lib/http.go
@@ -176,6 +176,9 @@ func (handler *httpHandler) turbotunnelMode(conn net.Conn, addr net.Addr) error
 		var p [2048]byte
 		for {
 			n, err := encapsulation.ReadData(conn, p[:])
+			if err == io.ErrShortBuffer {
+				err = nil
+			}
 			if err != nil {
 				return
 			}