diff --git a/http_client_test.go b/http_client_test.go index d970381..e91284e 100644 --- a/http_client_test.go +++ b/http_client_test.go @@ -75,6 +75,7 @@ func TestHTTPClientSend(t *testing.T) { wg.Done() })) + defer server.Close() client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: true}) @@ -103,6 +104,7 @@ func TestHTTPClientResponseBuffer(t *testing.T) { wg.Done() })) + defer server.Close() client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: false, ResponseBufferSize: 1024}) @@ -153,6 +155,7 @@ func TestHTTPClientHTTPSSend(t *testing.T) { wg.Done() })) + defer server.Close() client := NewHTTPClient(server.URL, &HTTPClientConfig{}) @@ -246,6 +249,7 @@ func TestHTTPClientRedirect(t *testing.T) { wg.Done() })) + defer server.Close() client := NewHTTPClient(server.URL, &HTTPClientConfig{FollowRedirects: 1, Debug: false}) @@ -277,6 +281,7 @@ func TestHTTPClientRedirectLimit(t *testing.T) { wg.Done() })) + defer server.Close() client := NewHTTPClient(server.URL, &HTTPClientConfig{FollowRedirects: 2, Debug: false}) @@ -300,6 +305,7 @@ func TestHTTPClientHandleHTTP10(t *testing.T) { wg.Done() })) + defer server.Close() client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: true}) diff --git a/input_raw_test.go b/input_raw_test.go index 38126f1..608c06e 100644 --- a/input_raw_test.go +++ b/input_raw_test.go @@ -22,11 +22,14 @@ func TestRAWInput(t *testing.T) { wg := new(sync.WaitGroup) quit := make(chan int) - listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {}) + origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) + defer origin.Close() + originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) + var respCounter, reqCounter int64 - input := NewRAWInput(listener.Addr().String(), testRawExpire) + input := NewRAWInput(originAddr, testRawExpire) output := NewTestOutput(func(data []byte) { if data[0] == '1' { atomic.AddInt64(&reqCounter, 1) @@ -42,9 +45,7 @@ func TestRAWInput(t *testing.T) { Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{output} - address := strings.Replace(listener.Addr().String(), "[::]", "127.0.0.1", -1) - - client := NewHTTPClient(address, &HTTPClientConfig{}) + client := NewHTTPClient(origin.URL, &HTTPClientConfig{}) time.Sleep(time.Millisecond) @@ -68,14 +69,15 @@ func TestInputRAW100Expect(t *testing.T) { fileContent, _ := ioutil.ReadFile("README.md") // Origing and Replay server initialization - origin := startHTTP(func(w http.ResponseWriter, req *http.Request) { - defer req.Body.Close() - ioutil.ReadAll(req.Body) + origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer r.Body.Close() + ioutil.ReadAll(r.Body) wg.Done() - }) + })) + defer origin.Close() - originAddr := strings.Replace(origin.Addr().String(), "[::]", "127.0.0.1", -1) + originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) input := NewRAWInput(originAddr, testRawExpire) @@ -92,20 +94,20 @@ func TestInputRAW100Expect(t *testing.T) { } }) - listener := startHTTP(func(w http.ResponseWriter, req *http.Request) { - defer req.Body.Close() - body, _ := ioutil.ReadAll(req.Body) + replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer r.Body.Close() + body, _ := ioutil.ReadAll(r.Body) if !bytes.Equal(body, fileContent) { - buf, _ := httputil.DumpRequest(req, true) + buf, _ := httputil.DumpRequest(r, true) t.Error("Wrong POST body:", string(buf)) } wg.Done() - }) - replayAddr := listener.Addr().String() + })) + defer replay.Close() - httpOutput := NewHTTPOutput(replayAddr, &HTTPOutputConfig{}) + httpOutput := NewHTTPOutput(replay.URL, &HTTPOutputConfig{}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{testOutput, httpOutput} @@ -131,31 +133,30 @@ func TestInputRAWChunkedEncoding(t *testing.T) { fileContent, _ := ioutil.ReadFile("README.md") // Origing and Replay server initialization - origin := startHTTP(func(w http.ResponseWriter, req *http.Request) { - defer req.Body.Close() - ioutil.ReadAll(req.Body) + origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer r.Body.Close() + ioutil.ReadAll(r.Body) wg.Done() - }) - - originAddr := strings.Replace(origin.Addr().String(), "[::]", "127.0.0.1", -1) + })) + originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) input := NewRAWInput(originAddr, testRawExpire) - listener := startHTTP(func(w http.ResponseWriter, req *http.Request) { - defer req.Body.Close() - body, _ := ioutil.ReadAll(req.Body) + replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer r.Body.Close() + body, _ := ioutil.ReadAll(r.Body) if !bytes.Equal(body, fileContent) { - buf, _ := httputil.DumpRequest(req, true) + buf, _ := httputil.DumpRequest(r, true) t.Error("Wrong POST body:", string(buf)) } wg.Done() - }) - replayAddr := listener.Addr().String() + })) + defer replay.Close() - httpOutput := NewHTTPOutput(replayAddr, &HTTPOutputConfig{Debug: true}) + httpOutput := NewHTTPOutput(replay.URL, &HTTPOutputConfig{Debug: true}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{httpOutput} diff --git a/middleware_test.go b/middleware_test.go index 4697d09..8458654 100644 --- a/middleware_test.go +++ b/middleware_test.go @@ -17,7 +17,7 @@ import ( type fakeServiceCb func(string, int, []byte) // Simple service that generate token on request, and require this token for accesing to secure area -func NewFakeSecureService(wg *sync.WaitGroup, cb fakeServiceCb) string { +func NewFakeSecureService(wg *sync.WaitGroup, cb fakeServiceCb) *httptest.Server { active_tokens := make([]string, 0) var mu sync.Mutex @@ -61,8 +61,7 @@ func NewFakeSecureService(wg *sync.WaitGroup, cb fakeServiceCb) string { wg.Done() })) - address := strings.Replace(server.Listener.Addr().String(), "[::]", "127.0.0.1", -1) - return address + return server } func TestFakeSecureService(t *testing.T) { @@ -70,12 +69,13 @@ func TestFakeSecureService(t *testing.T) { wg := new(sync.WaitGroup) - addr := NewFakeSecureService(wg, func(path string, status int, resp []byte) { + server := NewFakeSecureService(wg, func(path string, status int, resp []byte) { }) + defer server.Close() wg.Add(3) - client := NewHTTPClient("http://"+addr, &HTTPClientConfig{Debug: true}) + client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: true}) resp, _ = client.Get("/token") token = proto.Body(resp) @@ -102,11 +102,14 @@ func TestEchoMiddleware(t *testing.T) { w.Header().Set("RequestPath", r.URL.Path) wg.Done() })) + defer from.Close() + to := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Env", "test") w.Header().Set("RequestPath", r.URL.Path) wg.Done() })) + defer to.Close() quit := make(chan int) @@ -150,6 +153,8 @@ func TestTokenMiddleware(t *testing.T) { from := NewFakeSecureService(wg, func(path string, status int, tok []byte) { time.Sleep(10 * time.Millisecond) }) + defer from.Close() + to := NewFakeSecureService(wg, func(path string, status int, tok []byte) { switch path { case "/secure": @@ -160,16 +165,18 @@ func TestTokenMiddleware(t *testing.T) { time.Sleep(10 * time.Millisecond) }) + defer to.Close() quit := make(chan int) Settings.middleware = "go run ./examples/middleware/token_modifier.go" + fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1) // Catch traffic from one service - input := NewRAWInput(from, testRawExpire) + input := NewRAWInput(fromAddr, testRawExpire) // And redirect to another - output := NewHTTPOutput(to, &HTTPOutputConfig{Debug: true}) + output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: true}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{output} @@ -184,7 +191,7 @@ func TestTokenMiddleware(t *testing.T) { // Should receive 2 requests from original + 2 from replayed wg.Add(4) - client := NewHTTPClient("http://"+from, &HTTPClientConfig{Debug: false}) + client := NewHTTPClient(from.URL, &HTTPClientConfig{Debug: false}) // Sending traffic to original service resp, _ = client.Get("/token") diff --git a/output_http_test.go b/output_http_test.go index f4a4e6c..b97399e 100644 --- a/output_http_test.go +++ b/output_http_test.go @@ -3,7 +3,6 @@ package main import ( "io" "io/ioutil" - "net" "net/http" "net/http/httptest" _ "net/http/httputil" @@ -12,25 +11,13 @@ import ( "time" ) -func startHTTP(cb func(http.ResponseWriter, *http.Request)) net.Listener { - handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - cb(w, r) - }) - - listener, _ := net.Listen("tcp", ":0") - - go http.Serve(listener, handler) - - return listener -} - func TestHTTPOutput(t *testing.T) { wg := new(sync.WaitGroup) quit := make(chan int) input := NewTestInput() - listener := startHTTP(func(w http.ResponseWriter, req *http.Request) { + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { if req.Header.Get("User-Agent") != "Gor" { t.Error("Wrong header") } @@ -49,13 +36,14 @@ func TestHTTPOutput(t *testing.T) { } wg.Done() - }) + })) + defer server.Close() headers := HTTPHeaders{HTTPHeader{"User-Agent", "Gor"}} methods := HTTPMethods{[]byte("GET"), []byte("PUT"), []byte("POST")} Settings.modifierConfig = HTTPModifierConfig{headers: headers, methods: methods} - output := NewHTTPOutput(listener.Addr().String(), &HTTPOutputConfig{Debug: true}) + output := NewHTTPOutput(server.URL, &HTTPOutputConfig{Debug: true}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{output} @@ -82,18 +70,19 @@ func TestHTTPOutputKeepOriginalHost(t *testing.T) { input := NewTestInput() - listener := startHTTP(func(w http.ResponseWriter, req *http.Request) { + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { if req.Host != "custom-host.com" { t.Error("Wrong header", req.Host) } wg.Done() - }) + })) + defer server.Close() headers := HTTPHeaders{HTTPHeader{"Host", "custom-host.com"}} Settings.modifierConfig = HTTPModifierConfig{headers: headers} - output := NewHTTPOutput(listener.Addr().String(), &HTTPOutputConfig{Debug: false, OriginalHost: true}) + output := NewHTTPOutput(server.URL, &HTTPOutputConfig{Debug: false, OriginalHost: true}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{output} @@ -140,13 +129,14 @@ func BenchmarkHTTPOutput(b *testing.B) { wg := new(sync.WaitGroup) quit := make(chan int) - listener := startHTTP(func(w http.ResponseWriter, req *http.Request) { + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { time.Sleep(50 * time.Millisecond) wg.Done() - }) + })) + defer server.Close() input := NewTestInput() - output := NewHTTPOutput(listener.Addr().String(), &HTTPOutputConfig{}) + output := NewHTTPOutput(server.URL, &HTTPOutputConfig{}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{output}