pkg/ioutil: exact readcloser

NewExactReadCloser wraps a ReadCloser so it returns errors if exact number
of bytes are not read.

pkg/ioutil/readcloser.go

@@ -14,7 +14,10 @@
 
 package ioutil
 
-import "io"
+import (
+	"fmt"
+	"io"
+)
 
 // ReaderAndCloser implements io.ReadCloser interface by combining
 // reader and closer together.
@@ -22,3 +25,42 @@ type ReaderAndCloser struct {
 	io.Reader
 	io.Closer
 }
+
+var (
+	ErrShortRead = fmt.Errorf("ioutil: short read")
+	ErrExpectEOF = fmt.Errorf("ioutil: expect EOF")
+)
+
+// NewExactReadCloser returns a ReadCloser that returns errors if the underlying
+// reader does not read back exactly the requested number of bytes.
+func NewExactReadCloser(rc io.ReadCloser, totalBytes int64) io.ReadCloser {
+	return &exactReadCloser{rc: rc, totalBytes: totalBytes}
+}
+
+type exactReadCloser struct {
+	rc         io.ReadCloser
+	br         int64
+	totalBytes int64
+}
+
+func (e *exactReadCloser) Read(p []byte) (int, error) {
+	n, err := e.rc.Read(p)
+	e.br += int64(n)
+	if e.br > e.totalBytes {
+		return 0, ErrExpectEOF
+	}
+	if e.br < e.totalBytes && n == 0 {
+		return 0, ErrShortRead
+	}
+	return n, err
+}
+
+func (e *exactReadCloser) Close() error {
+	if err := e.rc.Close(); err != nil {
+		return err
+	}
+	if e.br < e.totalBytes {
+		return ErrShortRead
+	}
+	return nil
+}

pkg/ioutil/readcloser_test.go

@@ -0,0 +1,46 @@
+// Copyright 2016 The etcd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ioutil
+
+import (
+	"bytes"
+	"io"
+	"testing"
+)
+
+type readerNilCloser struct{ io.Reader }
+
+func (rc *readerNilCloser) Close() error { return nil }
+
+// TestExactReadCloserExpectEOF expects an eof when reading too much.
+func TestExactReadCloserExpectEOF(t *testing.T) {
+	buf := bytes.NewBuffer(make([]byte, 10))
+	rc := NewExactReadCloser(&readerNilCloser{buf}, 1)
+	if _, err := rc.Read(make([]byte, 10)); err != ErrExpectEOF {
+		t.Fatalf("expected %v, got %v", ErrExpectEOF, err)
+	}
+}
+
+// TestExactReadCloserShort expects an eof when reading too little
+func TestExactReadCloserShort(t *testing.T) {
+	buf := bytes.NewBuffer(make([]byte, 5))
+	rc := NewExactReadCloser(&readerNilCloser{buf}, 10)
+	if _, err := rc.Read(make([]byte, 10)); err != nil {
+		t.Fatalf("Read expected nil err, got %v", err)
+	}
+	if err := rc.Close(); err != ErrShortRead {
+		t.Fatalf("Close expected %v, got %v", ErrShortRead, err)
+	}
+}