Fix race marshalling a bitfield after a Have has also been posted

diff --git a/connection.go b/connection.go
index 5c105c72679b79c20e05c372075328213d536b05..6e0edf2b62e5fd72f1b4678d97dd67c28dc8ad93 100644 (file)
--- a/connection.go
+++ b/connection.go
@@ -217,8 +217,10 @@ func (cn *connection) Close() {
        cn.closed.Set()
        cn.discardPieceInclination()
        cn.pieceRequestOrder.Clear()
-       // TODO: This call blocks sometimes, why?
-       go cn.conn.Close()
+       if cn.conn != nil {
+               // TODO: This call blocks sometimes, why?
+               go cn.conn.Close()
+       }
 }
 
 func (cn *connection) PeerHasPiece(piece int) bool {
@@ -465,7 +467,10 @@ func (cn *connection) Bitfield(haves []bool) {
                Type:     pp.Bitfield,
                Bitfield: haves,
        })
-       cn.sentHaves = haves
+       // Make a copy of haves, as that's read when the message is marshalled
+       // without the lock. Also it obviously shouldn't change in the Msg due to
+       // changes in .sentHaves.
+       cn.sentHaves = append([]bool(nil), haves...)
 }
 
 func (cn *connection) updateRequests() {

diff --git a/connection_test.go b/connection_test.go
index 29d032accc57dde33801baa6fb707465fc599a82..2f6bed4e2b09bcbfad4a3084101d299a4badd02e 100644 (file)
--- a/connection_test.go
+++ b/connection_test.go
@@ -1,6 +1,7 @@
 package torrent
 
 import (
+       "container/list"
        "io"
        "io/ioutil"
        "net"
@@ -57,3 +58,37 @@ func TestCancelRequestOptimized(t *testing.T) {
        // trying to flush it, and then promptly close.
        require.EqualValues(t, "\x00\x00\x00\x00", string(b))
 }
+
+// Ensure that no race exists between sending a bitfield, and a subsequent
+// Have that would potentially alter it.
+func TestSendBitfieldThenHave(t *testing.T) {
+       r, w := io.Pipe()
+       c := &connection{
+               t: &Torrent{
+                       cl: &Client{},
+               },
+               rw: struct {
+                       io.Reader
+                       io.Writer
+               }{r, w},
+               outgoingUnbufferedMessages: list.New(),
+       }
+       go c.writer(time.Minute)
+       c.mu().Lock()
+       c.Bitfield([]bool{false, true, false})
+       c.mu().Unlock()
+       c.mu().Lock()
+       c.Have(2)
+       c.mu().Unlock()
+       b := make([]byte, 15)
+       n, err := io.ReadFull(r, b)
+       c.mu().Lock()
+       // This will cause connection.writer to terminate.
+       c.closed.Set()
+       c.mu().Unlock()
+       require.NoError(t, err)
+       require.EqualValues(t, 15, n)
+       // Here we see that the bitfield doesn't have piece 2 set, as that should
+       // arrive in the following Have message.
+       require.EqualValues(t, "\x00\x00\x00\x02\x05@\x00\x00\x00\x05\x04\x00\x00\x00\x02", string(b))
+}