From b8e47c5dea7dded7ee27a93b28771784421c2ed2 Mon Sep 17 00:00:00 2001
From: iulianpascalau <iulian@elrond.network>
Date: Thu, 14 Feb 2019 20:19:03 +0200
Subject: [PATCH] mock package: - changed GetStreams function from mock_conn.go
 to call allStreams which has the same functionality but is concurrent safe -
 changed protocol field from mock_stream to be concurrent safe. Taken the
 implementation from swarm_stream.go

---
 p2p/net/mock/mock_conn.go   | 6 +-----
 p2p/net/mock/mock_stream.go | 9 ++++++---
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/p2p/net/mock/mock_conn.go b/p2p/net/mock/mock_conn.go
index f4cf10358..3e0af7268 100644
--- a/p2p/net/mock/mock_conn.go
+++ b/p2p/net/mock/mock_conn.go
@@ -121,11 +121,7 @@ func (c *conn) NewStream() (inet.Stream, error) {
 }
 
 func (c *conn) GetStreams() []inet.Stream {
-	var out []inet.Stream
-	for e := c.streams.Front(); e != nil; e = e.Next() {
-		out = append(out, e.Value.(*stream))
-	}
-	return out
+	return c.allStreams()
 }
 
 // LocalMultiaddr is the Multiaddr on this side
diff --git a/p2p/net/mock/mock_stream.go b/p2p/net/mock/mock_stream.go
index 8d0150504..331b8960b 100644
--- a/p2p/net/mock/mock_stream.go
+++ b/p2p/net/mock/mock_stream.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"io"
 	"net"
+	"sync/atomic"
 	"time"
 
 	inet "github.com/libp2p/go-libp2p-net"
@@ -24,7 +25,7 @@ type stream struct {
 
 	writeErr error
 
-	protocol protocol.ID
+	protocol atomic.Value
 	stat     inet.Stat
 }
 
@@ -70,7 +71,9 @@ func (s *stream) Write(p []byte) (n int, err error) {
 }
 
 func (s *stream) Protocol() protocol.ID {
-	return s.protocol
+	// Ignore type error. It means that the protocol is unset.
+	p, _ := s.protocol.Load().(protocol.ID)
+	return p
 }
 
 func (s *stream) Stat() inet.Stat {
@@ -78,7 +81,7 @@ func (s *stream) Stat() inet.Stat {
 }
 
 func (s *stream) SetProtocol(proto protocol.ID) {
-	s.protocol = proto
+	s.protocol.Store(proto)
 }
 
 func (s *stream) Close() error {