From dcfee93f96d231b5590f991313b5d9f925757f52 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Wed, 11 May 2016 23:50:21 +1000
Subject: [PATCH] Fix race marshalling a bitfield after a Have has also been
 posted

---
 connection.go      | 11 ++++++++---
 connection_test.go | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/connection.go b/connection.go
index 5c105c72..6e0edf2b 100644
--- 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 29d032ac..2f6bed4e 100644
--- 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))
+}
-- 
2.48.1