github/goreplay
4
0
Fork 0
mirror of https://github.com/buger/goreplay.git synced 2024-04-21 12:32:02 +00:00
Code Issues Releases Wiki Activity

Refactor emitter.go and fix test accordingly.

Browse Source
This commit is contained in:
Arijit Das
2020-06-06 16:00:43 +05:30
parent 932c894b6b
commit db28858804
17 changed files with 225 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
emitter.go
+37 -19
View File
@@ -8,13 +8,24 @@ import (
"time"
)
var wg sync.WaitGroup
var closeOnce sync.Once
type emitter struct {
sync.WaitGroup
quit chan int
}
func NewEmitter(quit chan int) *emitter {
return &emitter{
quit: quit,
}
}
// Start initialize loop for sending data from inputs to outputs
func Start(plugins *InOutPlugins, stop chan int) {
if Settings.middleware != "" {
middleware := NewMiddleware(Settings.middleware)
func (e *emitter) Start(plugins *InOutPlugins, middlewareCmd string) {
e.Add(1)
defer e.Done()
if middlewareCmd != "" {
middleware := NewMiddleware(middlewareCmd)
for _, in := range plugins.Inputs {
middleware.ReadFrom(in)
@@ -26,31 +37,34 @@ func Start(plugins *InOutPlugins, stop chan int) {
middleware.ReadFrom(r)
}
}
wg.Add(1)
e.Add(1)
go func() {
defer e.Done()
if err := CopyMulty(middleware, plugins.Outputs...); err != nil {
log.Println("Error during copy: ", err)
Close(stop)
e.close()
}
}()
} else {
for _, in := range plugins.Inputs {
wg.Add(1)
e.Add(1)
go func(in io.Reader) {
defer e.Done()
if err := CopyMulty(in, plugins.Outputs...); err != nil {
log.Println("Error during copy: ", err)
Close(stop)
e.close()
}
}(in)
}
for _, out := range plugins.Outputs {
if r, ok := out.(io.Reader); ok {
wg.Add(1)
e.Add(1)
go func(r io.Reader) {
defer e.Done()
if err := CopyMulty(r, plugins.Outputs...); err != nil {
log.Println("Error during copy: ", err)
Close(stop)
e.close()
}
}(r)
}
@@ -59,7 +73,7 @@ func Start(plugins *InOutPlugins, stop chan int) {
for {
select {
case <-stop:
case <-e.quit:
finalize(plugins)
return
case <-time.After(100 * time.Millisecond):
@@ -67,17 +81,22 @@ func Start(plugins *InOutPlugins, stop chan int) {
}
}
func (e *emitter) close() {
select {
case <- e.quit:
default:
close(e.quit)
}
}
// Close closes all the goroutine and waits for it to finish.
func Close(quit chan int) {
closeOnce.Do(func() {
close(quit)
})
wg.Wait()
func (e *emitter) Close() {
e.close()
e.Wait()
}
// CopyMulty copies from 1 reader to multiple writers
func CopyMulty(src io.Reader, writers ...io.Writer) error {
defer wg.Done()
buf := make([]byte, Settings.copyBufferSize)
wIndex := 0
modifier := NewHTTPModifier(&Settings.modifierConfig)
@@ -192,5 +211,4 @@ func CopyMulty(src io.Reader, writers ...io.Writer) error {
i++
}
}
emitter_test.go
+18 -12
View File
@@ -21,17 +21,18 @@ func TestEmitter(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
for i := 0; i < 1000; i++ {
for i := 0; i < 1; i++ {
wg.Add(1)
input.EmitGET()
}
wg.Wait()
close(quit)
emitter.Close()
}
func TestEmitterFiltered(t *testing.T) {
@@ -49,10 +50,13 @@ func TestEmitterFiltered(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
methods := HTTPMethods{[]byte("GET")}
Settings.modifierConfig = HTTPModifierConfig{methods: methods}
go Start(plugins, quit)
emitter := &emitter{quit: quit}
go emitter.Start(plugins, "")
wg.Add(2)
@@ -77,8 +81,7 @@ func TestEmitterFiltered(t *testing.T) {
input.EmitBytes(respb)
wg.Wait()
Close(quit)
emitter.Close()
Settings.modifierConfig = HTTPModifierConfig{}
}
@@ -105,10 +108,12 @@ func TestEmitterRoundRobin(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output1, output2},
}
plugins.All = append(plugins.All, input, output1, output2)
Settings.splitOutput = true
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
for i := 0; i < 1000; i++ {
wg.Add(1)
@@ -116,8 +121,7 @@ func TestEmitterRoundRobin(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
if counter1 == 0 || counter2 == 0 {
t.Errorf("Round robin should split traffic equally: %d vs %d", counter1, counter2)
@@ -140,8 +144,10 @@ func BenchmarkEmitter(b *testing.B) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
b.ResetTimer()
@@ -151,5 +157,5 @@ func BenchmarkEmitter(b *testing.B) {
}
wg.Wait()
close(quit)
emitter.Close()
}
gor.go
+2 -1
View File
@@ -86,6 +86,7 @@ func main() {
}()
}
emitter := NewEmitter(closeCh)
c := make(chan os.Signal, 1)
signal.Notify(c, os.Interrupt, syscall.SIGTERM)
go func() {
@@ -103,7 +104,7 @@ func main() {
})
}
Start(plugins, closeCh)
emitter.Start(plugins, Settings.middleware)
}
func finalize(plugins *InOutPlugins) {
input_file.go
+14 -5
View File
@@ -113,7 +113,7 @@ type FileInput struct {
func NewFileInput(path string, loop bool) (i *FileInput) {
i = new(FileInput)
i.data = make(chan []byte, 1000)
i.exit = make(chan bool, 1)
i.exit = make(chan bool)
i.path = path
i.speedFactor = 1
i.loop = loop
@@ -153,7 +153,10 @@ func (i *FileInput) init() (err error) {
}
func (i *FileInput) Read(data []byte) (int, error) {
buf := <-i.data
buf, ok := <-i.data
if !ok {
return 0, os.ErrClosed
}
copy(data, buf)
return len(buf), nil
@@ -214,7 +217,13 @@ func (i *FileInput) emit() {
lastTime = reader.timestamp
}
i.data <- reader.ReadPayload()
// Recheck if we have exited since last check.
select {
case <-i.exit:
return
default:
i.data <- reader.ReadPayload()
}
}
log.Printf("FileInput: end of file '%s'\n", i.path)
@@ -231,8 +240,8 @@ func (i *FileInput) Close() error {
defer i.mu.Unlock()
i.mu.Lock()
i.exit <- true
close(i.exit)
close(i.data)
for _, r := range i.readers {
r.Close()
}
input_file_test.go
+11 -11
View File
@@ -18,7 +18,6 @@ import (
var _ = log.Println
func TestInputFileWithGET(t *testing.T) {
input := NewTestInput()
rg := NewRequestGenerator([]io.Reader{input}, func() { input.EmitGET() }, 1)
readPayloads := [][]byte{}
@@ -305,7 +304,6 @@ func (expectedCaptureFile *CaptureFile) PayloadsEqual(other [][]byte) bool {
}
func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile {
f, err := ioutil.TempFile("", "testmainconf")
if err != nil {
panic(err)
@@ -316,7 +314,6 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile {
readPayloads := [][]byte{}
output := NewTestOutput(func(data []byte) {
readPayloads = append(readPayloads, Duplicate(data))
requestGenerator.wg.Done()
})
@@ -326,23 +323,25 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile {
Inputs: requestGenerator.inputs,
Outputs: []io.Writer{output, outputFile},
}
for _, input := range requestGenerator.inputs {
plugins.All = append(plugins.All, input)
}
plugins.All = append(plugins.All, output, outputFile)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
requestGenerator.emit()
requestGenerator.wg.Wait()
time.Sleep(100 * time.Millisecond)
outputFile.Close()
close(quit)
emitter.Close()
return NewExpectedCaptureFile(readPayloads, f)
}
func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback) (err error) {
quit := make(chan int)
wg := new(sync.WaitGroup)
@@ -356,9 +355,11 @@ func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
wg.Add(count)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
done := make(chan int, 1)
go func() {
@@ -372,8 +373,7 @@ func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback
case <-time.After(2 * time.Second):
err = errors.New("Timed out")
}
close(quit)
emitter.close()
return
}
input_http.go
+10 -1
View File
@@ -5,6 +5,7 @@ import (
"net"
"net/http"
"net/http/httputil"
"os"
"time"
)
@@ -27,7 +28,10 @@ func NewHTTPInput(address string) (i *HTTPInput) {
}
func (i *HTTPInput) Read(data []byte) (int, error) {
buf := <-i.data
buf, ok := <-i.data
if !ok {
return 0, os.ErrClosed
}
header := payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1)
@@ -37,6 +41,11 @@ func (i *HTTPInput) Read(data []byte) (int, error) {
return len(buf) + len(header), nil
}
func (i *HTTPInput) Close() error {
close(i.data)
return nil
}
func (i *HTTPInput) handler(w http.ResponseWriter, r *http.Request) {
r.URL.Scheme = "http"
r.URL.Host = i.listener.Addr().String()
input_http_test.go
+8 -5
View File
@@ -25,8 +25,10 @@ func TestHTTPInput(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
address := strings.Replace(input.listener.Addr().String(), "[::]", "127.0.0.1", -1)
@@ -36,8 +38,7 @@ func TestHTTPInput(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
}
func TestInputHTTPLargePayload(t *testing.T) {
@@ -61,8 +62,10 @@ func TestInputHTTPLargePayload(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
wg.Add(1)
address := strings.Replace(input.listener.Addr().String(), "[::]", "127.0.0.1", -1)
@@ -73,5 +76,5 @@ func TestInputHTTPLargePayload(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
}
input_raw.go
+6 -1
View File
@@ -3,6 +3,7 @@ package main
import (
"log"
"net"
"os"
"time"
"github.com/buger/goreplay/proto"
@@ -52,7 +53,10 @@ func NewRAWInput(address string, engine int, trackResponse bool, expire time.Dur
}
func (i *RAWInput) Read(data []byte) (int, error) {
msg := <-i.data
msg, ok := <-i.data
if !ok {
return 0, os.ErrClosed
}
buf := msg.Bytes()
var header []byte
@@ -109,5 +113,6 @@ func (i *RAWInput) String() string {
func (i *RAWInput) Close() error {
i.listener.Close()
close(i.quit)
close(i.data)
return nil
}
input_raw_test.go
+28 -26
View File
@@ -45,8 +45,6 @@ func TestRAWInputIPv4(t *testing.T) {
var respCounter, reqCounter int64
input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "X-Real-IP", "", "", 0)
defer input.Close()
output := NewTestOutput(func(data []byte) {
if data[0] == '1' {
body := payloadBody(data)
@@ -69,10 +67,12 @@ func TestRAWInputIPv4(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
client := NewHTTPClient("http://"+listener.Addr().String(), &HTTPClientConfig{})
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
for i := 0; i < 100; i++ {
// request + response
@@ -82,8 +82,7 @@ func TestRAWInputIPv4(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
}
func TestRAWInputNoKeepAlive(t *testing.T) {
@@ -109,7 +108,6 @@ func TestRAWInputNoKeepAlive(t *testing.T) {
originAddr := listener.Addr().String()
input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "", "", 0)
defer input.Close()
output := NewTestOutput(func(data []byte) {
wg.Done()
@@ -119,10 +117,12 @@ func TestRAWInputNoKeepAlive(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
client := NewHTTPClient("http://"+listener.Addr().String(), &HTTPClientConfig{})
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
for i := 0; i < 100; i++ {
// request + response
@@ -132,8 +132,7 @@ func TestRAWInputNoKeepAlive(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
}
func TestRAWInputIPv6(t *testing.T) {
@@ -157,7 +156,6 @@ func TestRAWInputIPv6(t *testing.T) {
var respCounter, reqCounter int64
input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "", "", 0)
defer input.Close()
output := NewTestOutput(func(data []byte) {
if data[0] == '1' {
@@ -177,10 +175,12 @@ func TestRAWInputIPv6(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
client := NewHTTPClient("http://"+listener.Addr().String(), &HTTPClientConfig{})
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
for i := 0; i < 100; i++ {
// request + response
@@ -190,7 +190,7 @@ func TestRAWInputIPv6(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
}
func TestInputRAW100Expect(t *testing.T) {
@@ -210,7 +210,6 @@ func TestInputRAW100Expect(t *testing.T) {
originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "", "", "", 0)
defer input.Close()
// We will use it to get content of raw HTTP request
testOutput := NewTestOutput(func(data []byte) {
@@ -244,8 +243,10 @@ func TestInputRAW100Expect(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{testOutput, httpOutput},
}
plugins.All = append(plugins.All, input, testOutput, httpOutput)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
// Origin + Response/Request Test Output + Request Http Output
wg.Add(4)
@@ -256,7 +257,7 @@ func TestInputRAW100Expect(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
}
func TestInputRAWChunkedEncoding(t *testing.T) {
@@ -275,7 +276,6 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "", "", "", 0)
defer input.Close()
replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
defer r.Body.Close()
@@ -296,9 +296,10 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{httpOutput},
}
plugins.All = append(plugins.All, input, httpOutput)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
wg.Add(2)
curl := exec.Command("curl", "http://"+originAddr, "--header", "Transfer-Encoding: chunked", "--header", "Expect:", "--data-binary", "@README.md")
@@ -308,8 +309,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
}
func TestInputRAWLargePayload(t *testing.T) {
@@ -341,7 +341,6 @@ func TestInputRAWLargePayload(t *testing.T) {
originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "", "", 0)
defer input.Close()
replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
body, _ := ioutil.ReadAll(req.Body)
@@ -364,8 +363,10 @@ func TestInputRAWLargePayload(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{httpOutput},
}
plugins.All = append(plugins.All, input, httpOutput)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
wg.Add(2)
curl := exec.Command("curl", "http://"+originAddr, "--header", "Transfer-Encoding: chunked", "--header", "Expect:", "--data-binary", "@/tmp/large")
@@ -375,7 +376,7 @@ func TestInputRAWLargePayload(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
}
func BenchmarkRAWInput(b *testing.B) {
@@ -394,7 +395,6 @@ func BenchmarkRAWInput(b *testing.B) {
upstreamAddr := strings.Replace(upstream.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "", "", 0)
defer input.Close()
output := NewTestOutput(func(data []byte) {
if data[0] == '1' {
@@ -410,8 +410,10 @@ func BenchmarkRAWInput(b *testing.B) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output, httpOutput},
}
plugins.All = append(plugins.All, input, output, httpOutput)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
emitted := 0
fileContent, _ := ioutil.ReadFile("LICENSE.txt")
@@ -442,5 +444,5 @@ func BenchmarkRAWInput(b *testing.B) {
time.Sleep(400 * time.Millisecond)
log.Println("Emitted ", emitted, ", Captured ", reqCounter, "requests and ", respCounter, " responses", "and replayed", replayCounter)
close(quit)
emitter.Close()
}
input_tcp.go
+9 -1
View File
@@ -38,12 +38,20 @@ func NewTCPInput(address string, config *TCPInputConfig) (i *TCPInput) {
}
func (i *TCPInput) Read(data []byte) (int, error) {
buf := <-i.data
buf, ok := <-i.data
if !ok {
return 0, os.ErrClosed
}
copy(data, buf)
return len(buf), nil
}
func (i *TCPInput) Close() error {
close(i.data)
return nil
}
func (i *TCPInput) listen(address string) {
if i.config.secure {
cer, err := tls.LoadX509KeyPair(i.config.certificatePath, i.config.keyPath)
input_tcp_test.go
+8 -6
View File
@@ -31,8 +31,10 @@ func TestTCPInput(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
tcpAddr, err := net.ResolveTCPAddr("tcp", input.listener.Addr().String())
@@ -55,8 +57,7 @@ func TestTCPInput(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
}
func genCertificate(template *x509.Certificate) ([]byte, []byte) {
@@ -113,8 +114,10 @@ func TestTCPInputSecure(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
conf := &tls.Config{
InsecureSkipVerify: true,
@@ -135,6 +138,5 @@ func TestTCPInputSecure(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
}
limiter_test.go
+16 -12
View File
@@ -22,16 +22,17 @@ func TestOutputLimiter(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
for i := 0; i < 100; i++ {
input.EmitGET()
}
wg.Wait()
close(quit)
emitter.Close()
}
func TestInputLimiter(t *testing.T) {
@@ -48,16 +49,17 @@ func TestInputLimiter(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
for i := 0; i < 100; i++ {
input.(*Limiter).plugin.(*TestInput).EmitGET()
}
wg.Wait()
close(quit)
emitter.Close()
}
// Should limit all requests
@@ -74,16 +76,17 @@ func TestPercentLimiter1(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
for i := 0; i < 100; i++ {
input.EmitGET()
}
wg.Wait()
close(quit)
emitter.Close()
}
// Should not limit at all
@@ -101,14 +104,15 @@ func TestPercentLimiter2(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
for i := 0; i < 100; i++ {
input.EmitGET()
}
wg.Wait()
close(quit)
emitter.Close()
}
middleware_test.go
+8 -6
View File
@@ -119,7 +119,6 @@ func TestEchoMiddleware(t *testing.T) {
// Catch traffic from one service
fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "", "", "", 0)
defer input.Close()
// And redirect to another
output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: false})
@@ -128,9 +127,11 @@ func TestEchoMiddleware(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
// Start Gor
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
// Wait till middleware initialization
time.Sleep(100 * time.Millisecond)
@@ -148,7 +149,7 @@ func TestEchoMiddleware(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
time.Sleep(200 * time.Millisecond)
Settings.middleware = ""
@@ -183,7 +184,6 @@ func TestTokenMiddleware(t *testing.T) {
fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
// Catch traffic from one service
input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "", "", "", 0)
defer input.Close()
// And redirect to another
output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: true})
@@ -192,9 +192,11 @@ func TestTokenMiddleware(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
// Start Gor
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
// Wait for middleware to initialize
// Give go compiller time to build programm
@@ -219,7 +221,7 @@ func TestTokenMiddleware(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
time.Sleep(100 * time.Millisecond)
Settings.middleware = ""
}
output_file_test.go
+6 -7
View File
@@ -24,8 +24,10 @@ func TestFileOutput(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
for i := 0; i < 100; i++ {
wg.Add(2)
@@ -34,10 +36,7 @@ func TestFileOutput(t *testing.T) {
}
time.Sleep(100 * time.Millisecond)
output.flush()
close(quit)
quit = make(chan int)
emitter.Close()
var counter int64
input2 := NewFileInput("/tmp/test_requests.gor", false)
@@ -51,10 +50,10 @@ func TestFileOutput(t *testing.T) {
Outputs: []io.Writer{output2},
}
go Start(plugins2, quit)
go emitter.Start(plugins2, Settings.middleware)
wg.Wait()
close(quit)
emitter.Close()
}
func TestFileOutputWithNameCleaning(t *testing.T) {
output_http_test.go
+16 -13
View File
@@ -52,8 +52,10 @@ func TestHTTPOutput(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{http_output, output},
}
plugins.All = append(plugins.All, input, output, http_output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
for i := 0; i < 1; i++ {
// 2 http-output, 2 - test output request, 2 - test output http response
@@ -64,9 +66,7 @@ func TestHTTPOutput(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
Settings.modifierConfig = HTTPModifierConfig{}
}
@@ -94,16 +94,16 @@ func TestHTTPOutputKeepOriginalHost(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
wg.Add(1)
input.EmitGET()
wg.Wait()
close(quit)
emitter.Close()
Settings.modifierConfig = HTTPModifierConfig{}
}
@@ -123,8 +123,10 @@ func TestOutputHTTPSSL(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
wg.Add(2)
@@ -132,7 +134,7 @@ func TestOutputHTTPSSL(t *testing.T) {
input.EmitGET()
wg.Wait()
close(quit)
emitter.Close()
}
func BenchmarkHTTPOutput(b *testing.B) {
@@ -152,8 +154,10 @@ func BenchmarkHTTPOutput(b *testing.B) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
for i := 0; i < b.N; i++ {
wg.Add(1)
@@ -161,6 +165,5 @@ func BenchmarkHTTPOutput(b *testing.B) {
}
wg.Wait()
close(quit)
emitter.Close()
}
output_tcp_test.go
+8 -6
View File
@@ -24,8 +24,10 @@ func TestTCPOutput(t *testing.T) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
for i := 0; i < 100; i++ {
wg.Add(1)
@@ -33,8 +35,7 @@ func TestTCPOutput(t *testing.T) {
}
wg.Wait()
close(quit)
emitter.Close()
}
func startTCP(cb func([]byte)) net.Listener {
@@ -78,8 +79,10 @@ func BenchmarkTCPOutput(b *testing.B) {
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
}
plugins.All = append(plugins.All, input, output)
go Start(plugins, quit)
emitter := NewEmitter(quit)
go emitter.Start(plugins, Settings.middleware)
b.ResetTimer()
for i := 0; i < b.N; i++ {
@@ -88,8 +91,7 @@ func BenchmarkTCPOutput(b *testing.B) {
}
wg.Wait()
close(quit)
emitter.Close()
}
func TestStickyDisable(t *testing.T) {
test_input.go
+20 -16
View File
@@ -3,7 +3,7 @@ package main
import (
"crypto/rand"
"encoding/base64"
"fmt"
"io"
"time"
)
@@ -22,22 +22,26 @@ func NewTestInput() (i *TestInput) {
}
func (i *TestInput) Read(data []byte) (int, error) {
select {
case buf := <-i.data:
var header []byte
if !i.skipHeader {
header = payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1)
copy(data[0:len(header)], header)
copy(data[len(header):], buf)
} else {
copy(data, buf)
}
return len(buf) + len(header), nil
case <-time.After(10* time.Second):
return 0, fmt.Errorf("timed out waiting for read")
buf, ok := <-i.data
if !ok {
return 0, io.EOF
}
var header []byte
if !i.skipHeader {
header = payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1)
copy(data[0:len(header)], header)
copy(data[len(header):], buf)
} else {
copy(data, buf)
}
return len(buf) + len(header), nil
}
func (i *TestInput) Close() error {
close(i.data)
return nil
}
func (i *TestInput) EmitBytes(data []byte) {