fix stream header sending for concurrent Close / CancelWrite calls

Author: marten-seemann

stream.go

@@ -45,6 +45,9 @@ type SendStream struct {
 	// Might be initialized to nil if this sendStream is part of an incoming bidirectional stream.
 	streamHdr   []byte
 	streamHdrMu sync.Mutex
+	// Set to true when a goroutine is spawned to send the header asynchronously.
+	// This only happens if the stream is closed / reset immediately after creation.
+	sendingHdrAsync bool
 
 	onClose func() // to remove the stream from the streamsMap
 
@@ -147,15 +150,25 @@ func (s *SendStream) maybeSendStreamHeader() error {
 // Write will unblock immediately, and future calls to Write will fail.
 // When called multiple times it is a no-op.
 func (s *SendStream) CancelWrite(e StreamErrorCode) {
+	// if a Goroutine is already sending the header, return immediately
 	s.streamHdrMu.Lock()
+	if s.sendingHdrAsync {
+		s.streamHdrMu.Unlock()
+		return
+	}
+
 	if len(s.streamHdr) > 0 {
 		// Sending the stream header might block if we are blocked by flow control.
 		// Send a stream header async so that CancelWrite can return immediately.
-		go func() {
-			defer s.streamHdrMu.Unlock()
+		s.sendingHdrAsync = true
+		streamHdr := s.streamHdr
+		s.streamHdr = nil
+		s.streamHdrMu.Unlock()
 
+		go func() {
 			s.SetWriteDeadline(time.Time{})
-			_ = s.maybeSendStreamHeader()
+			_, _ = s.str.Write(streamHdr)
+			s.str.SetReliableBoundary()
 			s.str.CancelWrite(webtransportCodeToHTTPCode(e))
 			s.onClose()
 		}()
@@ -178,16 +191,26 @@ func (s *SendStream) closeWithSession(err error) {
 // Close closes the write-direction of the stream.
 // Future calls to Write are not permitted after calling Close.
 func (s *SendStream) Close() error {
+	// if a Goroutine is already sending the header, return immediately
 	s.streamHdrMu.Lock()
+	if s.sendingHdrAsync {
+		s.streamHdrMu.Unlock()
+		return nil
+	}
+
 	if len(s.streamHdr) > 0 {
 		// Sending the stream header might block if we are blocked by flow control.
-		// Send a stream header async so that Close can return immediately.
-		go func() {
-			defer s.streamHdrMu.Unlock()
+		// Send a stream header async so that CancelWrite can return immediately.
+		s.sendingHdrAsync = true
+		streamHdr := s.streamHdr
+		s.streamHdr = nil
+		s.streamHdrMu.Unlock()
 
+		go func() {
 			s.SetWriteDeadline(time.Time{})
-			_ = s.maybeSendStreamHeader()
-			s.str.Close()
+			_, _ = s.str.Write(streamHdr)
+			s.str.SetReliableBoundary()
+			_ = s.str.Close()
 			s.onClose()
 		}()
 		return nil