use number of bytes returned from read call

cenkalti · Mar 29, 2019 · da18d26 · da18d26
1 parent d410952
commit da18d26
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/internal/peerconn/peerwriter/peerwriter.go b/internal/peerconn/peerwriter/peerwriter.go
@@ -164,18 +164,19 @@ func (p *PeerWriter) messageWriter() {
 			// Reserve space for length and message ID
 			buf.Write([]byte{0, 0, 0, 0, 0})
 
+			var m int64
 			if wt, ok := msg.(io.WriterTo); ok {
-				_, err = wt.WriteTo(buf)
+				m, err = wt.WriteTo(buf)
 			} else {
-				_, err = buf.ReadFrom(msg)
+				m, err = buf.ReadFrom(msg)
 			}
 			if err != nil {
 				p.log.Errorf("cannot serialize message [%v]: %s", msg.ID(), err.Error())
 				return
 			}
 
 			// Put length
-			binary.BigEndian.PutUint32(buf.Bytes()[:4], uint32(1+buf.Len()-5))
+			binary.BigEndian.PutUint32(buf.Bytes()[:4], uint32(1+m))
 			// Put message ID
 			buf.Bytes()[4] = uint8(msg.ID())