优化代码

This commit is contained in:
GoEdgeLab
2024-04-15 12:20:31 +08:00
parent 184a0cef53
commit 62d7ae55b6
3 changed files with 19 additions and 11 deletions

View File

@@ -16,6 +16,7 @@ import (
"github.com/TeaOSLab/EdgeNode/internal/goman"
"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
"github.com/TeaOSLab/EdgeNode/internal/rpc"
"github.com/TeaOSLab/EdgeNode/internal/utils"
connutils "github.com/TeaOSLab/EdgeNode/internal/utils/conns"
"github.com/iwind/TeaGo/Tea"
"io"
@@ -245,7 +246,9 @@ func (this *HTTPCacheTaskManager) fetchKey(key *pb.HTTPCacheTaskKey) error {
}
// 读取内容,以便于生成缓存
_, err = io.Copy(io.Discard, resp.Body)
var buf = utils.BytePool16k.Get()
_, err = io.CopyBuffer(io.Discard, resp.Body, buf.Bytes)
utils.BytePool16k.Put(buf)
if err != nil {
if err != io.EOF {
err = this.simplifyErr(err)

View File

@@ -128,8 +128,8 @@ func (this *HTTPRequest) doWebsocket(requestHost string, isLastRetry bool) (shou
go func() {
// 读取第一个响应
var respReader = NewWebsocketResponseReader(originConn)
resp, err := http.ReadResponse(bufio.NewReader(respReader), this.RawReq)
if err != nil || resp == nil {
resp, respErr := http.ReadResponse(bufio.NewReader(respReader), this.RawReq)
if respErr != nil || resp == nil {
if resp != nil && resp.Body != nil {
_ = resp.Body.Close()
}
@@ -160,8 +160,8 @@ func (this *HTTPRequest) doWebsocket(requestHost string, isLastRetry bool) (shou
if headerIndex > 0 {
var leftBytes = headerBytes[headerIndex+4:]
if len(leftBytes) > 0 {
_, err = clientConn.Write(leftBytes)
if err != nil {
_, writeErr := clientConn.Write(leftBytes)
if writeErr != nil {
if resp.Body != nil {
_ = resp.Body.Close()
}
@@ -181,22 +181,25 @@ func (this *HTTPRequest) doWebsocket(requestHost string, isLastRetry bool) (shou
var buf = utils.BytePool4k.Get()
defer utils.BytePool4k.Put(buf)
for {
n, err := originConn.Read(buf.Bytes)
n, readErr := originConn.Read(buf.Bytes)
if n > 0 {
this.writer.sentBodyBytes += int64(n)
_, err = clientConn.Write(buf.Bytes[:n])
if err != nil {
_, writeErr := clientConn.Write(buf.Bytes[:n])
if writeErr != nil {
break
}
}
if err != nil {
if readErr != nil {
break
}
}
_ = clientConn.Close()
_ = originConn.Close()
}()
_, _ = io.Copy(originConn, clientConn)
var buf = utils.BytePool4k.Get()
_, _ = io.CopyBuffer(originConn, clientConn, buf.Bytes)
utils.BytePool4k.Put(buf)
return
}

View File

@@ -119,7 +119,9 @@ func (this *Post307Action) Perform(waf *WAF, group *RuleGroup, set *RuleSet, req
// 清空请求内容
var req = request.WAFRaw()
if req.ContentLength > 0 && req.Body != nil {
_, _ = io.Copy(io.Discard, req.Body)
var buf = utils.BytePool16k.Get()
_, _ = io.CopyBuffer(io.Discard, req.Body, buf.Bytes)
utils.BytePool16k.Put(buf)
_ = req.Body.Close()
}