Browse Source

Fix buffer leak in mkcp transport protocol (#3347)

yiitz 6 months ago
parent
commit
d2b2f8c1f3
2 changed files with 13 additions and 1 deletions
  1. 4 1
      transport/internet/kcp/connection.go
  2. 9 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
 	var b *buf.Buffer
-	defer b.Release()
+	defer func() {
+		b.Release()
+	}()
 
 
 	for {
 	for {
 		for {
 		for {
@@ -537,6 +539,7 @@ func (c *Connection) Terminate() {
 	c.closer.Close()
 	c.closer.Close()
 	c.sendingWorker.Release()
 	c.sendingWorker.Release()
 	c.receivingWorker.Release()
 	c.receivingWorker.Release()
+	c.output.Release()
 }
 }
 
 
 func (c *Connection) HandleOption(opt SegmentOption) {
 func (c *Connection) HandleOption(opt SegmentOption) {

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

@@ -10,6 +10,7 @@ import (
 
 
 type SegmentWriter interface {
 type SegmentWriter interface {
 	Write(seg Segment) error
 	Write(seg Segment) error
+	Release()
 }
 }
 
 
 type SimpleSegmentWriter struct {
 type SimpleSegmentWriter struct {
@@ -36,6 +37,10 @@ func (w *SimpleSegmentWriter) Write(seg Segment) error {
 	return err
 	return err
 }
 }
 
 
+func (w *SimpleSegmentWriter) Release() {
+	w.buffer.Release()
+}
+
 type RetryableWriter struct {
 type RetryableWriter struct {
 	writer SegmentWriter
 	writer SegmentWriter
 }
 }
@@ -51,3 +56,7 @@ func (w *RetryableWriter) Write(seg Segment) error {
 		return w.writer.Write(seg)
 		return w.writer.Write(seg)
 	})
 	})
 }
 }
+
+func (w *RetryableWriter) Release() {
+	w.writer.Release()
+}