goreplay/tcp_client.go

package main

import (
	"crypto/tls"
	"io"
	"log"
	"net"
	"runtime/debug"
	"syscall"
	"time"
)

type TCPClientConfig struct {
	Debug              bool
	ConnectionTimeout  time.Duration
	Timeout            time.Duration
	ResponseBufferSize int
	Secure             bool
}

type TCPClient struct {
	baseURL        string
	addr           string
	conn           net.Conn
	respBuf        []byte
	config         *TCPClientConfig
	redirectsCount int
}

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
}

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
}

func (c *TCPClient) Disconnect() {
	if c.conn != nil {
		c.conn.Close()
		c.conn = nil
		Debug("[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 {
		if c.config.Debug {
			Debug("[TCPClient] connection closed, reconnecting")
		}
		return false
	} else if err == syscall.EPIPE {
		Debug("Detected broken pipe.", err)
		return false
	}

	return true
}

func (c *TCPClient) Send(data []byte) (response []byte, err error) {
	// Don't exit on panic
	defer func() {
		if r := recover(); r != nil {
			Debug("[TCPClient]", r, string(data))

			if _, ok := r.(error); !ok {
				log.Println("[TCPClient] Failed to send request: ", string(data))
				log.Println("PANIC: pkg:", r, debug.Stack())
			}
		}
	}()

	if c.conn == nil || !c.isAlive() {
		Debug("[TCPClient] Connecting:", c.baseURL)
		if err = c.Connect(); err != nil {
			log.Println("[TCPClient] Connection error:", err)
			return
		}
	}

	timeout := time.Now().Add(c.config.Timeout)

	c.conn.SetWriteDeadline(timeout)

	if c.config.Debug {
		Debug("[TCPClient] Sending:", string(data))
	}

	if _, err = c.conn.Write(data); err != nil {
		Debug("[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("[TCPClient] Read the whole body error:", err, c.baseURL)
				break
			}

			readBytes += int(n)
		}

		if readBytes >= maxResponseSize {
			Debug("[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("[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("[TCPClient] Received:", string(payload))
	}

	return payload, err
}