From 32990ebf77a9222f17fae19b3c01866ed26408d3 Mon Sep 17 00:00:00 2001 From: nadoo <287492+nadoo@users.noreply.github.com> Date: Sat, 24 Oct 2020 18:55:47 +0800 Subject: [PATCH] proxy: optimize Copy for network connection --- go.mod | 4 ++-- go.sum | 8 ++++---- proxy/conn.go | 26 ++++++++++++++++---------- proxy/ws/frame.go | 22 +++++++++++----------- 4 files changed, 33 insertions(+), 27 deletions(-) diff --git a/go.mod b/go.mod index eccc3d0..79f85e0 100644 --- a/go.mod +++ b/go.mod @@ -11,9 +11,9 @@ require ( github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect github.com/xtaci/kcp-go/v5 v5.6.1 golang.org/x/crypto v0.0.0-20201016220609-9e8e0b390897 - golang.org/x/net v0.0.0-20201022231255-08b38378de70 // indirect + golang.org/x/net v0.0.0-20201024042810-be3efd7ff127 // indirect golang.org/x/sys v0.0.0-20201022201747-fb209a7c41cd // indirect - golang.org/x/tools v0.0.0-20201022215848-ffe8bce740af // indirect + golang.org/x/tools v0.0.0-20201023174141-c8cfbd0f21e6 // indirect gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b // indirect ) diff --git a/go.sum b/go.sum index 3a8f8d2..92eed3a 100644 --- a/go.sum +++ b/go.sum @@ -141,8 +141,8 @@ golang.org/x/net v0.0.0-20200707034311-ab3426394381 h1:VXak5I6aEWmAXeQjA+QSZzlgN golang.org/x/net v0.0.0-20200707034311-ab3426394381/go.mod h1:/O7V0waA8r7cgGh81Ro3o1hOxt32SMVPicZroKQ2sZA= golang.org/x/net v0.0.0-20201021035429-f5854403a974 h1:IX6qOQeG5uLjB/hjjwjedwfjND0hgjPMMyO1RoIXQNI= golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= -golang.org/x/net v0.0.0-20201022231255-08b38378de70 h1:Z6x4N9mAi4oF0TbHweCsH618MO6OI6UFgV0FP5n0wBY= -golang.org/x/net v0.0.0-20201022231255-08b38378de70/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.0.0-20201024042810-be3efd7ff127 h1:pZPp9+iYUqwYKLjht0SDBbRCRK/9gAXDy7pz5fRDpjo= +golang.org/x/net v0.0.0-20201024042810-be3efd7ff127/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20200625203802-6e8e738ad208/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= @@ -178,8 +178,8 @@ golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtn golang.org/x/tools v0.0.0-20200425043458-8463f397d07c h1:iHhCR0b26amDCiiO+kBguKZom9aMF+NrFxh9zeKR/XU= golang.org/x/tools v0.0.0-20200425043458-8463f397d07c/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE= golang.org/x/tools v0.0.0-20200808161706-5bf02b21f123/go.mod h1:njjCfa9FT2d7l9Bc6FUM5FLjQPp3cFF28FI3qnDFljA= -golang.org/x/tools v0.0.0-20201022215848-ffe8bce740af h1:L1EV/2WjPBRxjdgDu0bbZ5gWPuQ6yX9HjG6hSb2dFV4= -golang.org/x/tools v0.0.0-20201022215848-ffe8bce740af/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= +golang.org/x/tools v0.0.0-20201023174141-c8cfbd0f21e6 h1:rbvTkL9AkFts1cgI78+gG6Yu1pwaqX6hjSJAatB78E4= +golang.org/x/tools v0.0.0-20201023174141-c8cfbd0f21e6/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4= diff --git a/proxy/conn.go b/proxy/conn.go index 4aac230..98f9c07 100644 --- a/proxy/conn.go +++ b/proxy/conn.go @@ -6,6 +6,7 @@ import ( "io" "net" "os" + "runtime" "sync" "time" @@ -70,18 +71,20 @@ func Relay(left, right net.Conn) error { return nil } -func worthReadFrom(src io.Reader) bool { +func worthTry(src io.Reader) bool { switch v := src.(type) { case *net.TCPConn, *net.UnixConn: return true + case *io.LimitedReader: + return worthTry(v.R) + case *Conn: + return worthTry(v.Conn) case *os.File: fi, err := v.Stat() if err != nil { return false } return fi.Mode().IsRegular() - case *io.LimitedReader: - return worthReadFrom(v.R) default: return false } @@ -95,15 +98,18 @@ func underlyingWriter(c io.Writer) io.Writer { } // Copy copies from src to dst. -// it will try to avoid memory allocating by using WriteTo or ReadFrom method, -// if both failed, then it'll fallback to call CopyBuffer method. func Copy(dst io.Writer, src io.Reader) (written int64, err error) { - if wt, ok := src.(io.WriterTo); ok { - return wt.WriteTo(dst) - } dst = underlyingWriter(dst) - if rt, ok := dst.(io.ReaderFrom); ok && worthReadFrom(src) { - return rt.ReadFrom(src) + switch runtime.GOOS { + case "linux", "windows", "dragonfly", "freebsd", "solaris": + if _, ok := dst.(*net.TCPConn); ok && worthTry(src) { + if wt, ok := src.(io.WriterTo); ok { + return wt.WriteTo(dst) + } + if rt, ok := dst.(io.ReaderFrom); ok { + return rt.ReadFrom(src) + } + } } return CopyBuffer(dst, src) } diff --git a/proxy/ws/frame.go b/proxy/ws/frame.go index ba01f38..58038a2 100644 --- a/proxy/ws/frame.go +++ b/proxy/ws/frame.go @@ -66,27 +66,27 @@ func (w *frameWriter) Write(b []byte) (int, error) { hdr[1] = maskBit } - nPayload, nLenField := len(b), 0 + nPayload, nHeader := len(b), 2 switch { case nPayload <= 125: hdr[1] |= byte(nPayload) case nPayload < 65536: hdr[1] |= 126 - nLenField = 2 - binary.BigEndian.PutUint16(hdr[2:2+nLenField], uint16(nPayload)) + nHeader += 2 + binary.BigEndian.PutUint16(hdr[2:nHeader], uint16(nPayload)) default: hdr[1] |= 127 - nLenField = 8 - binary.BigEndian.PutUint64(hdr[2:2+nLenField], uint64(nPayload)) + nHeader += 8 + binary.BigEndian.PutUint64(hdr[2:nHeader], uint64(nPayload)) } - header := hdr[:2+nLenField] + header := hdr[:nHeader] if w.server { n, err := (&net.Buffers{header, b}).WriteTo(w.Writer) - if int(n) <= len(header) { - return 0, err + if int(n) > nHeader { + return int(n) - nHeader, err } - return int(n) - len(header), err + return 0, err } payload := pool.GetBuffer(nPayload) @@ -98,8 +98,8 @@ func (w *frameWriter) Write(b []byte) (int, error) { } n, err := (&net.Buffers{header, w.maskKey[:], payload}).WriteTo(w.Writer) - if int(n) > len(header)+4 { - return int(n) - len(header) - 4, err + if int(n) > nHeader+4 { + return int(n) - nHeader - 4, err } return 0, err }