Skip to content

Commit

Permalink
parse buffer from different bases and data unit: (#754)
Browse files Browse the repository at this point in the history
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)
  • Loading branch information
Urban Ishimwe authored May 26, 2020
1 parent 1a5ac70 commit 7e90a73
Show file tree
Hide file tree
Showing 9 changed files with 148 additions and 88 deletions.
5 changes: 3 additions & 2 deletions input_raw.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ type RAWInput struct {
listener *raw.Listener
bpfFilter string
timestampType string
bufferSize int
bufferSize int64
}

// Available engines for intercepting traffic
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down
8 changes: 5 additions & 3 deletions output_file.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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")
}

Expand Down Expand Up @@ -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") {
Expand Down
40 changes: 0 additions & 40 deletions output_file_settings.go

This file was deleted.

34 changes: 19 additions & 15 deletions output_file_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
}
Expand Down Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion output_null.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ package main
type NullOutput struct {
}

// NullOutput constructor for NullOutput
// NewNullOutput constructor for NullOutput
func NewNullOutput() (o *NullOutput) {
return new(NullOutput)
}
Expand Down
1 change: 1 addition & 0 deletions protocol.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"strconv"
)

// These constants help to indicate the type of payload
const (
RequestPayload = '1'
ResponsePayload = '2'
Expand Down
15 changes: 8 additions & 7 deletions raw_socket_listener/listener.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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 {
Expand Down
Loading

0 comments on commit 7e90a73

Please sign in to comment.