synchronise sending streamHeader

stream.go

@@ -46,6 +46,8 @@ type sendStream struct {
 	streamHdr []byte
 
 	onClose func()
+
+	once sync.Once
 }
 
 var _ SendStream = &sendStream{}
@@ -54,15 +56,15 @@ func newSendStream(str quic.SendStream, hdr []byte, onClose func()) *sendStream
 	return &sendStream{str: str, streamHdr: hdr, onClose: onClose}
 }
 
-func (s *sendStream) maybeSendStreamHeader() error {
-	if len(s.streamHdr) == 0 {
-		return nil
-	}
-	if _, err := s.str.Write(s.streamHdr); err != nil {
-		return err
-	}
-	s.streamHdr = nil
-	return nil
+func (s *sendStream) maybeSendStreamHeader() (err error) {
+	s.once.Do(func() {
+		if _, e := s.str.Write(s.streamHdr); e != nil {
+			err = e
+			return
+		}
+		s.streamHdr = nil
+	})
+	return
 }
 
 func (s *sendStream) Write(b []byte) (int, error) {

webtransport_test.go

@@ -575,3 +575,40 @@ func TestCloseStreamsOnSessionClose(t *testing.T) {
 	_, err = ustr.Read([]byte{0})
 	require.Error(t, err)
 }
+
+func TestWriteCloseRace(t *testing.T) {
+	ch := make(chan struct{})
+	sess, closeServer := establishSession(t, func(sess *webtransport.Session) {
+		str, err := sess.AcceptStream(context.Background())
+		if err != nil {
+			return
+		}
+		defer str.Close()
+		<-ch
+	})
+	defer closeServer()
+	str, err := sess.OpenStream()
+	require.NoError(t, err)
+	ready := make(chan struct{}, 2)
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		ready <- struct{}{}
+		wg.Wait()
+		str.Write([]byte("foobar"))
+		ready <- struct{}{}
+	}()
+	go func() {
+		ready <- struct{}{}
+		wg.Wait()
+		str.Close()
+		ready <- struct{}{}
+	}()
+	<-ready
+	<-ready
+	wg.Add(-2)
+	<-ready
+	<-ready
+	close(ch)
+}