Add synchronization to prevent post-melt collects

This fixes a race condition in which snowflakes.End() is called while
snowflakes.Collect() is in progress resulting in a write to a closed
channel. We now wait for all in-progress collections to finish and add
an extra check before proceeding with a collection.

2 files changed, 13 insertions, 3 deletions

diff --git a/client/lib/peers.go b/client/lib/peers.go
index d864fc8..d02eed3 100644
--- a/client/lib/peers.go
+++ b/client/lib/peers.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"sync"
 )
 
 // Container which keeps track of multiple WebRTC remote peers.
@@ -25,7 +26,10 @@ type Peers struct {
 	snowflakeChan chan *WebRTCPeer
 	activePeers   *list.List
 
-	melt chan struct{}
+	melt   chan struct{}
+	melted bool
+
+	collection sync.WaitGroup
 }
 
 // Construct a fresh container of remote peers.
@@ -45,6 +49,11 @@ func NewPeers(tongue Tongue) (*Peers, error) {
 // As part of |SnowflakeCollector| interface.
 func (p *Peers) Collect() (*WebRTCPeer, error) {
 	// Engage the Snowflake Catching interface, which must be available.
+	p.collection.Add(1)
+	defer p.collection.Done()
+	if p.melted {
+		return nil, fmt.Errorf("Snowflakes have melted")
+	}
 	if nil == p.Tongue {
 		return nil, errors.New("missing Tongue to catch Snowflakes with")
 	}
@@ -110,8 +119,10 @@ func (p *Peers) purgeClosedPeers() {
 
 // Close all Peers contained here.
 func (p *Peers) End() {
-	close(p.snowflakeChan)
 	close(p.melt)
+	p.melted = true
+	p.collection.Wait()
+	close(p.snowflakeChan)
 	cnt := p.Count()
 	for e := p.activePeers.Front(); e != nil; {
 		next := e.Next()
diff --git a/client/lib/snowflake.go b/client/lib/snowflake.go
index 0ba5667..e888160 100644
--- a/client/lib/snowflake.go
+++ b/client/lib/snowflake.go
@@ -181,7 +181,6 @@ func Handler(socks net.Conn, tongue Tongue) error {
 // transfer to the Tor SOCKS handler when needed.
 func connectLoop(snowflakes SnowflakeCollector) {
 	for {
-		// Check if ending is necessary.
 		_, err := snowflakes.Collect()
 		if err != nil {
 			log.Printf("WebRTC: %v  Retrying in %v...",