session/grpchijack: fix race

Signed-off-by: Akihiro Suda <akihiro.suda.cz@hco.ntt.co.jp>

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