handle ContextErr inside WriteHeader

internal/transport/http2_server.go

@@ -933,9 +933,14 @@ func (t *http2Server) checkForHeaderListSize(it interface{}) bool {
 
 // WriteHeader sends the header metadata md back to the client.
 func (t *http2Server) WriteHeader(s *Stream, md metadata.MD) error {
-	if s.updateHeaderSent() || s.getState() == streamDone {
-		return ErrIllegalHeaderWrite
+	if s.getState() == streamDone {
+		return ContextErr(s.ctx.Err())
 	}
+
+	if s.updateHeaderSent() {
+		status.Errorf(codes.Internal, "transport: %v", ErrIllegalHeaderWrite)
+	}
+
 	s.hdrMu.Lock()
 	if md.Len() > 0 {
 		if s.header.Len() > 0 {
@@ -1062,14 +1067,7 @@ func (t *http2Server) WriteStatus(s *Stream, st *status.Status) error {
 func (t *http2Server) Write(s *Stream, hdr []byte, data []byte, opts *Options) error {
 	if !s.isHeaderSent() { // Headers haven't been written yet.
 		if err := t.WriteHeader(s, nil); err != nil {
-			if _, ok := err.(ConnectionError); ok {
-				return err
-			}
-			if s.ctx.Err() != nil {
-				return ContextErr(s.ctx.Err())
-			}
-			// TODO(mmukhi, dfawley): Make sure this is the right code to return.
-			return status.Errorf(codes.Internal, "transport: %v", err)
+			return err
 		}
 	} else {
 		// Writing headers checks for this condition.