From 2cf220a1373d8903e3218c9eae02ec83298b3724 Mon Sep 17 00:00:00 2001 From: Conner Fromknecht Date: Thu, 10 Jan 2019 19:49:44 -0800 Subject: [PATCH] peer: handle message encoding errors on write In this commit, we modify the peer's writeMessage method to properly handle errors returned from encoding an lnwire message and from setting the write deadline on the connection. Since an error would likely result in an empty byte slice, the worse case seems to be that we may have tried to send an empty message on the wire. Lastly, we correct the way we compute bytes sent on the wire to properly count the number of bytes *written*, and not just the length of the encoded message. --- peer.go | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/peer.go b/peer.go index 41f32299..62a8d297 100644 --- a/peer.go +++ b/peer.go @@ -1367,13 +1367,25 @@ func (p *peer) writeMessage(msg lnwire.Message) error { // With the temp buffer created and sliced properly (length zero, full // capacity), we'll now encode the message directly into this buffer. - n, err := lnwire.WriteMessage(b, msg, 0) - atomic.AddUint64(&p.bytesSent, uint64(n)) + _, err := lnwire.WriteMessage(b, msg, 0) + if err != nil { + return err + } - p.conn.SetWriteDeadline(time.Now().Add(writeMessageTimeout)) + // Compute and set the write deadline we will impose on the remote peer. + writeDeadline := time.Now().Add(writeMessageTimeout) + err = p.conn.SetWriteDeadline(writeDeadline) + if err != nil { + return err + } // Finally, write the message itself in a single swoop. - _, err = p.conn.Write(b.Bytes()) + n, err := p.conn.Write(b.Bytes()) + + // Regardless of the error returned, record how many bytes were written + // to the wire. + atomic.AddUint64(&p.bytesSent, uint64(n)) + return err }