From 03a605b0eccba2ef118c06dc447f3bd03a996e25 Mon Sep 17 00:00:00 2001
From: Cuong Do <cdo@cockroachlabs.com>
Date: Wed, 2 Mar 2016 14:28:20 -0500
Subject: [PATCH] Fix #119 - CloseNotify race with ServeHTTP

The generated code for the mux handlers can panic if the call to
ServeHTTP() finishes before the goroutine finishes. This is because
CloseNotify() can only be called before ServeHTTP() finishes. I've
fixed this by calling CloseNotify() before starting the goroutine. The
returned channel can safely be selected any time after that.
---
 .../examplepb/a_bit_of_everything.pb.gw.go    | 33 +++++---
 examples/examplepb/echo_service.pb.gw.go      |  6 +-
 examples/examplepb/flow_combination.pb.gw.go  | 78 ++++++++++++-------
 .../gengateway/template.go                    |  3 +-
 4 files changed, 80 insertions(+), 40 deletions(-)

diff --git a/examples/examplepb/a_bit_of_everything.pb.gw.go b/examples/examplepb/a_bit_of_everything.pb.gw.go
index 93f9b1ae4fa..963cdbcc234 100644
--- a/examples/examplepb/a_bit_of_everything.pb.gw.go
+++ b/examples/examplepb/a_bit_of_everything.pb.gw.go
@@ -495,10 +495,11 @@ func RegisterABitOfEverythingServiceHandler(ctx context.Context, mux *runtime.Se
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -519,10 +520,11 @@ func RegisterABitOfEverythingServiceHandler(ctx context.Context, mux *runtime.Se
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -543,10 +545,11 @@ func RegisterABitOfEverythingServiceHandler(ctx context.Context, mux *runtime.Se
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -567,10 +570,11 @@ func RegisterABitOfEverythingServiceHandler(ctx context.Context, mux *runtime.Se
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -591,10 +595,11 @@ func RegisterABitOfEverythingServiceHandler(ctx context.Context, mux *runtime.Se
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -615,10 +620,11 @@ func RegisterABitOfEverythingServiceHandler(ctx context.Context, mux *runtime.Se
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -639,10 +645,11 @@ func RegisterABitOfEverythingServiceHandler(ctx context.Context, mux *runtime.Se
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -663,10 +670,11 @@ func RegisterABitOfEverythingServiceHandler(ctx context.Context, mux *runtime.Se
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -687,10 +695,11 @@ func RegisterABitOfEverythingServiceHandler(ctx context.Context, mux *runtime.Se
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -711,10 +720,11 @@ func RegisterABitOfEverythingServiceHandler(ctx context.Context, mux *runtime.Se
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -735,10 +745,11 @@ func RegisterABitOfEverythingServiceHandler(ctx context.Context, mux *runtime.Se
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
diff --git a/examples/examplepb/echo_service.pb.gw.go b/examples/examplepb/echo_service.pb.gw.go
index 9cfeb0c150f..a9b765b1c1e 100644
--- a/examples/examplepb/echo_service.pb.gw.go
+++ b/examples/examplepb/echo_service.pb.gw.go
@@ -104,10 +104,11 @@ func RegisterEchoServiceHandler(ctx context.Context, mux *runtime.ServeMux, conn
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -128,10 +129,11 @@ func RegisterEchoServiceHandler(ctx context.Context, mux *runtime.ServeMux, conn
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
diff --git a/examples/examplepb/flow_combination.pb.gw.go b/examples/examplepb/flow_combination.pb.gw.go
index 9f3cc1e0e4e..10d69f98830 100644
--- a/examples/examplepb/flow_combination.pb.gw.go
+++ b/examples/examplepb/flow_combination.pb.gw.go
@@ -1008,10 +1008,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1032,10 +1033,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1056,10 +1058,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1080,10 +1083,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1104,10 +1108,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1128,10 +1133,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1152,10 +1158,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1176,10 +1183,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1200,10 +1208,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1224,10 +1233,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1248,10 +1258,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1272,10 +1283,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1296,10 +1308,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1320,10 +1333,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1344,10 +1358,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1368,10 +1383,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1392,10 +1408,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1416,10 +1433,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1440,10 +1458,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1464,10 +1483,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1488,10 +1508,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1512,10 +1533,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1536,10 +1558,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1560,10 +1583,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1584,10 +1608,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
@@ -1608,10 +1633,11 @@ func RegisterFlowCombinationHandler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()
diff --git a/protoc-gen-grpc-gateway/gengateway/template.go b/protoc-gen-grpc-gateway/gengateway/template.go
index 08c615fa2ff..27ac672c6d0 100644
--- a/protoc-gen-grpc-gateway/gengateway/template.go
+++ b/protoc-gen-grpc-gateway/gengateway/template.go
@@ -273,10 +273,11 @@ func Register{{$svc.GetName}}Handler(ctx context.Context, mux *runtime.ServeMux,
 		defer cancel()
 		closeNotifier, ok := w.(http.CloseNotifier)
 		if ok {
+			closeNotify := closeNotifier.CloseNotify()
 			go func() {
 				select {
 				case <-ctx.Done():
-				case <-closeNotifier.CloseNotify():
+				case <-closeNotify:
 					cancel()
 				}
 			}()