packer/rpc: make sure we read all the data

This commit is contained in:
Mitchell Hashimoto 2014-01-01 21:53:36 -08:00
parent 3a1908bbb3
commit d5bf5d7f39
1 changed files with 18 additions and 5 deletions

View File

@ -284,10 +284,13 @@ func (m *MuxConn) loop() {
// TODO(mitchellh): probably would be better to re-use a buffer... // TODO(mitchellh): probably would be better to re-use a buffer...
data := make([]byte, length) data := make([]byte, length)
if length > 0 { n := 0
if _, err := m.rwc.Read(data); err != nil { for n < int(length) {
if n2, err := m.rwc.Read(data); err != nil {
log.Printf("[ERR] Error reading data: %s", err) log.Printf("[ERR] Error reading data: %s", err)
return return
} else {
n += n2
} }
} }
@ -434,10 +437,20 @@ func (m *MuxConn) write(from muxPacketFrom, id uint32, dataType muxPacketType, p
if err := binary.Write(m.rwc, binary.BigEndian, int32(len(p))); err != nil { if err := binary.Write(m.rwc, binary.BigEndian, int32(len(p))); err != nil {
return 0, err return 0, err
} }
if len(p) == 0 {
return 0, nil // Write all the bytes. If we don't write all the bytes, report an error
var err error = nil
n := 0
for n < len(p) {
var n2 int
n2, err = m.rwc.Write(p)
n += n2
if err != nil {
break
}
} }
return m.rwc.Write(p)
return n, err
} }
// Stream is a single stream of data and implements io.ReadWriteCloser. // Stream is a single stream of data and implements io.ReadWriteCloser.