fix race condition in lazy negotiate

multistream.go

@@ -161,11 +161,15 @@ func (msm *MultistreamMuxer) NegotiateLazy(rwc io.ReadWriteCloser) (Multistream,
 	lzc := &lazyConn{
 		con:        rwc,
 		rhandshake: true,
+		rhsync:     true,
 	}
 
+	// take lock here to prevent a race condition where the reads below from
+	// finishing and taking the write lock before this goroutine can
+	lzc.whlock.Lock()
+
 	go func() {
 		defer close(writeErr)
-		lzc.whlock.Lock()
 		defer lzc.whlock.Unlock()
 		lzc.whsync = true
 
@@ -182,6 +186,7 @@ func (msm *MultistreamMuxer) NegotiateLazy(rwc io.ReadWriteCloser) (Multistream,
 				return
 			}
 		}
+		lzc.whandshake = true
 	}()
 
 	line, err := ReadNextToken(rwc)

multistream_test.go

@@ -86,6 +86,64 @@ func TestProtocolNegotiationLazy(t *testing.T) {
 	verifyPipe(t, ac, b)
 }
 
+func TestNegLazyStress(t *testing.T) {
+	count := 1000
+
+	mux := NewMultistreamMuxer()
+	mux.AddHandler("/a", nil)
+	mux.AddHandler("/b", nil)
+	mux.AddHandler("/c", nil)
+
+	message := []byte("this is the message")
+	listener := make(chan io.ReadWriteCloser)
+	go func() {
+		for rwc := range listener {
+			m, selected, _, err := mux.NegotiateLazy(rwc)
+			if err != nil {
+				t.Error(err)
+				return
+			}
+
+			if selected != "/a" {
+				t.Error("incorrect protocol selected")
+				return
+			}
+
+			_, err = m.Read(nil)
+			if err != nil {
+				t.Error(err)
+				return
+			}
+
+			_, err = m.Write(message)
+			if err != nil {
+				t.Error(err)
+				return
+			}
+		}
+	}()
+
+	for i := 0; i < count; i++ {
+		a, b := net.Pipe()
+		listener <- a
+
+		ms := NewMSSelect(b, "/a")
+
+		buf := make([]byte, len(message))
+		_, err := io.ReadFull(ms, buf)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if !bytes.Equal(message, buf) {
+			t.Fatal("incorrect output: ", buf)
+		}
+
+		a.Close()
+		b.Close()
+	}
+}
+
 func TestInvalidProtocol(t *testing.T) {
 	a, b := net.Pipe()