diff --git a/client/main.go b/client/main.go index f569762..5d5ec15 100644 --- a/client/main.go +++ b/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" diff --git a/generic/copy.go b/generic/copy.go index 023556b..8a1105a 100644 --- a/generic/copy.go +++ b/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) diff --git a/go.mod b/go.mod index dc029aa..b77f76d 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index da333ff..5da8b6f 100644 --- a/go.sum +++ b/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= diff --git a/server/main.go b/server/main.go index 011312a..5aff7ae 100644 --- a/server/main.go +++ b/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"