diff --git a/input_raw.go b/input_raw.go index d022c731..48089459 100644 --- a/input_raw.go +++ b/input_raw.go @@ -21,7 +21,7 @@ type RAWInput struct { listener *raw.Listener bpfFilter string timestampType string - bufferSize int + bufferSize int64 } // Available engines for intercepting traffic @@ -32,7 +32,7 @@ const ( ) // NewRAWInput constructor for RAWInput. Accepts address with port as argument. -func NewRAWInput(address string, engine int, trackResponse bool, expire time.Duration, realIPHeader string, bpfFilter string, timestampType string, bufferSize int) (i *RAWInput) { +func NewRAWInput(address string, engine int, trackResponse bool, expire time.Duration, realIPHeader string, bpfFilter string, timestampType string, bufferSize int64) (i *RAWInput) { i = new(RAWInput) i.data = make(chan *raw.TCPMessage) i.address = address @@ -105,6 +105,7 @@ func (i *RAWInput) String() string { return "Intercepting traffic from: " + i.address } +// Close closes the input raw listener func (i *RAWInput) Close() error { i.listener.Close() close(i.quit) diff --git a/output_file.go b/output_file.go index de97186c..60581158 100644 --- a/output_file.go +++ b/output_file.go @@ -29,10 +29,11 @@ var dateFileNameFuncs = map[string]func(*FileOutput) string{ "%t": func(o *FileOutput) string { return string(o.payloadType) }, } +// FileOutputConfig ... type FileOutputConfig struct { flushInterval time.Duration - sizeLimit unitSizeVar - outputFileMaxSize unitSizeVar + sizeLimit int64 + outputFileMaxSize int64 queueLimit int append bool } @@ -219,7 +220,7 @@ func (o *FileOutput) Write(data []byte) (n int, err error) { o.totalFileSize += int64(len(data) + len(payloadSeparator)) o.queueLength++ - if Settings.outputFileConfig.outputFileMaxSize > 0 && o.totalFileSize >= int64(Settings.outputFileConfig.outputFileMaxSize) { + if Settings.outputFileConfig.outputFileMaxSize > 0 && o.totalFileSize >= Settings.outputFileConfig.outputFileMaxSize { return len(data), errors.New("File output reached size limit") } @@ -256,6 +257,7 @@ func (o *FileOutput) String() string { return "File output: " + o.file.Name() } +// Close closes the output file func (o *FileOutput) Close() error { if o.file != nil { if strings.HasSuffix(o.currentName, ".gz") { diff --git a/output_file_settings.go b/output_file_settings.go deleted file mode 100644 index 58cf656d..00000000 --- a/output_file_settings.go +++ /dev/null @@ -1,40 +0,0 @@ -package main - -import ( - "strconv" - "strings" -) - -var dataUnitMap = map[byte]int64{ - 'k': 1024, - 'm': 1024 * 1024, - 'g': 1024 * 1024 * 1024, -} - -func parseDataUnit(s string) int64 { - // Allow kb, mb, gb - if strings.HasSuffix(s, "b") { - s = s[:len(s)-1] - } - - if unit, ok := dataUnitMap[s[len(s)-1]]; ok { - size, _ := strconv.ParseInt(s[:len(s)-1], 10, 64) - return unit * size - } - - // If no unit specified use bytes - size, _ := strconv.ParseInt(s, 10, 64) - return size - -} - -type unitSizeVar int64 - -func (u unitSizeVar) String() string { - return strconv.Itoa(int(u)) -} - -func (u *unitSizeVar) Set(s string) error { - *u = unitSizeVar(parseDataUnit(s)) - return nil -} diff --git a/output_file_test.go b/output_file_test.go index 3dfa564a..604af29c 100644 --- a/output_file_test.go +++ b/output_file_test.go @@ -158,21 +158,25 @@ func TestFileOutputCompression(t *testing.T) { } func TestParseDataUnit(t *testing.T) { - var tests = []struct { - value string - size int64 - }{ - {"100kb", dataUnitMap['k'] * 100}, - {"100k", dataUnitMap['k'] * 100}, - {"1kb", dataUnitMap['k']}, - {"1g", dataUnitMap['g']}, - {"10m", dataUnitMap['m'] * 10}, - {"zsaa312", 0}, + var d = map[string]int64{ + "42mb": 42 << 20, + "4_2": 42, + "00": 0, + "\n\n 0.0\r\t\f": 0, + "0_600tb": 384 << 40, + "0600Tb": 384 << 40, + "0o12Mb": 10 << 20, + "0b_10010001111_1kb": 2335 << 10, + "1024": 1 << 10, + "0b111": 7, + "0x12gB": 18 << 30, + "0x_67_7a_2f_cc_40_c6": 113774485586118, + "121562380192901": 121562380192901, } - - for _, c := range tests { - if parseDataUnit(c.value) != c.size { - t.Error(c.value, "should be", c.size, "instead", parseDataUnit(c.value)) + for k, v := range d { + n, err := bufferParser(k, "0") + if err != nil || n != v { + t.Errorf("Error parsing %s: %v", k, err) } } } @@ -321,7 +325,7 @@ func TestFileOutputAppendSizeLimitOverflow(t *testing.T) { messageSize := len(message) + len(payloadSeparator) - output := NewFileOutput(name, &FileOutputConfig{append: false, flushInterval: time.Minute, sizeLimit: unitSizeVar(2 * messageSize)}) + output := NewFileOutput(name, &FileOutputConfig{append: false, flushInterval: time.Minute, sizeLimit: 2 * int64(messageSize)}) output.Write([]byte("1 1 1\r\ntest")) name1 := output.file.Name() diff --git a/output_null.go b/output_null.go index 4b756db0..ee5fdaab 100644 --- a/output_null.go +++ b/output_null.go @@ -4,7 +4,7 @@ package main type NullOutput struct { } -// NullOutput constructor for NullOutput +// NewNullOutput constructor for NullOutput func NewNullOutput() (o *NullOutput) { return new(NullOutput) } diff --git a/protocol.go b/protocol.go index 63c4560a..b7851e36 100644 --- a/protocol.go +++ b/protocol.go @@ -7,6 +7,7 @@ import ( "strconv" ) +// These constants help to indicate the type of payload const ( RequestPayload = '1' ResponsePayload = '2' diff --git a/raw_socket_listener/listener.go b/raw_socket_listener/listener.go index 0d5fa6d0..97f65d21 100644 --- a/raw_socket_listener/listener.go +++ b/raw_socket_listener/listener.go @@ -72,12 +72,12 @@ type Listener struct { trackResponse bool messageExpire time.Duration - bpfFilter string - timestampType string + bpfFilter string + timestampType string overrideSnapLen bool - immediateMode bool + immediateMode bool - bufferSize int + bufferSize int64 conn net.PacketConn pcapHandles []*pcap.Handle @@ -100,7 +100,7 @@ const ( ) // NewListener creates and initializes new Listener object -func NewListener(addr string, port string, engine int, trackResponse bool, expire time.Duration, bpfFilter string, timestampType string, bufferSize int, overrideSnapLen bool, immediateMode bool) (l *Listener) { +func NewListener(addr string, port string, engine int, trackResponse bool, expire time.Duration, bpfFilter string, timestampType string, bufferSize int64, overrideSnapLen bool, immediateMode bool) (l *Listener) { l = &Listener{} l.packetsChan = make(chan *packet, 10000) @@ -367,7 +367,7 @@ func (t *Listener) readPcap() { log.Println("Setting immediate mode") } if t.bufferSize > 0 { - inactive.SetBufferSize(t.bufferSize) + inactive.SetBufferSize(int(t.bufferSize)) } handle, herr := inactive.Activate() @@ -889,11 +889,12 @@ func (t *Listener) IsReady() bool { } } -// Receive TCP messages from the listener channel +// Receiver TCP messages from the listener channel func (t *Listener) Receiver() chan *TCPMessage { return t.messagesChan } +// Close tcp listener func (t *Listener) Close() { close(t.quit) if t.conn != nil { diff --git a/settings.go b/settings.go index f4865858..ea48f0ff 100644 --- a/settings.go +++ b/settings.go @@ -3,13 +3,14 @@ package main import ( "flag" "fmt" + "log" "os" + "regexp" + "strconv" "sync" "time" ) -var VERSION string - // MultiOption allows to specify multiple flags with same name and collects all values into array type MultiOption []string @@ -57,9 +58,9 @@ type AppSettings struct { inputRAWExpire time.Duration inputRAWBpfFilter string inputRAWTimestampType string - copyBufferSize int + copyBufferSize int64 inputRAWImmediateMode bool - inputRawBufferSize int + inputRawBufferSize int64 inputRAWOverrideSnapLen bool middleware string @@ -80,13 +81,16 @@ type AppSettings struct { var Settings AppSettings func usage() { - fmt.Printf("Gor is a simple http traffic replication tool written in Go. Its main goal is to replay traffic from production servers to staging and dev environments.\nProject page: https://github.com/buger/gor\nAuthor: leonsbox@gmail.com\nCurrent Version: %s\n\n", VERSION) + fmt.Printf("Gor is a simple http traffic replication tool written in Go. Its main goal is to replay traffic from production servers to staging and dev environments.\nProject page: https://github.com/buger/gor\nAuthor: leonsbox@gmail.com\nCurrent Version: v%s\n\n", VERSION) flag.PrintDefaults() os.Exit(2) } func init() { flag.Usage = usage + var ( + inputRawBufferSize, outputFileMaxSize, copyBufferSize, outputFileSize string + ) flag.StringVar(&Settings.pprof, "http-pprof", "", "Enable profiling. Starts http server on specified port, exposing special /debug/pprof endpoint. Example: `:8181`") flag.BoolVar(&Settings.verbose, "verbose", false, "Turn on more verbose output") @@ -118,13 +122,23 @@ func init() { flag.Var(&Settings.outputFile, "output-file", "Write incoming requests to file: \n\tgor --input-raw :80 --output-file ./requests.gor") flag.DurationVar(&Settings.outputFileConfig.flushInterval, "output-file-flush-interval", time.Second, "Interval for forcing buffer flush to the file, default: 1s.") flag.BoolVar(&Settings.outputFileConfig.append, "output-file-append", false, "The flushed chunk is appended to existence file or not. ") - - // Set default - Settings.outputFileConfig.sizeLimit.Set("32mb") - flag.Var(&Settings.outputFileConfig.sizeLimit, "output-file-size-limit", "Size of each chunk. Default: 32mb") + flag.StringVar(&outputFileSize, "output-file-size-limit", "32mb", "Size of each chunk. Default: 32mb") + { + n, err := bufferParser(outputFileSize, "32MB") + if err != nil { + log.Fatalf("output-file-size-limit error: %v\n", err) + } + Settings.outputFileConfig.sizeLimit = n + } flag.IntVar(&Settings.outputFileConfig.queueLimit, "output-file-queue-limit", 256, "The length of the chunk queue. Default: 256") - Settings.outputFileConfig.outputFileMaxSize.Set("-1") - flag.Var(&Settings.outputFileConfig.outputFileMaxSize, "output-file-max-size-limit", "Max size of output file, Default: 1TB") + flag.StringVar(&outputFileMaxSize, "output-file-max-size-limit", "1TB", "Max size of output file, Default: 1TB") + { + n, err := bufferParser(outputFileMaxSize, "1TB") + if err != nil { + log.Fatalf("output-file-max-size-limit error: %v\n", err) + } + Settings.outputFileConfig.outputFileMaxSize = n + } flag.BoolVar(&Settings.prettifyHTTP, "prettify-http", false, "If enabled, will automatically decode requests and responses with: Content-Encodning: gzip and Transfer-Encoding: chunked. Useful for debugging, in conjuction with --output-stdout") @@ -141,11 +155,25 @@ func init() { flag.StringVar(&Settings.inputRAWBpfFilter, "input-raw-bpf-filter", "", "BPF filter to write custom expressions. Can be useful in case of non standard network interfaces like tunneling or SPAN port. Example: --input-raw-bpf-filter 'dst port 80'") flag.StringVar(&Settings.inputRAWTimestampType, "input-raw-timestamp-type", "", "Possible values: PCAP_TSTAMP_HOST, PCAP_TSTAMP_HOST_LOWPREC, PCAP_TSTAMP_HOST_HIPREC, PCAP_TSTAMP_ADAPTER, PCAP_TSTAMP_ADAPTER_UNSYNCED. This values not supported on all systems, GoReplay will tell you available values of you put wrong one.") - flag.IntVar(&Settings.copyBufferSize, "copy-buffer-size", 5*1024*1024, "Set the buffer size for an individual request (default 5M)") + flag.StringVar(©BufferSize, "copy-buffer-size", "5mb", "Set the buffer size for an individual request (default 5MB)") + { + n, err := bufferParser(copyBufferSize, "5mb") + if err != nil { + log.Fatalf("copy-buffer-size error: %v\n", err) + } + Settings.copyBufferSize = n + } flag.BoolVar(&Settings.inputRAWOverrideSnapLen, "input-raw-override-snaplen", false, "Override the capture snaplen to be 64k. Required for some Virtualized environments") flag.BoolVar(&Settings.inputRAWImmediateMode, "input-raw-immediate-mode", false, "Set pcap interface to immediate mode.") - flag.IntVar(&Settings.inputRawBufferSize, "input-raw-buffer-size", 0, "Controls size of the OS buffer (in bytes) which holds packets until they dispatched. Default value depends by system: in Linux around 2MB. If you see big package drop, increase this value.") + flag.StringVar(&inputRawBufferSize, "input-raw-buffer-size", "", "Controls size of the OS buffer which holds packets until they dispatched. Default value depends by system: in Linux around 2MB. If you see big package drop, increase this value.") + { + n, err := bufferParser(inputRawBufferSize, "0") + if err != nil { + log.Fatalf("input-raw-buffer-size error: %v\n", err) + } + Settings.inputRawBufferSize = n + } flag.StringVar(&Settings.middleware, "middleware", "", "Used for modifying traffic using external command") @@ -210,19 +238,78 @@ func init() { flag.Var(&Settings.modifierConfig.paramHashFilters, "http-param-limiter", "Takes a fraction of requests, consistently taking or rejecting a request based on the FNV32-1A hash of a specific GET param:\n\t gor --input-raw :8080 --output-http staging.com --http-param-limiter user_id:25%") } -var previousDebugTime int64 +var previousDebugTime = time.Now() var debugMutex sync.Mutex +var pID = os.Getpid() -// Debug gets called only if --verbose flag specified +// Debug take an effect only if --verbose flag specified func Debug(args ...interface{}) { if Settings.verbose { debugMutex.Lock() + defer debugMutex.Unlock() now := time.Now() - diff := float64(now.UnixNano()-previousDebugTime) / 1000000 - previousDebugTime = now.UnixNano() - debugMutex.Unlock() - - fmt.Printf("[DEBUG][PID %d][%d][%fms] ", os.Getpid(), now.UnixNano(), diff) + diff := now.Sub(previousDebugTime).String() + previousDebugTime = now + fmt.Printf("[DEBUG][PID %d][%s][elapsed %s] ", pID, now.Format(time.StampNano), diff) fmt.Println(args...) } } + +// bufferParser parses buffer to bytes from different bases and data units +// size is the buffer in string, rpl act as a replacement for empty buffer. +// e.g: (--output-file-size-limit "") may override default 32mb with empty buffer, +// which can be solved by setting rpl by bufferParser(buffer, "32mb") +func bufferParser(size, rpl string) (buffer int64, err error) { + const ( + _ = 1 << (iota * 10) + KB + MB + GB + TB + ) + + var ( + // the following regexes follow Go semantics https://golang.org/ref/spec#Letters_and_digits + rB = regexp.MustCompile(`(?i)^(?:0b|0x|0o)?[\da-f_]+$`) + rKB = regexp.MustCompile(`(?i)^(?:0b|0x|0o)?[\da-f_]+kb$`) + rMB = regexp.MustCompile(`(?i)^(?:0b|0x|0o)?[\da-f_]+mb$`) + rGB = regexp.MustCompile(`(?i)^(?:0b|0x|0o)?[\da-f_]+gb$`) + rTB = regexp.MustCompile(`(?i)^(?:0b|0x|0o)?[\da-f_]+tb$`) + empt = regexp.MustCompile(`^[\n\t\r 0.\f\a]*$`) + + lmt = len(size) - 2 + s = []byte(size) + ) + + if empt.Match(s) { + size = rpl + s = []byte(size) + } + + // recover, especially when buffer size overflows int64 i.e ~8019PBs + defer func() { + if e, ok := recover().(error); ok { + err = e.(error) + } + }() + + switch { + case rB.Match(s): + buffer, err = strconv.ParseInt(size, 0, 64) + case rKB.Match(s): + buffer, err = strconv.ParseInt(size[:lmt], 0, 64) + buffer *= KB + case rMB.Match(s): + buffer, err = strconv.ParseInt(size[:lmt], 0, 64) + buffer *= MB + case rGB.Match(s): + buffer, err = strconv.ParseInt(size[:lmt], 0, 64) + buffer *= GB + case rTB.Match(s): + buffer, err = strconv.ParseInt(size[:lmt], 0, 64) + buffer *= TB + default: + return 0, fmt.Errorf("invalid buffer %q", size) + } + return +} diff --git a/version.go b/version.go new file mode 100644 index 00000000..b7942301 --- /dev/null +++ b/version.go @@ -0,0 +1,4 @@ +package main + +// VERSION the current version of goreplay +const VERSION = "1.0.0"