upgrade to github.com/xtaci/smux@v1.4.5 github.com/xtaci/smux/v2@v2.0.14, optimize Copy to use io.WriteTo

client/main.go

@@ -8,7 +8,6 @@ import (
 	"math/rand"
 	"net"
 	"os"
-	"sync"
 	"time"
 
 	"golang.org/x/crypto/pbkdf2"
@@ -30,9 +29,6 @@ const maxSmuxVer = 2
 // VERSION is injected by buildflags
 var VERSION = "SELFBUILD"
 
-// A pool for stream copying
-var xmitBuf sync.Pool
-
 func handleClient(mux generic.Mux, p1 net.Conn, quiet bool) {
 	logln := func(v ...interface{}) {
 		if !quiet {
@@ -57,8 +53,7 @@ func handleClient(mux generic.Mux, p1 net.Conn, quiet bool) {
 	streamCopy := func(dst io.Writer, src io.ReadCloser) chan struct{} {
 		die := make(chan struct{})
 		go func() {
-			buf := xmitBuf.Get().([]byte)
-			if _, err := generic.CopyBuffer(dst, src, buf); err != nil {
+			if _, err := generic.Copy(dst, src); err != nil {
 				if s2, ok := p2.(generic.Stream); ok {
 					// verbose error handling
 					cause := err
@@ -74,7 +69,6 @@ func handleClient(mux generic.Mux, p1 net.Conn, quiet bool) {
 					}
 				}
 			}
-			xmitBuf.Put(buf)
 			close(die)
 		}()
 		return die
@@ -99,9 +93,6 @@ func main() {
 		// add more log flags for debugging
 		log.SetFlags(log.LstdFlags | log.Lshortfile)
 	}
-	xmitBuf.New = func() interface{} {
-		return make([]byte, 4096)
-	}
 
 	myApp := cli.NewApp()
 	myApp.Name = "kcptun"

generic/copy.go

@@ -2,12 +2,28 @@ package generic
 
 import "io"
 
-// io.CopyBuffer has extra tests for interface like io.ReaderFrom and io.WriterTo
-// which is not efficient in memory management from tests
-func CopyBuffer(dst io.Writer, src io.Reader, buf []byte) (written int64, err error) {
-	if buf != nil && len(buf) == 0 {
-		panic("empty buffer in copyBuffer")
+// Memory optimized Copy function specified for this library
+func Copy(dst io.Writer, src io.Reader) (written int64, err error) {
+	// If the reader has a WriteTo method, use it to do the copy.
+	// Avoids an allocation and a copy.
+	if wt, ok := src.(io.WriterTo); ok {
+		return wt.WriteTo(dst)
 	}
+	// Similarly, if the writer has a ReadFrom method, use it to do the copy.
+	if rt, ok := dst.(io.ReaderFrom); ok {
+		return rt.ReadFrom(src)
+	}
+
+	// limited to 4K per stream
+	size := 4096
+	if l, ok := src.(*io.LimitedReader); ok && int64(size) > l.N {
+		if l.N < 1 {
+			size = 1
+		} else {
+			size = int(l.N)
+		}
+	}
+	buf := make([]byte, size)
 
 	for {
 		nr, er := src.Read(buf)

go.mod

@@ -13,8 +13,8 @@ require (
 	github.com/urfave/cli v1.21.0
 	github.com/xtaci/kcp-go v5.4.19+incompatible
 	github.com/xtaci/lossyconn v0.0.0-20190602105132-8df528c0c9ae // indirect
-	github.com/xtaci/smux v1.4.4
-	github.com/xtaci/smux/v2 v2.0.13
+	github.com/xtaci/smux v1.4.5
+	github.com/xtaci/smux/v2 v2.0.14
 	github.com/xtaci/tcpraw v1.2.25
 	golang.org/x/crypto v0.0.0-20190909091759-094676da4a83
 	golang.org/x/net v0.0.0-20190909003024-a7b16738d86b // indirect

go.sum

@@ -37,10 +37,14 @@ github.com/xtaci/lossyconn v0.0.0-20190602105132-8df528c0c9ae h1:J0GxkO96kL4WF+A
 github.com/xtaci/lossyconn v0.0.0-20190602105132-8df528c0c9ae/go.mod h1:gXtu8J62kEgmN++bm9BVICuT/e8yiLI2KFobd/TRFsE=
 github.com/xtaci/smux v1.4.4 h1:FukIfahko+KHhS9Gxppkp6756opZymvPOLNmpny1is4=
 github.com/xtaci/smux v1.4.4/go.mod h1:LuA3S0xssf4fmGRJ7ow3EehgmDUzib4EcobaFNKvlMA=
+github.com/xtaci/smux v1.4.5 h1:Pi64WSFFIDrpMSzRZzqt5e4YO/m4a6BAKGjxSa+KU0A=
+github.com/xtaci/smux v1.4.5/go.mod h1:LuA3S0xssf4fmGRJ7ow3EehgmDUzib4EcobaFNKvlMA=
 github.com/xtaci/smux/v2 v2.0.11 h1:thVWmgGRciZ8iaATwpY2B/51aHzmMI6wrF7DfcJSckU=
 github.com/xtaci/smux/v2 v2.0.11/go.mod h1:Iqy5a3Gax2p7WCKHOHkSNo/COthNFXd3/vqrcKNtzqI=
 github.com/xtaci/smux/v2 v2.0.13 h1:T3VtVS8CQ2nbsMttkzTB48acij5HAhi8x8gYe9SodJA=
 github.com/xtaci/smux/v2 v2.0.13/go.mod h1:Iqy5a3Gax2p7WCKHOHkSNo/COthNFXd3/vqrcKNtzqI=
+github.com/xtaci/smux/v2 v2.0.14 h1:XfmTsDX7lNTLPP8bCLRQuviP1AEix3NYefM/QVUg5X4=
+github.com/xtaci/smux/v2 v2.0.14/go.mod h1:Iqy5a3Gax2p7WCKHOHkSNo/COthNFXd3/vqrcKNtzqI=
 github.com/xtaci/tcpraw v1.2.25 h1:VDlqo0op17JeXBM6e2G9ocCNLOJcw9mZbobMbJjo0vk=
 github.com/xtaci/tcpraw v1.2.25/go.mod h1:dKyZ2V75s0cZ7cbgJYdxPvms7af0joIeOyx1GgJQbLk=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=

server/main.go

@@ -32,9 +32,6 @@ const maxSmuxVer = 2
 // VERSION is injected by buildflags
 var VERSION = "SELFBUILD"
 
-// A pool for stream copying
-var xmitBuf sync.Pool
-
 // handle multiplex-ed connection
 func handleMux(conn net.Conn, config *Config) {
 	// check if target is unix domain socket
@@ -121,8 +118,7 @@ func handleClient(p1 io.ReadWriteCloser, p2 net.Conn, quiet bool) {
 	streamCopy := func(dst io.Writer, src io.ReadCloser) chan struct{} {
 		die := make(chan struct{})
 		go func() {
-			buf := xmitBuf.Get().([]byte)
-			if _, err := generic.CopyBuffer(dst, src, buf); err != nil {
+			if _, err := generic.Copy(dst, src); err != nil {
 				if s1, ok := p1.(generic.Stream); ok {
 					// verbose error handling
 					cause := err
@@ -138,7 +134,6 @@ func handleClient(p1 io.ReadWriteCloser, p2 net.Conn, quiet bool) {
 					}
 				}
 			}
-			xmitBuf.Put(buf)
 			close(die)
 		}()
 		return die
@@ -163,9 +158,6 @@ func main() {
 		// add more log flags for debugging
 		log.SetFlags(log.LstdFlags | log.Lshortfile)
 	}
-	xmitBuf.New = func() interface{} {
-		return make([]byte, 4096)
-	}
 
 	myApp := cli.NewApp()
 	myApp.Name = "kcptun"