加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
此仓库是为了提升国内下载速度的镜像仓库,每日同步一次。 原始仓库: https://github.com/buger/goreplay
克隆/下载
tcp_client.go 3.95 KB
一键复制 编辑 原始数据 按行查看 历史
Dima Golomozy 提交于 2023-01-12 17:24 . goreplay-cli package (#1148)
package goreplay
import (
"crypto/tls"
"io"
"net"
"runtime/debug"
"syscall"
"time"
)
// TCPClientConfig client configuration
type TCPClientConfig struct {
Debug bool
ConnectionTimeout time.Duration
Timeout time.Duration
ResponseBufferSize int
Secure bool
}
// TCPClient client connection properties
type TCPClient struct {
baseURL string
addr string
conn net.Conn
respBuf []byte
config *TCPClientConfig
redirectsCount int
}
// NewTCPClient returns new TCPClient
func NewTCPClient(addr string, config *TCPClientConfig) *TCPClient {
if config.Timeout.Nanoseconds() == 0 {
config.Timeout = 5 * time.Second
}
config.ConnectionTimeout = config.Timeout
if config.ResponseBufferSize == 0 {
config.ResponseBufferSize = 100 * 1024 // 100kb
}
client := &TCPClient{config: config, addr: addr}
client.respBuf = make([]byte, config.ResponseBufferSize)
return client
}
// Connect creates a tcp connection of the client
func (c *TCPClient) Connect() (err error) {
c.Disconnect()
c.conn, err = net.DialTimeout("tcp", c.addr, c.config.ConnectionTimeout)
if c.config.Secure {
tlsConn := tls.Client(c.conn, &tls.Config{InsecureSkipVerify: true})
if err = tlsConn.Handshake(); err != nil {
return
}
c.conn = tlsConn
}
return
}
// Disconnect closes the client connection
func (c *TCPClient) Disconnect() {
if c.conn != nil {
c.conn.Close()
c.conn = nil
Debug(1, "[TCPClient] Disconnected: ", c.baseURL)
}
}
func (c *TCPClient) isAlive() bool {
one := make([]byte, 1)
// Ready 1 byte from socket without timeout to check if it not closed
c.conn.SetReadDeadline(time.Now().Add(time.Millisecond))
_, err := c.conn.Read(one)
if err == nil {
return true
} else if err == io.EOF {
Debug(1, "[TCPClient] connection closed, reconnecting")
return false
} else if err == syscall.EPIPE {
Debug(1, "Detected broken pipe.", err)
return false
}
return true
}
// Send sends data over created tcp connection
func (c *TCPClient) Send(data []byte) (response []byte, err error) {
// Don't exit on panic
defer func() {
if r := recover(); r != nil {
Debug(1, "[TCPClient]", r, string(data))
if _, ok := r.(error); !ok {
Debug(1, "[TCPClient] Failed to send request: ", string(data))
Debug(1, "PANIC: pkg:", r, debug.Stack())
}
}
}()
if c.conn == nil || !c.isAlive() {
Debug(1, "[TCPClient] Connecting:", c.baseURL)
if err = c.Connect(); err != nil {
Debug(1, "[TCPClient] Connection error:", err)
return
}
}
timeout := time.Now().Add(c.config.Timeout)
c.conn.SetWriteDeadline(timeout)
if c.config.Debug {
Debug(1, "[TCPClient] Sending:", string(data))
}
if _, err = c.conn.Write(data); err != nil {
Debug(1, "[TCPClient] Write error:", err, c.baseURL)
return
}
var readBytes, n int
var currentChunk []byte
timeout = time.Now().Add(c.config.Timeout)
for {
c.conn.SetReadDeadline(timeout)
if readBytes < len(c.respBuf) {
n, err = c.conn.Read(c.respBuf[readBytes:])
readBytes += n
if err != nil {
if err == io.EOF {
err = nil
}
break
}
} else {
if currentChunk == nil {
currentChunk = make([]byte, readChunkSize)
}
n, err = c.conn.Read(currentChunk)
if err == io.EOF {
break
} else if err != nil {
Debug(1, "[TCPClient] Read the whole body error:", err, c.baseURL)
break
}
readBytes += int(n)
}
if readBytes >= maxResponseSize {
Debug(1, "[TCPClient] Body is more than the max size", maxResponseSize,
c.baseURL)
break
}
// For following chunks expect less timeout
timeout = time.Now().Add(c.config.Timeout / 5)
}
if err != nil {
Debug(1, "[TCPClient] Response read error", err, c.conn, readBytes)
return
}
if readBytes > len(c.respBuf) {
readBytes = len(c.respBuf)
}
payload := make([]byte, readBytes)
copy(payload, c.respBuf[:readBytes])
if c.config.Debug {
Debug(1, "[TCPClient] Received:", string(payload))
}
return payload, err
}
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化