parse buffer from different bases and data unit: (#754)

All flags that expect buffer as input i.e. `--output-file-size-limit`, `--output-file-max-size-limit`, `--copy-buffer-size` and `input-raw-buffer-size` can now parse  inputs from differents bases and data units like: `10mb`, `10kb`, `100gb`,   `18tb`, `11839023`.... 
data units and bases are case insensitive, the parser accepts only the format of [Go integer literals](https://golang.org/ref/spec#Integer_literals)
This commit is contained in:
Urban Ishimwe
2020-05-26 18:05:59 +03:00
committed by GitHub
parent 1a5ac70c24
commit 7e90a73b81
9 changed files with 148 additions and 88 deletions
+3 -2
View File
@@ -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)
+5 -3
View File
@@ -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") {
-40
View File
@@ -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
}
+19 -15
View File
@@ -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()
+1 -1
View File
@@ -4,7 +4,7 @@ package main
type NullOutput struct {
}
// NullOutput constructor for NullOutput
// NewNullOutput constructor for NullOutput
func NewNullOutput() (o *NullOutput) {
return new(NullOutput)
}
+1
View File
@@ -7,6 +7,7 @@ import (
"strconv"
)
// These constants help to indicate the type of payload
const (
RequestPayload = '1'
ResponsePayload = '2'
+8 -7
View File
@@ -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 {
+107 -20
View File
@@ -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: <Leonid Bugaev> 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: <Leonid Bugaev> 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(&copyBufferSize, "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
}
+4
View File
@@ -0,0 +1,4 @@
package main
// VERSION the current version of goreplay
const VERSION = "1.0.0"