Browse Source

Fix mkcp unreleased buffer (#3411)

* Fix buffer leak in mkcp transport protocol

* fix Write panic caused by concurrent state check
yiitz 5 months ago
parent
commit
7c5e4c7a4b
2 changed files with 20 additions and 1 deletions
  1. 4 1
      transport/internet/kcp/connection.go
  2. 16 0
      transport/internet/kcp/output.go

+ 4 - 1
transport/internet/kcp/connection.go

@@ -390,7 +390,9 @@ func (c *Connection) writeMultiBufferInternal(reader io.Reader) error {
 	}()
 
 	var b *buf.Buffer
-	defer b.Release()
+	defer func() {
+		b.Release()
+	}()
 
 	for {
 		for {
@@ -537,6 +539,7 @@ func (c *Connection) Terminate() {
 	c.closer.Close()
 	c.sendingWorker.Release()
 	c.receivingWorker.Release()
+	c.output.Release()
 }
 
 func (c *Connection) HandleOption(opt SegmentOption) {

+ 16 - 0
transport/internet/kcp/output.go

@@ -10,12 +10,14 @@ import (
 
 type SegmentWriter interface {
 	Write(seg Segment) error
+	Release()
 }
 
 type SimpleSegmentWriter struct {
 	sync.Mutex
 	buffer *buf.Buffer
 	writer io.Writer
+	closed bool
 }
 
 func NewSegmentWriter(writer io.Writer) SegmentWriter {
@@ -28,6 +30,9 @@ func NewSegmentWriter(writer io.Writer) SegmentWriter {
 func (w *SimpleSegmentWriter) Write(seg Segment) error {
 	w.Lock()
 	defer w.Unlock()
+	if w.closed {
+		return io.ErrClosedPipe
+	}
 
 	w.buffer.Clear()
 	rawBytes := w.buffer.Extend(seg.ByteSize())
@@ -36,6 +41,13 @@ func (w *SimpleSegmentWriter) Write(seg Segment) error {
 	return err
 }
 
+func (w *SimpleSegmentWriter) Release() {
+	w.Lock()
+	defer w.Unlock()
+	w.buffer.Release()
+	w.closed = true
+}
+
 type RetryableWriter struct {
 	writer SegmentWriter
 }
@@ -51,3 +63,7 @@ func (w *RetryableWriter) Write(seg Segment) error {
 		return w.writer.Write(seg)
 	})
 }
+
+func (w *RetryableWriter) Release() {
+	w.writer.Release()
+}