diff --git a/input_file_test.go b/input_file_test.go index cbe79ba..53cba82 100644 --- a/input_file_test.go +++ b/input_file_test.go @@ -161,7 +161,7 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile { requestGenerator.wg.Done() }) - outputFile := NewFileOutput(f.Name()) + outputFile := NewFileOutput(f.Name(), time.Minute) Plugins.Inputs = requestGenerator.inputs Plugins.Outputs = []io.Writer{output, outputFile} diff --git a/output_file.go b/output_file.go index 5ccc93c..fec6bf4 100644 --- a/output_file.go +++ b/output_file.go @@ -1,34 +1,70 @@ package main import ( - "io" + "bufio" + "fmt" "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 *bufio.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) + if err := o.writer.Flush(); err != nil { + break + } + } + }() + + 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,37 @@ 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() { + if o.file != nil { + o.writer.Flush() + o.file.Close() + } + + o.file, err = os.OpenFile(o.currentName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660) + 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 { + o.writer.Flush() + } +} + func (o *FileOutput) String() string { - return "File output: " + o.path + return "File output: " + o.file.Name() +} + +func (o *FileOutput) Close() { + o.writer.Flush() + o.file.Close() } diff --git a/output_file_test.go b/output_file_test.go index 4a8b0ae..51a5109 100644 --- a/output_file_test.go +++ b/output_file_test.go @@ -1,9 +1,14 @@ package main import ( + "fmt" "io" + "log" + "os" "sync" + "sync/atomic" "testing" + "time" ) func TestFileOutput(t *testing.T) { @@ -11,7 +16,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 +28,18 @@ func TestFileOutput(t *testing.T) { input.EmitGET() input.EmitPOST() } + time.Sleep(100 * time.Millisecond) + output.Flush() + close(quit) quit = make(chan int) + var counter int64 input2 := NewFileInput("/tmp/test_requests.gor") output2 := NewTestOutput(func(data []byte) { + atomic.AddInt64(&counter, 1) + log.Println(counter) wg.Done() }) @@ -40,3 +51,44 @@ 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.Write([]byte("1 1 1\r\ntest")) + name3 := output.file.Name() + + if name2 != name1 { + t.Errorf("Fast changes should happen in same file:", name1, name2) + } + + if name3 == name1 { + t.Errorf("File name should change:", name1, name3) + } + + os.Remove(name1) + os.Remove(name3) +} diff --git a/plugins.go b/plugins.go index d09bfe8..b593d78 100644 --- a/plugins.go +++ b/plugins.go @@ -117,7 +117,7 @@ func InitPlugins() { } 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..2296959 100644 --- a/settings.go +++ b/settings.go @@ -39,8 +39,9 @@ type AppSettings struct { outputTCP MultiOption outputTCPStats bool - inputFile MultiOption - outputFile MultiOption + inputFile MultiOption + outputFile MultiOption + outputFileFlushInterval time.Duration inputRAW MultiOption inputRAWEngine string @@ -84,6 +85,7 @@ func init() { flag.Var(&Settings.inputFile, "input-file", "Read requests from file: \n\tgor --input-file ./requests.gor --output-http staging.com") 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")