diff --git a/pkg/ioutil/readcloser.go b/pkg/ioutil/readcloser.go index 4aeb0c7b0..d3efcfe3d 100644 --- a/pkg/ioutil/readcloser.go +++ b/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 +} diff --git a/pkg/ioutil/readcloser_test.go b/pkg/ioutil/readcloser_test.go new file mode 100644 index 000000000..6d13bdcec --- /dev/null +++ b/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) + } +}