From 77baeb8d7959628bf113a9cde2e052a5be785227 Mon Sep 17 00:00:00 2001
From: Josh Rickmar <jrick@conformal.com>
Date: Mon, 15 Jun 2015 11:38:25 -0400
Subject: [PATCH] Fix data race.

This synchronizes access to the (*Client).disconnect channel, which
may change during reconnects.

Fixes #54.
---
 infrastructure.go | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/infrastructure.go b/infrastructure.go
index 1ae078a0..673bd8a8 100644
--- a/infrastructure.go
+++ b/infrastructure.go
@@ -394,6 +394,16 @@ out:
 	log.Tracef("RPC client input handler done for %s", c.config.Host)
 }
 
+// disconnectChan returns a copy of the current disconnect channel.  The channel
+// is read protected by the client mutex, and is safe to call while the channel
+// is being reassigned during a reconnect.
+func (c *Client) disconnectChan() <-chan struct{} {
+	c.mtx.Lock()
+	ch := c.disconnect
+	c.mtx.Unlock()
+	return ch
+}
+
 // wsOutHandler handles all outgoing messages for the websocket connection.  It
 // uses a buffered channel to serialize output messages while allowing the
 // sender to continue running asynchronously.  It must be run as a goroutine.
@@ -410,7 +420,7 @@ out:
 				break out
 			}
 
-		case <-c.disconnect:
+		case <-c.disconnectChan():
 			break out
 		}
 	}
@@ -436,7 +446,7 @@ func (c *Client) sendMessage(marshalledJSON []byte) {
 	// Don't send the message if disconnected.
 	select {
 	case c.sendChan <- marshalledJSON:
-	case <-c.disconnect:
+	case <-c.disconnectChan():
 		return
 	}
 }
@@ -618,9 +628,9 @@ out:
 			// has happened.
 			c.wsConn = wsConn
 			c.retryCount = 0
-			c.disconnect = make(chan struct{})
 
 			c.mtx.Lock()
+			c.disconnect = make(chan struct{})
 			c.disconnected = false
 			c.mtx.Unlock()