diff --git a/src/frp/models/msg/process.go b/src/frp/models/msg/process.go index 43e34ab..55d99ba 100644 --- a/src/frp/models/msg/process.go +++ b/src/frp/models/msg/process.go @@ -26,6 +26,7 @@ import ( "frp/utils/conn" "frp/utils/log" "frp/utils/pcrypto" + "frp/utils/pool" ) // will block until connection close @@ -115,7 +116,10 @@ func pipeDecrypt(r *conn.Conn, w *conn.Conn, conf config.BaseConf, needRecord bo return fmt.Errorf("Pcrypto Init error: %v", err) } - buf := make([]byte, 5*1024+4) + // get []byte from buffer pool + buf := pool.GetBuf(5*1024 + 4) + defer pool.PutBuf(buf) + var left, res []byte var cnt int = -1 @@ -198,7 +202,10 @@ func pipeEncrypt(r *conn.Conn, w *conn.Conn, conf config.BaseConf, needRecord bo }() } - buf := make([]byte, 5*1024) + // get []byte from buffer pool + buf := pool.GetBuf(5*1024 + 4) + defer pool.PutBuf(buf) + for { n, err := r.Read(buf) if err != nil { diff --git a/src/frp/utils/pool/pool.go b/src/frp/utils/pool/pool.go new file mode 100644 index 0000000..5b7c6d2 --- /dev/null +++ b/src/frp/utils/pool/pool.go @@ -0,0 +1,58 @@ +// Copyright 2016 fatedier, fatedier@gmail.com +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package pool + +import "sync" + +var ( + bufPool5k sync.Pool + bufPool2k sync.Pool + bufPool1k sync.Pool + bufPool sync.Pool +) + +func GetBuf(size int) []byte { + var x interface{} + if size >= 5*1024 { + x = bufPool5k.Get() + } else if size >= 2*1024 { + x = bufPool2k.Get() + } else if size >= 1*1024 { + x = bufPool1k.Get() + } else { + x = bufPool.Get() + } + if x == nil { + return make([]byte, size) + } + buf := x.([]byte) + if cap(buf) < size { + return make([]byte, size) + } + return buf[:size] +} + +func PutBuf(buf []byte) { + size := cap(buf) + if size >= 5*1024 { + bufPool5k.Put(buf) + } else if size >= 2*1024 { + bufPool2k.Put(buf) + } else if size >= 1*1024 { + bufPool1k.Put(buf) + } else { + bufPool.Put(buf) + } +} diff --git a/src/frp/utils/vhost/http.go b/src/frp/utils/vhost/http.go index 4bc720c..aa66b6a 100644 --- a/src/frp/utils/vhost/http.go +++ b/src/frp/utils/vhost/http.go @@ -26,6 +26,7 @@ import ( "time" "frp/utils/conn" + "frp/utils/pool" ) type HttpMuxer struct { @@ -61,9 +62,11 @@ func HttpHostNameRewrite(c *conn.Conn, rewriteHost string) (_ net.Conn, err erro } func hostNameRewrite(request io.Reader, rewriteHost string) (_ []byte, err error) { - buffer := make([]byte, 1024) - request.Read(buffer) - retBuffer, err := parseRequest(buffer, rewriteHost) + buf := pool.GetBuf(1024) + defer pool.PutBuf(buf) + + request.Read(buf) + retBuffer, err := parseRequest(buf, rewriteHost) return retBuffer, err } diff --git a/src/frp/utils/vhost/https.go b/src/frp/utils/vhost/https.go index eedfab3..7f96365 100644 --- a/src/frp/utils/vhost/https.go +++ b/src/frp/utils/vhost/https.go @@ -22,6 +22,7 @@ import ( "time" "frp/utils/conn" + "frp/utils/pool" ) const ( @@ -52,7 +53,9 @@ func NewHttpsMuxer(listener *conn.Listener, timeout time.Duration) (*HttpsMuxer, } func readHandshake(rd io.Reader) (host string, err error) { - data := make([]byte, 1024) + data := pool.GetBuf(1024) + origin := data + defer pool.PutBuf(origin) length, err := rd.Read(data) if err != nil { return