From 22c3f92f5faea8db492fb0f5ae4daf0d2752b19e Mon Sep 17 00:00:00 2001
From: Gyu-Ho Lee <gyuhox@gmail.com>
Date: Wed, 4 Oct 2017 14:44:57 -0700
Subject: [PATCH] transport: fix racey send to writes channel in WriteStatus
 (#1546)

Concurrent 'SendMsg' calls to stream lead to
multiple 'WriteStatus' calls, while closing
'writes' channel is not synchronized.

This patch marks 'streamDone' first before 'ht.do',
so that following 'WriteStatus' does not trigger panic
on 'writes' channel.

Signed-off-by: Gyu-Ho Lee <gyuhox@gmail.com>
---
 transport/handler_server.go      |  5 +----
 transport/handler_server_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/transport/handler_server.go b/transport/handler_server.go
index 0489fada..f1f6caf8 100644
--- a/transport/handler_server.go
+++ b/transport/handler_server.go
@@ -173,7 +173,6 @@ func (ht *serverHandlerTransport) do(fn func()) error {
 		case <-ht.closedCh:
 			return ErrConnClosing
 		}
-
 	}
 }
 
@@ -183,6 +182,7 @@ func (ht *serverHandlerTransport) WriteStatus(s *Stream, st *status.Status) erro
 		ht.mu.Unlock()
 		return nil
 	}
+	ht.streamDone = true
 	ht.mu.Unlock()
 	err := ht.do(func() {
 		ht.writeCommonHeaders(s)
@@ -223,9 +223,6 @@ func (ht *serverHandlerTransport) WriteStatus(s *Stream, st *status.Status) erro
 		}
 	})
 	close(ht.writes)
-	ht.mu.Lock()
-	ht.streamDone = true
-	ht.mu.Unlock()
 	return err
 }
 
diff --git a/transport/handler_server_test.go b/transport/handler_server_test.go
index 262e6016..06fe813c 100644
--- a/transport/handler_server_test.go
+++ b/transport/handler_server_test.go
@@ -26,6 +26,7 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"reflect"
+	"sync"
 	"testing"
 	"time"
 
@@ -390,6 +391,30 @@ func TestHandlerTransport_HandleStreams_Timeout(t *testing.T) {
 	}
 }
 
+func TestHandlerTransport_HandleStreams_MultiWriteStatus(t *testing.T) {
+	st := newHandleStreamTest(t)
+	handleStream := func(s *Stream) {
+		if want := "/service/foo.bar"; s.method != want {
+			t.Errorf("stream method = %q; want %q", s.method, want)
+		}
+		st.bodyw.Close() // no body
+
+		var wg sync.WaitGroup
+		wg.Add(5)
+		for i := 0; i < 5; i++ {
+			go func() {
+				defer wg.Done()
+				st.ht.WriteStatus(s, status.New(codes.OK, ""))
+			}()
+		}
+		wg.Wait()
+	}
+	st.ht.HandleStreams(
+		func(s *Stream) { go handleStream(s) },
+		func(ctx context.Context, method string) context.Context { return ctx },
+	)
+}
+
 func TestHandlerTransport_HandleStreams_ErrDetails(t *testing.T) {
 	errDetails := []proto.Message{
 		&epb.RetryInfo{