transport/server: fix race that could cause a stray header to be sent

internal/transport/http2_server.go

@@ -945,15 +945,16 @@ func (t *http2Server) streamContextErr(s *Stream) error {
 
 // WriteHeader sends the header metadata md back to the client.
 func (t *http2Server) WriteHeader(s *Stream, md metadata.MD) error {
-	if s.updateHeaderSent() {
-		return ErrIllegalHeaderWrite
-	}
-
+	s.hdrMu.Lock()
+	defer s.hdrMu.Unlock()
 	if s.getState() == streamDone {
 		return t.streamContextErr(s)
 	}
 
-	s.hdrMu.Lock()
+	if s.updateHeaderSent() {
+		return ErrIllegalHeaderWrite
+	}
+
 	if md.Len() > 0 {
 		if s.header.Len() > 0 {
 			s.header = metadata.Join(s.header, md)
@@ -962,10 +963,8 @@ func (t *http2Server) WriteHeader(s *Stream, md metadata.MD) error {
 		}
 	}
 	if err := t.writeHeaderLocked(s); err != nil {
-		s.hdrMu.Unlock()
 		return status.Convert(err).Err()
 	}
-	s.hdrMu.Unlock()
 	return nil
 }
 
@@ -1013,17 +1012,19 @@ func (t *http2Server) writeHeaderLocked(s *Stream) error {
 // TODO(zhaoq): Now it indicates the end of entire stream. Revisit if early
 // OK is adopted.
 func (t *http2Server) WriteStatus(s *Stream, st *status.Status) error {
+	s.hdrMu.Lock()
+	defer s.hdrMu.Unlock()
+
 	if s.getState() == streamDone {
 		return nil
 	}
-	s.hdrMu.Lock()
+
 	// TODO(mmukhi): Benchmark if the performance gets better if count the metadata and other header fields
 	// first and create a slice of that exact size.
 	headerFields := make([]hpack.HeaderField, 0, 2) // grpc-status and grpc-message will be there if none else.
 	if !s.updateHeaderSent() {                      // No headers have been sent.
 		if len(s.header) > 0 { // Send a separate header frame.
 			if err := t.writeHeaderLocked(s); err != nil {
-				s.hdrMu.Unlock()
 				return err
 			}
 		} else { // Send a trailer only response.
@@ -1052,7 +1053,7 @@ func (t *http2Server) WriteStatus(s *Stream, st *status.Status) error {
 		endStream: true,
 		onWrite:   t.setResetPingStrikes,
 	}
-	s.hdrMu.Unlock()
+
 	success, err := t.controlBuf.execute(t.checkForHeaderListSize, trailingHeader)
 	if !success {
 		if err != nil {

test/end2end_test.go

@@ -8077,3 +8077,33 @@ func (s) TestUnexpectedEOF(t *testing.T) {
 		}
 	}
 }
+
+// TestRecvWhileReturningStatus performs a Recv in a service handler while the
+// handler returns its status.  A race condition could result in the server
+// sending headers twice (once as a result of the status and once as a result
+// of the failed Recv call).
+func (s) TestRecvWhileReturningStatus(t *testing.T) {
+	ss := &stubserver.StubServer{
+		FullDuplexCallF: func(stream testpb.TestService_FullDuplexCallServer) error {
+			// The client never sends, so this Recv blocks until the server
+			// returns and causes stream operations to return errors.
+			go stream.Recv()
+			return nil
+		},
+	}
+	if err := ss.Start(nil); err != nil {
+		t.Fatalf("Error starting endpoint server: %v", err)
+	}
+	defer ss.Stop()
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	for i := 0; i < 100; i++ {
+		stream, err := ss.Client.FullDuplexCall(ctx)
+		if err != nil {
+			t.Fatalf("Error while creating stream: %v", err)
+		}
+		if _, err := stream.Recv(); err != io.EOF {
+			t.Fatalf("stream.Recv() = %v, want io.EOF", err)
+		}
+	}
+}