snap: return errors if Message's snapshot is not entirely read

release-3.0
Anthony Romano 2016-06-04 20:37:03 -07:00
parent 3aadb25c31
commit 93f114c76c
1 changed files with 5 additions and 2 deletions

View File

@ -17,6 +17,7 @@ package snap
import (
"io"
"github.com/coreos/etcd/pkg/ioutil"
"github.com/coreos/etcd/raft/raftpb"
)
@ -38,7 +39,7 @@ type Message struct {
func NewMessage(rs raftpb.Message, rc io.ReadCloser, rcSize int64) *Message {
return &Message{
Message: rs,
ReadCloser: rc,
ReadCloser: ioutil.NewExactReadCloser(rc, rcSize),
TotalSize: int64(rs.Size()) + rcSize,
closeC: make(chan bool, 1),
}
@ -52,7 +53,9 @@ func (m Message) CloseNotify() <-chan bool {
}
func (m Message) CloseWithError(err error) {
m.ReadCloser.Close()
if cerr := m.ReadCloser.Close(); cerr != nil {
err = cerr
}
if err == nil {
m.closeC <- true
} else {