From 2991d2c1f27f7b47d1f50e464d7fac45d64daa09 Mon Sep 17 00:00:00 2001
From: Akihiro Suda <suda.akihiro@lab.ntt.co.jp>
Date: Sun, 14 Apr 2019 20:48:42 +0900
Subject: [PATCH] session/grpchijack: fix race

Signed-off-by: Akihiro Suda <akihiro.suda.cz@hco.ntt.co.jp>
---
 session/grpchijack/dial.go | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/session/grpchijack/dial.go b/session/grpchijack/dial.go
index 151ab5498f3ab..4b424e9e6593b 100644
--- a/session/grpchijack/dial.go
+++ b/session/grpchijack/dial.go
@@ -45,14 +45,14 @@ type conn struct {
 	lastBuf []byte
 
 	closedOnce sync.Once
-	readMu     sync.Mutex
+	mu         sync.Mutex
 	err        error
 	closeCh    chan struct{}
 }
 
 func (c *conn) Read(b []byte) (n int, err error) {
-	c.readMu.Lock()
-	defer c.readMu.Unlock()
+	c.mu.Lock()
+	defer c.mu.Unlock()
 
 	if c.lastBuf != nil {
 		n := copy(b, c.lastBuf)
@@ -79,6 +79,8 @@ func (c *conn) Read(b []byte) (n int, err error) {
 }
 
 func (c *conn) Write(b []byte) (int, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	m := &controlapi.BytesMessage{Data: b}
 	if err := c.stream.SendMsg(m); err != nil {
 		return 0, err
@@ -87,6 +89,8 @@ func (c *conn) Write(b []byte) (int, error) {
 }
 
 func (c *conn) Close() (err error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 	c.closedOnce.Do(func() {
 		defer func() {
 			close(c.closeCh)
@@ -99,7 +103,6 @@ func (c *conn) Close() (err error) {
 			}
 		}
 
-		c.readMu.Lock()
 		for {
 			m := new(controlapi.BytesMessage)
 			m.Data = c.buf
@@ -114,7 +117,6 @@ func (c *conn) Close() (err error) {
 			c.buf = m.Data[:cap(m.Data)]
 			c.lastBuf = append(c.lastBuf, c.buf...)
 		}
-		c.readMu.Unlock()
 
 	})
 	return nil