diff --git a/Makefile b/Makefile index 7372f13..4eb8a86 100644 --- a/Makefile +++ b/Makefile @@ -32,7 +32,7 @@ race: $(RUN) go test ./... $(ARGS) -v -race -timeout 15s test: - $(RUN) go test ./. -race -timeout 60s $(LDFLAGS) $(ARGS) -v + $(RUN) go test ./. -timeout 60s $(LDFLAGS) $(ARGS) -v test_all: $(RUN) go test ./... -timeout 60s $(LDFLAGS) $(ARGS) -v diff --git a/input_file.go b/input_file.go index 37f677f..281a1a1 100644 --- a/input_file.go +++ b/input_file.go @@ -3,42 +3,103 @@ package main import ( "bufio" "bytes" + "compress/gzip" + "errors" "io" "log" "os" + "path/filepath" + "sort" "strconv" + "strings" "time" ) // FileInput can read requests generated by FileOutput type FileInput struct { - data chan []byte - path string - file *os.File - speedFactor float64 + data chan []byte + path string + currentFile *os.File + currentReader *bufio.Reader + speedFactor float64 + loop bool } // NewFileInput constructor for FileInput. Accepts file path as argument. -func NewFileInput(path string) (i *FileInput) { +func NewFileInput(path string, loop bool) (i *FileInput) { i = new(FileInput) i.data = make(chan []byte) i.path = path i.speedFactor = 1 - i.init(path) + i.loop = loop + + if err := i.updateFile(); err != nil { + return + } go i.emit() return } -func (i *FileInput) init(path string) { - file, err := os.Open(path) +type NextFileNotFound struct{} - if err != nil { - log.Fatal(i, "Cannot open file %q. Error: %s", path, err) +func (_ *NextFileNotFound) Error() string { + return "There is no new files" +} + +// path can be a pattern +// It sort paths lexicographically and tries to choose next one +func (i *FileInput) updateFile() (err error) { + var matches []string + + if matches, err = filepath.Glob(i.path); err != nil { + log.Println("Wrong file pattern", i.path, err) + return } - i.file = file + if len(matches) == 0 { + log.Println("No files match pattern: ", i.path) + return errors.New("No matching files") + } + + sort.Strings(matches) + + // Just pick first file, if there is many, and we are just started + if i.currentFile == nil { + if i.currentFile, err = os.Open(matches[0]); err != nil { + log.Println("Can't read file ", matches[0], err) + return + } + } else { + found := false + for idx, p := range matches { + if p == i.currentFile.Name() && idx != len(matches)-1 { + if i.currentFile, err = os.Open(matches[idx+1]); err != nil { + log.Println("Can't read file ", matches[idx+1], err) + return + } + + found = true + } + } + + if !found { + return new(NextFileNotFound) + } + } + + if strings.HasSuffix(i.currentFile.Name(), ".gz") { + gzReader, err := gzip.NewReader(i.currentFile) + if err != nil { + log.Fatal(err) + } + i.currentReader = bufio.NewReader(gzReader) + } else { + i.currentReader = bufio.NewReader(i.currentFile) + } + + return nil } func (i *FileInput) Read(data []byte) (int, error) { @@ -56,18 +117,35 @@ func (i *FileInput) emit() { var lastTime int64 payloadSeparatorAsBytes := []byte(payloadSeparator) - reader := bufio.NewReader(i.file) + var buffer bytes.Buffer for { - line, err := reader.ReadBytes('\n') + line, err := i.currentReader.ReadBytes('\n') if err != nil { if err != io.EOF { log.Fatal(err) } - break + // If our path pattern match multiple files, try to find them + if err == io.EOF { + if e := i.updateFile(); e != nil { + if _, ok := e.(*NextFileNotFound); ok && i.loop { + // Start from the first file + i.currentFile = nil + i.currentReader = nil + lastTime = 0 + i.updateFile() + + continue + } else { + break + } + } + + continue + } } if bytes.Equal(payloadSeparatorAsBytes[1:], line) { @@ -105,3 +183,7 @@ func (i *FileInput) emit() { log.Printf("FileInput: end of file '%s'\n", i.path) } + +func (i *FileInput) Close() { + i.currentFile.Close() +} diff --git a/input_file_test.go b/input_file_test.go index cbe79ba..78e0d56 100644 --- a/input_file_test.go +++ b/input_file_test.go @@ -3,8 +3,11 @@ package main import ( "bytes" "errors" + "fmt" "io" "io/ioutil" + "log" + "math/rand" "os" "sync" "syscall" @@ -12,6 +15,8 @@ import ( "time" ) +var _ = log.Println + func TestInputFileWithGET(t *testing.T) { input := NewTestInput() @@ -35,7 +40,6 @@ func TestInputFileWithGET(t *testing.T) { t.Error("Request read back from file should match") } } - } func TestInputFileWithPayloadLargerThan64Kb(t *testing.T) { @@ -93,6 +97,90 @@ func TestInputFileWithGETAndPOST(t *testing.T) { } +func TestInputFileMultipleFiles(t *testing.T) { + rnd := rand.Int63() + + file1, _ := os.OpenFile(fmt.Sprintf("/tmp/%d_0", rnd), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660) + file1.Write([]byte("1 1 1\ntest1")) + file1.Write([]byte(payloadSeparator)) + file1.Write([]byte("1 1 2\ntest2")) + file1.Write([]byte(payloadSeparator)) + file1.Close() + + file2, _ := os.OpenFile(fmt.Sprintf("/tmp/%d_1", rnd), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660) + file2.Write([]byte("1 1 3\ntest3")) + file2.Write([]byte(payloadSeparator)) + file2.Close() + + input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd), false) + buf := make([]byte, 1000) + n, _ := input.Read(buf) + if buf[10] != '1' { + t.Error("Shound emit requests in right order", string(buf[:n])) + } + input.Read(buf) + if buf[10] != '2' { + t.Error("Shound emit requests in right order", string(buf[:n])) + } + + input.Read(buf) + if buf[10] != '3' { + t.Error("Shound emit requests from second file", string(buf[:n])) + } + + os.Remove(file1.Name()) + os.Remove(file2.Name()) +} + +func TestInputFileLoop(t *testing.T) { + rnd := rand.Int63() + + file, _ := os.OpenFile(fmt.Sprintf("/tmp/%d", rnd), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660) + file.Write([]byte("1 1 1\ntest1")) + file.Write([]byte(payloadSeparator)) + file.Write([]byte("1 1 2\ntest2")) + file.Write([]byte(payloadSeparator)) + file.Close() + + input := NewFileInput(fmt.Sprintf("/tmp/%d", rnd), true) + buf := make([]byte, 1000) + + // Even if we have just 2 requests in file, it should indifinitly loop + for i := 0; i < 1000; i++ { + input.Read(buf) + } + input.Close() + + os.Remove(file.Name()) +} + +func TestInputFileCompressed(t *testing.T) { + rnd := rand.Int63() + + output := NewFileOutput(fmt.Sprintf("/tmp/%d_0.gz", rnd), time.Minute) + for i := 0; i < 1000; i++ { + output.Write([]byte("1 1 1\r\ntest")) + } + name1 := output.file.Name() + output.Close() + + output2 := NewFileOutput(fmt.Sprintf("/tmp/%d_1.gz", rnd), time.Minute) + for i := 0; i < 1000; i++ { + output2.Write([]byte("1 1 1\r\ntest")) + } + name2 := output2.file.Name() + output2.Close() + + input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd), false) + buf := make([]byte, 1000) + for i := 0; i < 2000; i++ { + input.Read(buf) + } + + os.Remove(name1) + os.Remove(name2) +} + type CaptureFile struct { data [][]byte file *os.File @@ -155,13 +243,12 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile { readPayloads := [][]byte{} output := NewTestOutput(func(data []byte) { - readPayloads = append(readPayloads, Duplicate(data)) requestGenerator.wg.Done() }) - outputFile := NewFileOutput(f.Name()) + outputFile := NewFileOutput(f.Name(), time.Minute) Plugins.Inputs = requestGenerator.inputs Plugins.Outputs = []io.Writer{output, outputFile} @@ -171,6 +258,9 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile { requestGenerator.emit() requestGenerator.wg.Wait() + time.Sleep(100 * time.Millisecond) + outputFile.Close() + close(quit) return NewExpectedCaptureFile(readPayloads, f) @@ -182,7 +272,7 @@ func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback quit := make(chan int) wg := new(sync.WaitGroup) - input := NewFileInput(captureFile.Name()) + input := NewFileInput(captureFile.Name(), false) output := NewTestOutput(func(data []byte) { callback(data) wg.Done() diff --git a/output_file.go b/output_file.go index 5ccc93c..8dcc11a 100644 --- a/output_file.go +++ b/output_file.go @@ -1,34 +1,70 @@ package main import ( + "bufio" + "compress/gzip" + "fmt" "io" "log" "os" + "strings" + "time" ) +var dateFileNameFuncs = map[string]func() string{ + "%Y": func() string { return time.Now().Format("2006") }, + "%m": func() string { return time.Now().Format("01") }, + "%d": func() string { return time.Now().Format("02") }, + "%H": func() string { return time.Now().Format("15") }, + "%M": func() string { return time.Now().Format("04") }, + "%S": func() string { return time.Now().Format("05") }, + "%NS": func() string { return fmt.Sprint(time.Now().Nanosecond()) }, +} + // FileOutput output plugin type FileOutput struct { - path string - file *os.File + pathTemplate string + currentName string + file *os.File + writer io.Writer } // NewFileOutput constructor for FileOutput, accepts path -func NewFileOutput(path string) io.Writer { +func NewFileOutput(pathTemplate string, flushInterval time.Duration) *FileOutput { o := new(FileOutput) - o.path = path - o.init(path) + o.pathTemplate = pathTemplate + o.updateName() + + // Force flushing every minute + go func() { + for { + time.Sleep(flushInterval) + o.flush() + } + }() + + go func() { + for { + time.Sleep(time.Second) + o.updateName() + } + }() return o } -func (o *FileOutput) init(path string) { - var err error +func (o *FileOutput) filename() string { + path := o.pathTemplate - o.file, err = os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660) - - if err != nil { - log.Fatal(o, "Cannot open file %q. Error: %s", path, err) + for name, fn := range dateFileNameFuncs { + path = strings.Replace(path, name, fn(), -1) } + + return path +} + +func (o *FileOutput) updateName() { + o.currentName = o.filename() } func (o *FileOutput) Write(data []byte) (n int, err error) { @@ -36,12 +72,49 @@ func (o *FileOutput) Write(data []byte) (n int, err error) { return len(data), nil } - o.file.Write(data) - o.file.Write([]byte(payloadSeparator)) + if o.file == nil || o.currentName != o.file.Name() { + o.Close() + + o.file, err = os.OpenFile(o.currentName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660) + + if strings.HasSuffix(o.currentName, ".gz") { + o.writer = gzip.NewWriter(o.file) + } else { + o.writer = bufio.NewWriter(o.file) + } + + if err != nil { + log.Fatal(o, "Cannot open file %q. Error: %s", o.currentName, err) + } + } + + o.writer.Write(data) + o.writer.Write([]byte(payloadSeparator)) return len(data), nil } +func (o *FileOutput) flush() { + if o.file != nil { + if strings.HasSuffix(o.currentName, ".gz") { + o.writer.(*gzip.Writer).Flush() + } else { + o.writer.(*bufio.Writer).Flush() + } + } +} + func (o *FileOutput) String() string { - return "File output: " + o.path + return "File output: " + o.file.Name() +} + +func (o *FileOutput) Close() { + if o.file != nil { + if strings.HasSuffix(o.currentName, ".gz") { + o.writer.(*gzip.Writer).Close() + } else { + o.writer.(*bufio.Writer).Flush() + } + o.file.Close() + } } diff --git a/output_file_test.go b/output_file_test.go index 4a8b0ae..5db0f8d 100644 --- a/output_file_test.go +++ b/output_file_test.go @@ -1,9 +1,13 @@ package main import ( + "fmt" "io" + "os" "sync" + "sync/atomic" "testing" + "time" ) func TestFileOutput(t *testing.T) { @@ -11,7 +15,7 @@ func TestFileOutput(t *testing.T) { quit := make(chan int) input := NewTestInput() - output := NewFileOutput("/tmp/test_requests.gor") + output := NewFileOutput("/tmp/test_requests.gor", time.Minute) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{output} @@ -23,12 +27,17 @@ func TestFileOutput(t *testing.T) { input.EmitGET() input.EmitPOST() } + time.Sleep(100 * time.Millisecond) + output.flush() + close(quit) quit = make(chan int) - input2 := NewFileInput("/tmp/test_requests.gor") + var counter int64 + input2 := NewFileInput("/tmp/test_requests.gor", false) output2 := NewTestOutput(func(data []byte) { + atomic.AddInt64(&counter, 1) wg.Done() }) @@ -40,3 +49,67 @@ func TestFileOutput(t *testing.T) { wg.Wait() close(quit) } + +func TestFileOutputPathTemplate(t *testing.T) { + output := &FileOutput{pathTemplate: "/tmp/log-%Y-%m-%d-%S"} + now := time.Now() + expectedPath := fmt.Sprintf("/tmp/log-%s-%s-%s-%s", now.Format("2006"), now.Format("01"), now.Format("02"), now.Format("05")) + path := output.filename() + + if expectedPath != path { + t.Errorf("Expected path %s but got %s", expectedPath, path) + } +} + +func TestFileOutputMultipleFiles(t *testing.T) { + output := NewFileOutput("/tmp/log-%Y-%m-%d-%S", time.Minute) + + if output.file != nil { + t.Error("Should not initialize file if no writes") + } + + output.Write([]byte("1 1 1\r\ntest")) + name1 := output.file.Name() + + output.Write([]byte("1 1 1\r\ntest")) + name2 := output.file.Name() + + time.Sleep(time.Second) + output.updateName() + + output.Write([]byte("1 1 1\r\ntest")) + name3 := output.file.Name() + + if name2 != name1 { + t.Error("Fast changes should happen in same file:", name1, name2, name3) + } + + if name3 == name1 { + t.Error("File name should change:", name1, name2, name3) + } + + os.Remove(name1) + os.Remove(name3) +} + +func TestFileOutputCompression(t *testing.T) { + output := NewFileOutput("/tmp/log-%Y-%m-%d-%S.gz", time.Minute) + + if output.file != nil { + t.Error("Should not initialize file if no writes") + } + + for i := 0; i < 1000; i++ { + output.Write([]byte("1 1 1\r\ntest")) + } + + name := output.file.Name() + output.Close() + + s, _ := os.Stat(name) + if s.Size() == 12*1000 { + t.Error("Should be compressed file:", s.Size()) + } + + os.Remove(name) +} diff --git a/plugins.go b/plugins.go index d09bfe8..805a086 100644 --- a/plugins.go +++ b/plugins.go @@ -113,11 +113,11 @@ func InitPlugins() { } for _, options := range Settings.inputFile { - registerPlugin(NewFileInput, options) + registerPlugin(NewFileInput, options, Settings.inputFileLoop) } for _, options := range Settings.outputFile { - registerPlugin(NewFileOutput, options) + registerPlugin(NewFileOutput, options, Settings.outputFileFlushInterval) } for _, options := range Settings.inputHTTP { diff --git a/settings.go b/settings.go index 24b253c..a49bd26 100644 --- a/settings.go +++ b/settings.go @@ -39,8 +39,10 @@ type AppSettings struct { outputTCP MultiOption outputTCPStats bool - inputFile MultiOption - outputFile MultiOption + inputFile MultiOption + inputFileLoop bool + outputFile MultiOption + outputFileFlushInterval time.Duration inputRAW MultiOption inputRAWEngine string @@ -83,7 +85,10 @@ func init() { flag.BoolVar(&Settings.outputTCPStats, "output-tcp-stats", false, "Report TCP output queue stats to console every 5 seconds.") flag.Var(&Settings.inputFile, "input-file", "Read requests from file: \n\tgor --input-file ./requests.gor --output-http staging.com") + flag.BoolVar(&Settings.inputFileLoop, "input-file-loop", false, "Loop input files, useful for performance testing.") + flag.Var(&Settings.outputFile, "output-file", "Write incoming requests to file: \n\tgor --input-raw :80 --output-file ./requests.gor") + flag.DurationVar(&Settings.outputFileFlushInterval, "output-file-flush-interval", time.Minute, "Interval for forcing buffer flush to the file, default: 60s.") flag.Var(&Settings.inputRAW, "input-raw", "Capture traffic from given port (use RAW sockets and require *sudo* access):\n\t# Capture traffic from 8080 port\n\tgor --input-raw :8080 --output-http staging.com")