Fix unsafe header re-use on close and timeout (#2)

mafredri · web-flow · commit 12254a01acb0 · 2022-07-11T12:15:06.000+03:00
This commit takes a minimal approach to fixing unsafe header re-use during close and timeout by reallocating the header in those cases. Partially fixes: coder/coder#2429 Related: hashicorp#40
diff --git a/stream.go b/stream.go
@@ -2,6 +2,7 @@ package yamux
 
 import (
 	"bytes"
+	"errors"
 	"io"
 	"sync"
 	"sync/atomic"
@@ -200,6 +201,10 @@ START:
 	// Send the header
 	s.sendHdr.encode(typeData, flags, s.id, max)
 	if err = s.session.waitForSendErr(s.sendHdr, body, s.sendErr); err != nil {
+		if errors.Is(err, ErrSessionShutdown) || errors.Is(err, ErrConnectionWriteTimeout) {
+			// Message left in ready queue, header re-use is unsafe.
+			s.sendHdr = header(make([]byte, headerSize))
+		}
 		return 0, err
 	}
 
@@ -273,6 +278,10 @@ func (s *Stream) sendWindowUpdate() error {
 	// Send the header
 	s.controlHdr.encode(typeWindowUpdate, flags, s.id, delta)
 	if err := s.session.waitForSendErr(s.controlHdr, nil, s.controlErr); err != nil {
+		if errors.Is(err, ErrSessionShutdown) || errors.Is(err, ErrConnectionWriteTimeout) {
+			// Message left in ready queue, header re-use is unsafe.
+			s.controlHdr = header(make([]byte, headerSize))
+		}
 		return err
 	}
 	return nil
@@ -287,6 +296,10 @@ func (s *Stream) sendClose() error {
 	flags |= flagFIN
 	s.controlHdr.encode(typeWindowUpdate, flags, s.id, 0)
 	if err := s.session.waitForSendErr(s.controlHdr, nil, s.controlErr); err != nil {
+		if errors.Is(err, ErrSessionShutdown) || errors.Is(err, ErrConnectionWriteTimeout) {
+			// Message left in ready queue, header re-use is unsafe.
+			s.controlHdr = header(make([]byte, headerSize))
+		}
 		return err
 	}
 	return nil
@@ -362,8 +375,9 @@ func (s *Stream) closeTimeout() {
 	// Send a RST so the remote side closes too.
 	s.sendLock.Lock()
 	defer s.sendLock.Unlock()
-	s.sendHdr.encode(typeWindowUpdate, flagRST, s.id, 0)
-	s.session.sendNoWait(s.sendHdr)
+	hdr := header(make([]byte, headerSize))
+	hdr.encode(typeWindowUpdate, flagRST, s.id, 0)
+	s.session.sendNoWait(hdr)
 }
 
 // forceClose is used for when the session is exiting
