add mutex for write/close

libp2p · Dec 17, 2019 · f6ca84a · f6ca84a
1 parent 3098bba
commit f6ca84a
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/conn_native.go b/conn_native.go
@@ -17,6 +17,7 @@ type Conn struct {
 	DefaultMessageType int
 	reader             io.Reader
 	closeOnce          sync.Once
+	mx                 sync.Mutex
 }
 
 func (c *Conn) Read(b []byte) (int, error) {
@@ -67,6 +68,9 @@ func (c *Conn) prepNextReader() error {
 }
 
 func (c *Conn) Write(b []byte) (n int, err error) {
+	c.mx.Lock()
+	defer c.mx.Unlock()
+
 	if err := c.Conn.WriteMessage(c.DefaultMessageType, b); err != nil {
 		return 0, err
 	}
@@ -78,6 +82,9 @@ func (c *Conn) Write(b []byte) (n int, err error) {
 // close error, subsequent and concurrent calls will return nil.
 // This method is thread-safe.
 func (c *Conn) Close() error {
+	c.mx.Lock()
+	defer c.mx.Unlock()
+
 	var err error
 	c.closeOnce.Do(func() {
 		err1 := c.Conn.WriteControl(