diff --git a/Makefile b/Makefile index 475e122..12db096 100644 --- a/Makefile +++ b/Makefile @@ -1,31 +1,33 @@ SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go input_tcp.go limiter.go output_dummy.go output_file.go input_http.go output_http.go output_tcp.go plugins.go settings.go settings_header_filters.go settings_header_hash_filters.go settings_headers.go settings_methods.go settings_option.go settings_url_regexp.go test_input.go elasticsearch.go settings_url_map.go +SOURCE_PATH = /gopath/src/github.com/buger/gor/ + release: release-x86 release-x64 release-x64: - docker run -v `pwd`:/gopath/src/gor -t --env GOOS=linux --env GOARCH=amd64 --env CGO_ENABLED=0 -i gor go build && tar -czf gor_x64.tar.gz gor && rm gor + docker run -v `pwd`:$(SOURCE_PATH) -t --env GOOS=linux --env GOARCH=amd64 --env CGO_ENABLED=0 -i gor go build && tar -czf gor_x64.tar.gz gor && rm gor release-x86: - docker run -v `pwd`:/gopath/src/gor -t --env GOOS=linux --env GOARCH=386 --env CGO_ENABLED=0 -i gor go build && tar -czf gor_x86.tar.gz gor && rm gor + docker run -v `pwd`:$(SOURCE_PATH) -t --env GOOS=linux --env GOARCH=386 --env CGO_ENABLED=0 -i gor go build && tar -czf gor_x86.tar.gz gor && rm gor dbuild: docker build -t gor . dtest: - docker run -v `pwd`:/gopath/src/gor -t -i --env GORACE="halt_on_error=1" gor go test $(ARGS) -race -v + docker run -v `pwd`:$(SOURCE_PATH) -t -i --env GORACE="halt_on_error=1" gor go test $(ARGS) -race -v dfmt: - docker run -v `pwd`:/gopath/src/gor -t -i gor go fmt + docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go fmt dvet: - docker run -v `pwd`:/gopath/src/gor -t -i gor go vet + docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go vet dbench: - docker run -v `pwd`:/gopath/src/gor -t -i gor go test -v -run NOT_EXISTING -bench HTTP + docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go test -v -run NOT_EXISTING -bench HTTP # Used mainly for debugging, because docker container do not have access to parent machine ports drun: - docker run -v `pwd`:/gopath/src/gor -t -i gor go run $(SOURCE) --input-dummy=0 --input-http=:9000 --output-http="http://localhost:9000" --verbose + docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go run $(SOURCE) --input-dummy=0 --input-http=:9000 --output-http="http://localhost:9000" --verbose dbash: - docker run -v `pwd`:/gopath/src/gor -t -i gor /bin/bash \ No newline at end of file + docker run -v `pwd`:$(SOURCE_PATH) -t -i gor /bin/bash \ No newline at end of file diff --git a/http_client.go b/http_client.go index ae36b2d..b9b8e13 100644 --- a/http_client.go +++ b/http_client.go @@ -2,9 +2,11 @@ package main import ( "crypto/tls" + "io" "net" "net/url" "strings" + "time" ) var defaultPorts = map[string]string{ @@ -15,13 +17,17 @@ var defaultPorts = map[string]string{ type HTTPClient struct { baseURL *url.URL conn net.Conn - buf []byte + respBuf []byte } func NewHTTPClient(baseURL string) *HTTPClient { + if !strings.HasPrefix(baseURL, "http") { + baseURL = "http://" + baseURL + } + client := new(HTTPClient) client.baseURL, _ = url.Parse(baseURL) - client.buf = make([]byte, 4096*10) + client.respBuf = make([]byte, 4096*10) if !strings.Contains(client.baseURL.Host, ":") { client.baseURL.Host += ":" + defaultPorts[client.baseURL.Scheme] @@ -31,11 +37,17 @@ func NewHTTPClient(baseURL string) *HTTPClient { } func (c *HTTPClient) Connect() (err error) { + c.Disconnect() + c.conn, err = net.Dial("tcp", c.baseURL.Host) if c.baseURL.Scheme == "https" { tlsConn := tls.Client(c.conn, &tls.Config{InsecureSkipVerify: true}) - err = tlsConn.Handshake() + + if err = tlsConn.Handshake(); err != nil { + return + } + c.conn = tlsConn } @@ -43,19 +55,46 @@ func (c *HTTPClient) Connect() (err error) { } func (c *HTTPClient) Disconnect() { - c.conn.Close() - c.conn = nil + if c.conn != nil { + c.conn.Close() + c.conn = nil + Debug("Disconnected: ", c.baseURL) + } +} + +func (c *HTTPClient) 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)) + if _, err := c.conn.Read(one); err == io.EOF { + return false + } + + return true } func (c *HTTPClient) Send(data []byte) (response []byte, err error) { - if c.conn == nil { + if c.conn == nil || !c.isAlive() { + Debug("Connecting:", c.baseURL) c.Connect() } - _, err = c.conn.Write(data) - n, err := c.conn.Read(c.buf) + timeout := time.Now().Add(5 * time.Second) - Debug(string(c.buf[:n])) + c.conn.SetWriteDeadline(timeout) + if _, err = c.conn.Write(data); err != nil { + Debug("Write error:", err, c.baseURL) + return + } - return c.buf[:n], err + c.conn.SetReadDeadline(timeout) + n, err := c.conn.Read(c.respBuf) + + if err != nil { + Debug("READ ERRORR!", err, c.conn) + return + } + + return c.respBuf[:n], err } diff --git a/http_client_test.go b/http_client_test.go index 63504b1..ce369c6 100644 --- a/http_client_test.go +++ b/http_client_test.go @@ -2,11 +2,13 @@ package main import ( "io/ioutil" + "net" "net/http" "net/http/httptest" "net/http/httputil" "sync" "testing" + _ "time" ) func TestHTTPClientURLPort(t *testing.T) { @@ -24,6 +26,11 @@ func TestHTTPClientURLPort(t *testing.T) { if c3.baseURL.String() != "https://example.com:1" { t.Error("Sould use specified port:", c3.baseURL.String()) } + + c4 := NewHTTPClient("example.com") + if c4.baseURL.String() != "http://example.com:80" { + t.Error("Sould add default protocol:", c4.baseURL.String()) + } } func TestHTTPClientSend(t *testing.T) { @@ -111,3 +118,66 @@ func TestHTTPClientHTTPSSend(t *testing.T) { wg.Wait() } + +func TestHTTPClientServerInstantDisconnect(t *testing.T) { + wg := new(sync.WaitGroup) + + GET_payload := []byte("GET / HTTP/1.1\r\n\r\n") + + ln, _ := net.Listen("tcp", ":0") + + go func() { + for { + conn, _ := ln.Accept() + conn.Close() + + wg.Done() + } + }() + + client := NewHTTPClient(ln.Addr().String()) + + wg.Add(2) + client.Send(GET_payload) + client.Send(GET_payload) + + wg.Wait() +} + +func TestHTTPClientServerNoKeepAlive(t *testing.T) { + wg := new(sync.WaitGroup) + + GET_payload := []byte("GET / HTTP/1.1\r\n\r\n") + + ln, _ := net.Listen("tcp", ":0") + + go func() { + for { + conn, err := ln.Accept() + if err != nil { + // handle error + } + + buf := make([]byte, 4096) + reqLen, err := conn.Read(buf) + if err != nil { + t.Error("Error reading:", err.Error()) + } + Debug("Received: ", string(buf[0:reqLen])) + conn.Write([]byte("OK")) + + // No keep-alive connections + conn.Close() + + wg.Done() + } + }() + + client := NewHTTPClient(ln.Addr().String()) + + wg.Add(2) + client.Send(GET_payload) + client.Send(GET_payload) + + wg.Wait() +}