diff --git a/emitter.go b/emitter.go index 49db03a..0b2080a 100644 --- a/emitter.go +++ b/emitter.go @@ -2,9 +2,9 @@ package main import ( "bytes" + "hash/fnv" "io" "time" - "hash/fnv" ) // Start initialize loop for sending data from inputs to outputs diff --git a/emitter_test.go b/emitter_test.go index 6991d71..e9fc4bb 100644 --- a/emitter_test.go +++ b/emitter_test.go @@ -1,11 +1,11 @@ package main import ( + "bytes" "io" "sync" "sync/atomic" "testing" - "bytes" ) func TestEmitter(t *testing.T) { diff --git a/http_client.go b/http_client.go index aeb322f..f78a819 100644 --- a/http_client.go +++ b/http_client.go @@ -61,7 +61,7 @@ func NewHTTPClient(baseURL string, config *HTTPClientConfig) *HTTPClient { if config.Timeout == 0 { config.Timeout = time.Second - } + } config.ConnectionTimeout = config.Timeout @@ -87,7 +87,7 @@ func (c *HTTPClient) Connect() (err error) { c.Disconnect() if !strings.Contains(c.host, ":") { - c.conn, err = net.DialTimeout("tcp", c.host + ":" + defaultPorts[c.scheme], c.config.ConnectionTimeout) + c.conn, err = net.DialTimeout("tcp", c.host+":"+defaultPorts[c.scheme], c.config.ConnectionTimeout) } else { c.conn, err = net.DialTimeout("tcp", c.host, c.config.ConnectionTimeout) } diff --git a/input_file_test.go b/input_file_test.go index 61e9212..6e514ab 100644 --- a/input_file_test.go +++ b/input_file_test.go @@ -4,6 +4,8 @@ import ( "bytes" "errors" "fmt" + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/s3" "io" "io/ioutil" "log" @@ -379,4 +381,50 @@ func Duplicate(data []byte) (duplicate []byte) { copy(duplicate, data) return -} \ No newline at end of file +} + +func TestInputFileFromS3(t *testing.T) { + rnd := rand.Int63() + path := fmt.Sprintf("s3://test-gor/%d/requests.gz", rnd) + + output := NewS3Output(path, + &S3OutputConfig{ + bufferConfig: FileOutputConfig{queueLimit: 5000}, + }, + ) + output.closeC = make(chan struct{}, 2) + + for i := 0; i <= 10000; i++ { + output.Write([]byte("1 1 1\ntest")) + + if i%5000 == 0 { + output.buffer.updateName() + } + } + output.Write([]byte("1 1 1\ntest")) + + for i := 0; i < 2; i++ { + <-output.closeC + } + + input := NewFileInput(fmt.Sprintf("s3://test-gor/%d", rnd), false) + + buf := make([]byte, 1000) + for i := 0; i <= 10000; i++ { + input.Read(buf) + } + + // Cleanup artifacts + bucket := aws.String("test-gor") + svc := s3.New(output.session) + params := &s3.ListObjectsInput{ + Bucket: bucket, + Prefix: aws.String(fmt.Sprintf("%d", rnd)), + } + + resp, _ := svc.ListObjects(params) + + for _, c := range resp.Contents { + svc.DeleteObject(&s3.DeleteObjectInput{Bucket: bucket, Key: c.Key}) + } +} diff --git a/input_s3.go b/input_s3.go deleted file mode 100644 index 9f73ba5..0000000 --- a/input_s3.go +++ /dev/null @@ -1,12 +0,0 @@ -package main - -import ( -) - -type S3InputConfig struct { - // bufferConfig FileInputConfig - - bufferPath string - region string - endpoint string -} diff --git a/output_binary.go b/output_binary.go index 6134e65..4e55217 100644 --- a/output_binary.go +++ b/output_binary.go @@ -1,175 +1,175 @@ package main import ( - "io" - "sync/atomic" - "time" + "io" + "sync/atomic" + "time" ) // BinaryOutputConfig struct for holding binary output configuration type BinaryOutputConfig struct { - workers int - Timeout time.Duration - BufferSize int - Debug bool - TrackResponses bool + workers int + Timeout time.Duration + BufferSize int + Debug bool + TrackResponses bool } // BinaryOutput plugin manage pool of workers which send request to replayed server // By default workers pool is dynamic and starts with 10 workers // You can specify fixed number of workers using `--output-tcp-workers` type BinaryOutput struct { - // Keep this as first element of struct because it guarantees 64bit - // alignment. atomic.* functions crash on 32bit machines if operand is not - // aligned at 64bit. See https://github.com/golang/go/issues/599 - activeWorkers int64 + // Keep this as first element of struct because it guarantees 64bit + // alignment. atomic.* functions crash on 32bit machines if operand is not + // aligned at 64bit. See https://github.com/golang/go/issues/599 + activeWorkers int64 - address string - queue chan []byte + address string + queue chan []byte - responses chan response + responses chan response - needWorker chan int + needWorker chan int - config *BinaryOutputConfig + config *BinaryOutputConfig - queueStats *GorStat + queueStats *GorStat } // NewBinaryOutput constructor for BinaryOutput // Initialize workers func NewBinaryOutput(address string, config *BinaryOutputConfig) io.Writer { - o := new(BinaryOutput) + o := new(BinaryOutput) - o.address = address - o.config = config + o.address = address + o.config = config - o.queue = make(chan []byte, 1000) - o.responses = make(chan response, 1000) - o.needWorker = make(chan int, 1) + o.queue = make(chan []byte, 1000) + o.responses = make(chan response, 1000) + o.needWorker = make(chan int, 1) - // Initial workers count - if o.config.workers == 0 { - o.needWorker <- initialDynamicWorkers - } else { - o.needWorker <- o.config.workers - } + // Initial workers count + if o.config.workers == 0 { + o.needWorker <- initialDynamicWorkers + } else { + o.needWorker <- o.config.workers + } - if len(Settings.middleware) > 0 { - o.config.TrackResponses = true - } + if len(Settings.middleware) > 0 { + o.config.TrackResponses = true + } - go o.workerMaster() + go o.workerMaster() - return o + return o } func (o *BinaryOutput) workerMaster() { - for { - newWorkers := <-o.needWorker - for i := 0; i < newWorkers; i++ { - go o.startWorker() - } + for { + newWorkers := <-o.needWorker + for i := 0; i < newWorkers; i++ { + go o.startWorker() + } - // Disable dynamic scaling if workers poll fixed size - if o.config.workers != 0 { - return - } - } + // Disable dynamic scaling if workers poll fixed size + if o.config.workers != 0 { + return + } + } } func (o *BinaryOutput) startWorker() { - client := NewTCPClient(o.address, &TCPClientConfig{ - Debug: o.config.Debug, - Timeout: o.config.Timeout, - ResponseBufferSize: o.config.BufferSize, - }) + client := NewTCPClient(o.address, &TCPClientConfig{ + Debug: o.config.Debug, + Timeout: o.config.Timeout, + ResponseBufferSize: o.config.BufferSize, + }) - deathCount := 0 + deathCount := 0 - atomic.AddInt64(&o.activeWorkers, 1) + atomic.AddInt64(&o.activeWorkers, 1) - for { - select { - case data := <-o.queue: - o.sendRequest(client, data) - deathCount = 0 - case <-time.After(time.Millisecond * 100): - // When dynamic scaling enabled workers die after 2s of inactivity - if o.config.workers == 0 { - deathCount++ - } else { - continue - } + for { + select { + case data := <-o.queue: + o.sendRequest(client, data) + deathCount = 0 + case <-time.After(time.Millisecond * 100): + // When dynamic scaling enabled workers die after 2s of inactivity + if o.config.workers == 0 { + deathCount++ + } else { + continue + } - if deathCount > 20 { - workersCount := atomic.LoadInt64(&o.activeWorkers) + if deathCount > 20 { + workersCount := atomic.LoadInt64(&o.activeWorkers) - // At least 1 startWorker should be alive - if workersCount != 1 { - atomic.AddInt64(&o.activeWorkers, -1) - return - } - } - } - } + // At least 1 startWorker should be alive + if workersCount != 1 { + atomic.AddInt64(&o.activeWorkers, -1) + return + } + } + } + } } func (o *BinaryOutput) Write(data []byte) (n int, err error) { - if !isRequestPayload(data) { - return len(data), nil - } + if !isRequestPayload(data) { + return len(data), nil + } - buf := make([]byte, len(data)) - copy(buf, data) + buf := make([]byte, len(data)) + copy(buf, data) - o.queue <- buf + o.queue <- buf - if o.config.workers == 0 { - workersCount := atomic.LoadInt64(&o.activeWorkers) + if o.config.workers == 0 { + workersCount := atomic.LoadInt64(&o.activeWorkers) - if len(o.queue) > int(workersCount) { - o.needWorker <- len(o.queue) - } - } + if len(o.queue) > int(workersCount) { + o.needWorker <- len(o.queue) + } + } - return len(data), nil + return len(data), nil } func (o *BinaryOutput) Read(data []byte) (int, error) { - resp := <-o.responses + resp := <-o.responses - Debug("[OUTPUT-TCP] Received response:", string(resp.payload)) + Debug("[OUTPUT-TCP] Received response:", string(resp.payload)) - header := payloadHeader(ReplayedResponsePayload, resp.uuid, resp.startedAt, resp.roundTripTime) - copy(data[0:len(header)], header) - copy(data[len(header):], resp.payload) + header := payloadHeader(ReplayedResponsePayload, resp.uuid, resp.startedAt, resp.roundTripTime) + copy(data[0:len(header)], header) + copy(data[len(header):], resp.payload) - return len(resp.payload) + len(header), nil + return len(resp.payload) + len(header), nil } func (o *BinaryOutput) sendRequest(client *TCPClient, request []byte) { - meta := payloadMeta(request) - if len(meta) < 2 { - return - } - uuid := meta[1] + meta := payloadMeta(request) + if len(meta) < 2 { + return + } + uuid := meta[1] - body := payloadBody(request) + body := payloadBody(request) - start := time.Now() - resp, err := client.Send(body) - stop := time.Now() + start := time.Now() + resp, err := client.Send(body) + stop := time.Now() - if err != nil { - Debug("Request error:", err) - } + if err != nil { + Debug("Request error:", err) + } - if o.config.TrackResponses { - o.responses <- response{resp, uuid, start.UnixNano(), stop.UnixNano() - start.UnixNano()} - } + if o.config.TrackResponses { + o.responses <- response{resp, uuid, start.UnixNano(), stop.UnixNano() - start.UnixNano()} + } } func (o *BinaryOutput) String() string { - return "TCP output: " + o.address + return "TCP output: " + o.address } diff --git a/output_http.go b/output_http.go index 161661c..50cf8a1 100644 --- a/output_http.go +++ b/output_http.go @@ -1,10 +1,10 @@ package main import ( + "fmt" "io" "sync/atomic" "time" - "fmt" "github.com/buger/gor-pro/proto" ) @@ -14,11 +14,11 @@ var _ = fmt.Println const initialDynamicWorkers = 10 type httpWorker struct { - output *HTTPOutput - client *HTTPClient + output *HTTPOutput + client *HTTPClient lastActivity time.Time - queue chan []byte - stop chan bool + queue chan []byte + stop chan bool } func newHTTPWorker(output *HTTPOutput, queue chan []byte) *httpWorker { @@ -38,12 +38,12 @@ func newHTTPWorker(output *HTTPOutput, queue chan []byte) *httpWorker { } w.stop = make(chan bool) - go func(){ + go func() { for { select { case payload := <-w.queue: output.sendRequest(client, payload) - case <- w.stop: + case <-w.stop: return } } @@ -182,7 +182,7 @@ func (o *HTTPOutput) sessionWorkerMaster() { now := time.Now() for id, w := range o.workerSessions { - if !w.lastActivity.IsZero() && now.Sub(w.lastActivity) >= 60 * time.Second { + if !w.lastActivity.IsZero() && now.Sub(w.lastActivity) >= 60*time.Second { w.stop <- true delete(o.workerSessions, id) atomic.AddInt64(&o.activeWorkers, -1) diff --git a/output_http_test.go b/output_http_test.go index 154e5c5..0b15e71 100644 --- a/output_http_test.go +++ b/output_http_test.go @@ -151,7 +151,6 @@ func TestHTTPOutputSessions(t *testing.T) { uuid1 := []byte("1234567890123456789a0000") uuid2 := []byte("1234567890123456789d0000") - for i := 0; i < 100; i++ { wg.Add(1) // OPTIONS should be ignored copy(uuid1[20:], randByte(4)) diff --git a/output_s3.go b/output_s3.go index 752eff1..1117480 100644 --- a/output_s3.go +++ b/output_s3.go @@ -30,6 +30,7 @@ type S3Output struct { buffer *FileOutput session *session.Session config *S3OutputConfig + closeC chan struct{} } // NewFileOutput constructor for FileOutput, accepts path @@ -84,13 +85,19 @@ func (o *S3Output) Close() { o.buffer.Close() } -func (o *S3Output) keyPath(idx int) (bucket, key string) { - path := o.pathTemplate[5:] // stripping `s3://` +func parseS3Url(path string) (bucket, key string) { + path = path[5:] // stripping `s3://` sep := strings.IndexByte(path, '/') bucket = path[:sep] key = path[sep+1:] + return bucket, key +} + +func (o *S3Output) keyPath(idx int) (bucket, key string) { + bucket, key = parseS3Url(o.pathTemplate) + for name, fn := range dateFileNameFuncs { key = strings.Replace(key, name, fn(), -1) } @@ -119,4 +126,8 @@ func (o *S3Output) onBufferUpdate(path string) { } os.Remove(path) + + if o.closeC != nil { + o.closeC <- struct{}{} + } } diff --git a/output_s3_test.go b/output_s3_test.go index 15693b5..748a083 100644 --- a/output_s3_test.go +++ b/output_s3_test.go @@ -58,3 +58,55 @@ func TestS3Output(t *testing.T) { os.Remove(m) } } + +func TestS3OutputQueueLimit(t *testing.T) { + bucket := aws.String("test-gor") + rnd := rand.Int63() + path := fmt.Sprintf("s3://test-gor/%d/requests.gz", rnd) + + output := NewS3Output(path, + &S3OutputConfig{ + bufferConfig: FileOutputConfig{queueLimit: 100}, + }, + ) + output.closeC = make(chan struct{}, 3) + + svc := s3.New(output.session) + + for i := 0; i < 3; i++ { + for i := 0; i < 100; i++ { + output.Write([]byte("1 1 1\ntest")) + } + output.buffer.updateName() + } + output.buffer.updateName() + output.Write([]byte("1 1 1\ntest")) + + for i := 0; i < 3; i++ { + <-output.closeC + } + + params := &s3.ListObjectsInput{ + Bucket: bucket, + Prefix: aws.String(fmt.Sprintf("%d", rnd)), + } + + resp, _ := svc.ListObjects(params) + if len(resp.Contents) != 3 { + t.Error("Should create 3 object", len(resp.Contents)) + } else { + if *resp.Contents[0].Key != fmt.Sprintf("%d/requests_0.gz", rnd) || + *resp.Contents[1].Key != fmt.Sprintf("%d/requests_1.gz", rnd) { + t.Error("Should assign proper names", resp.Contents) + } + } + + for _, c := range resp.Contents { + svc.DeleteObject(&s3.DeleteObjectInput{Bucket: bucket, Key: c.Key}) + } + + matches, _ := filepath.Glob(fmt.Sprintf("/tmp/gor_output_s3_*")) + for _, m := range matches { + os.Remove(m) + } +} diff --git a/protocol.go b/protocol.go index 2307e1b..3bf47e0 100644 --- a/protocol.go +++ b/protocol.go @@ -14,7 +14,7 @@ const ( ) func randByte(len int) []byte { - b := make([]byte, len / 2) + b := make([]byte, len/2) rand.Read(b) h := make([]byte, len) @@ -100,7 +100,7 @@ func payloadID(payload []byte) []byte { return []byte{} } - return payload[2: 2 + idx] + return payload[2 : 2+idx] } func isOriginPayload(payload []byte) bool { diff --git a/raw_socket_listener/listener.go b/raw_socket_listener/listener.go index a851fb5..8ef1a3f 100644 --- a/raw_socket_listener/listener.go +++ b/raw_socket_listener/listener.go @@ -34,9 +34,9 @@ import ( var _ = fmt.Println type packet struct { - srcIP []byte - data []byte - timestamp time.Time + srcIP []byte + data []byte + timestamp time.Time } // Listener handle traffic capture @@ -641,9 +641,9 @@ func (t *Listener) buildPacket(packetSrcIP []byte, packetData []byte, timestamp copy(copyPacketData, packetSrcIP) return &packet{ - srcIP: packetSrcIP, - data: packetData, - timestamp:timestamp, + srcIP: packetSrcIP, + data: packetData, + timestamp: timestamp, } } diff --git a/raw_socket_listener/tcp_message.go b/raw_socket_listener/tcp_message.go index f1bbc1b..3a55864 100644 --- a/raw_socket_listener/tcp_message.go +++ b/raw_socket_listener/tcp_message.go @@ -17,8 +17,8 @@ var _ = log.Println type TCPProtocol uint8 const ( - ProtocolHTTP TCPProtocol = 0 - ProtocolBinary TCPProtocol = 1 + ProtocolHTTP TCPProtocol = 0 + ProtocolBinary TCPProtocol = 1 ) // TCPMessage ensure that all TCP packets for given request is received, and processed in right sequence @@ -365,7 +365,7 @@ func (t *TCPMessage) updateBodyType() { var lengthB, encB, connB []byte - proto.ParseHeaders(t.packetsData(), func(header, value []byte)bool{ + proto.ParseHeaders(t.packetsData(), func(header, value []byte) bool { if proto.HeadersEqual(header, []byte("Content-Length")) { lengthB = value return false @@ -446,7 +446,7 @@ func (t *TCPMessage) check100Continue() { } var expectB []byte - proto.ParseHeaders(t.packetsData(), func(header, value []byte)bool{ + proto.ParseHeaders(t.packetsData(), func(header, value []byte) bool { if proto.HeadersEqual(header, bExpectHeader) { expectB = value return false diff --git a/raw_socket_listener/tcp_message_test.go b/raw_socket_listener/tcp_message_test.go index 14c4e6e..e05f058 100644 --- a/raw_socket_listener/tcp_message_test.go +++ b/raw_socket_listener/tcp_message_test.go @@ -255,4 +255,3 @@ func TestTcpMessageStart(t *testing.T) { t.Error("Message timestamp should be equal to the lowest related packet timestamp", start, msg.Start) } } - diff --git a/raw_socket_listener/tcp_packet.go b/raw_socket_listener/tcp_packet.go index f2ba82c..9427a3c 100644 --- a/raw_socket_listener/tcp_packet.go +++ b/raw_socket_listener/tcp_packet.go @@ -36,11 +36,11 @@ type TCPPacket struct { DataOffset uint8 IsFIN bool - Raw []byte - Data []byte - Addr []byte + Raw []byte + Data []byte + Addr []byte timestamp time.Time - ID tcpID + ID tcpID } // ParseTCPPacket takes address and tcp payload and returns parsed TCPPacket @@ -85,7 +85,7 @@ func (t *TCPPacket) ParseBasic() { func (t *TCPPacket) dump() *packet { packetSrcIP := make([]byte, 16) - packetData := make([]byte, len(t.Data) + 16) + packetData := make([]byte, len(t.Data)+16) copy(packetSrcIP, t.Addr) @@ -104,9 +104,9 @@ func (t *TCPPacket) dump() *packet { copy(packetData[16:], t.Data) return &packet{ - srcIP: packetSrcIP, - data:packetData, - timestamp:t.timestamp, + srcIP: packetSrcIP, + data: packetData, + timestamp: t.timestamp, } } diff --git a/settings.go b/settings.go index f6be3c7..76d5faa 100644 --- a/settings.go +++ b/settings.go @@ -30,7 +30,7 @@ type AppSettings struct { stats bool exitAfter time.Duration - splitOutput bool + splitOutput bool recognizeTCPSessions bool inputDummy MultiOption @@ -58,15 +58,15 @@ type AppSettings struct { middleware string - inputHTTP MultiOption + inputHTTP MultiOption - outputHTTP MultiOption + outputHTTP MultiOption outputHTTPConfig HTTPOutputConfig - outputBinary MultiOption + outputBinary MultiOption outputBinaryConfig BinaryOutputConfig - modifierConfig HTTPModifierConfig + modifierConfig HTTPModifierConfig outputKafkaConfig KafkaConfig } @@ -136,7 +136,7 @@ func init() { flag.Var(&Settings.outputHTTP, "output-http", "Forwards incoming requests to given http address.\n\t# Redirect all incoming requests to staging.com address \n\tgor --input-raw :80 --output-http http://staging.com") - /* outputHTTPConfig */ + /* outputHTTPConfig */ flag.IntVar(&Settings.outputHTTPConfig.BufferSize, "output-http-response-buffer", 0, "HTTP response buffer size, all data after this size will be discarded.") flag.IntVar(&Settings.outputHTTPConfig.workers, "output-http-workers", 0, "Gor uses dynamic worker scaling by default. Enter a number to run a set number of workers.") flag.IntVar(&Settings.outputHTTPConfig.redirectLimit, "output-http-redirects", 0, "Enable how often redirects should be followed.") @@ -148,7 +148,6 @@ func init() { flag.StringVar(&Settings.outputHTTPConfig.elasticSearch, "output-http-elasticsearch", "", "Send request and response stats to ElasticSearch:\n\tgor --input-raw :8080 --output-http staging.com --output-http-elasticsearch 'es_host:api_port/index_name'") /* outputHTTPConfig */ - flag.Var(&Settings.outputBinary, "output-binary", "Forwards incoming binary payloads to given address.\n\t# Redirect all incoming requests to staging.com address \n\tgor --input-raw :80 --input-raw-protocol binary --output-binary staging.com:80") /* outputBinaryConfig */ flag.IntVar(&Settings.outputBinaryConfig.BufferSize, "output-tcp-response-buffer", 0, "TCP response buffer size, all data after this size will be discarded.") diff --git a/tcp_client.go b/tcp_client.go index 677c2d4..8235b9a 100644 --- a/tcp_client.go +++ b/tcp_client.go @@ -1,190 +1,190 @@ package main import ( - "crypto/tls" - "io" - "log" - "net" - "runtime/debug" - "syscall" - "time" + "crypto/tls" + "io" + "log" + "net" + "runtime/debug" + "syscall" + "time" ) type TCPClientConfig struct { - Debug bool - ConnectionTimeout time.Duration - Timeout time.Duration - ResponseBufferSize int - Secure bool + 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 + 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 - } + if config.Timeout.Nanoseconds() == 0 { + config.Timeout = 5 * time.Second + } - config.ConnectionTimeout = config.Timeout + config.ConnectionTimeout = config.Timeout - if config.ResponseBufferSize == 0 { - config.ResponseBufferSize = 100 * 1024 // 100kb - } + if config.ResponseBufferSize == 0 { + config.ResponseBufferSize = 100 * 1024 // 100kb + } - client := &TCPClient{config: config, addr: addr} - client.respBuf = make([]byte, config.ResponseBufferSize) + client := &TCPClient{config: config, addr: addr} + client.respBuf = make([]byte, config.ResponseBufferSize) - return client + return client } func (c *TCPClient) Connect() (err error) { - c.Disconnect() + c.Disconnect() - c.conn, err = net.DialTimeout("tcp", c.addr, c.config.ConnectionTimeout) + 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 c.config.Secure { + tlsConn := tls.Client(c.conn, &tls.Config{InsecureSkipVerify: true}) - if err = tlsConn.Handshake(); err != nil { - return - } + if err = tlsConn.Handshake(); err != nil { + return + } - c.conn = tlsConn - } + c.conn = tlsConn + } - return + return } func (c *TCPClient) Disconnect() { - if c.conn != nil { - c.conn.Close() - c.conn = nil - Debug("[TCPClient] Disconnected: ", c.baseURL) - } + if c.conn != nil { + c.conn.Close() + c.conn = nil + Debug("[TCPClient] Disconnected: ", c.baseURL) + } } func (c *TCPClient) isAlive() bool { - one := make([]byte, 1) + 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) + // 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 - } + 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 + 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)) + // 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 _, 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 - } - } + 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) + timeout := time.Now().Add(c.config.Timeout) - c.conn.SetWriteDeadline(timeout) + c.conn.SetWriteDeadline(timeout) - if c.config.Debug { - Debug("[TCPClient] Sending:", string(data)) - } + 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 - } + 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) + var readBytes, n int + var currentChunk []byte + timeout = time.Now().Add(c.config.Timeout) - for { - c.conn.SetReadDeadline(timeout) + for { + c.conn.SetReadDeadline(timeout) - if readBytes < len(c.respBuf) { - n, err = c.conn.Read(c.respBuf[readBytes:]) - readBytes += n + 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) - } + if err != nil { + if err == io.EOF { + err = nil + } + break + } + } else { + if currentChunk == nil { + currentChunk = make([]byte, readChunkSize) + } - n, err = c.conn.Read(currentChunk) + 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 - } + if err == io.EOF { + break + } else if err != nil { + Debug("[TCPClient] Read the whole body error:", err, c.baseURL) + break + } - readBytes += int(n) - } + readBytes += int(n) + } - if readBytes >= maxResponseSize { - Debug("[TCPClient] Body is more than the max size", maxResponseSize, - c.baseURL) - break - } + 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) - } + // 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 err != nil { + Debug("[TCPClient] Response read error", err, c.conn, readBytes) + return + } - if readBytes > len(c.respBuf) { - readBytes = len(c.respBuf) - } + if readBytes > len(c.respBuf) { + readBytes = len(c.respBuf) + } - payload := make([]byte, readBytes) - copy(payload, c.respBuf[:readBytes]) + payload := make([]byte, readBytes) + copy(payload, c.respBuf[:readBytes]) - if c.config.Debug { - Debug("[TCPClient] Received:", string(payload)) - } + if c.config.Debug { + Debug("[TCPClient] Received:", string(payload)) + } - return payload, err + return payload, err } diff --git a/test_input.go b/test_input.go index b903359..8da3cc8 100644 --- a/test_input.go +++ b/test_input.go @@ -8,7 +8,7 @@ import ( // TestInput used for testing purpose, it allows emitting requests on demand type TestInput struct { - data chan []byte + data chan []byte disableHeaders bool } @@ -45,7 +45,6 @@ func (i *TestInput) EmitGET() { i.data <- []byte("GET / HTTP/1.1\r\n\r\n") } - // EmitPOST emits POST request with Content-Length func (i *TestInput) EmitPOST() { i.data <- []byte("POST /pub/WWW/ HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")