diff --git a/capture/capture.go b/capture/capture.go index b1aaa10..3fb0abe 100644 --- a/capture/capture.go +++ b/capture/capture.go @@ -18,6 +18,7 @@ import ( "github.com/google/gopacket" "github.com/google/gopacket/layers" "github.com/google/gopacket/pcap" + "github.com/pbnjay/memory" "golang.org/x/sys/unix" ) @@ -50,7 +51,8 @@ type Listener struct { port uint16 // src or/and dst port trackResponse bool - host string // pcap file name or interface (name, hardware addr, index or ip address) + host string // pcap file name or interface (name, hardware addr, index or ip address) + maxAllowedMem int closeDone chan struct{} quit chan struct{} @@ -98,7 +100,7 @@ func (eng *EngineType) String() (e string) { // NewListener creates and initialize a new Listener. if transport or/and engine are invalid/unsupported // is "tcp" and "pcap", are assumed. l.Engine and l.Transport can help to get the values used. // if there is an error it will be associated with getting network interfaces -func NewListener(host string, port uint16, transport string, engine EngineType, trackResponse bool) (l *Listener, err error) { +func NewListener(host string, port uint16, transport string, engine EngineType, trackResponse bool, maxAllowedMem int) (l *Listener, err error) { l = &Listener{} l.host = host @@ -112,6 +114,7 @@ func NewListener(host string, port uint16, transport string, engine EngineType, l.closeDone = make(chan struct{}) l.quit = make(chan struct{}) l.Reading = make(chan bool) + l.maxAllowedMem = maxAllowedMem switch engine { default: l.Engine = EnginePcap @@ -326,12 +329,13 @@ func (l *Listener) read(handler PacketHandler) { case <-l.quit: return case <-ticker.C: - usedMemory, totalMemory := memUsage() - if 100*float64(usedMemory)/float64(totalMemory) >= 95 { - fmt.Println("Using more memory then allowed. Enabling throttling.") - time.Sleep(time.Second) + if l.maxAllowedMem > 0 { + usedMemory, _ := memUsage() + if int(usedMemory) > l.maxAllowedMem { + log.Printf("[WARNING] Using more memory (%vmb) than allowed (%vmb). Enabling throttling for 1 second\n", bToMb(usedMemory), bToMb(uint64(l.maxAllowedMem))) + time.Sleep(time.Second) + } } - continue default: data, ci, err := hndl.ZeroCopyReadPacketData() if err == nil { @@ -520,3 +524,23 @@ func pcapLinkTypeLength(lType int) (int, bool) { return 0, false } } + +func memUsage() (uint64, uint64) { + var m runtime.MemStats + runtime.ReadMemStats(&m) + // fmt.Printf("Current memory usage: ") + // // For info on each, see: https://golang.org/pkg/runtime/#MemStats + // fmt.Printf("Alloc = %v MiB", bToMb(m.Alloc)) + // fmt.Printf("\tTotalAlloc = %v MiB", bToMb(m.TotalAlloc)) + // fmt.Printf("\tSys = %v MiB", bToMb(m.Sys)) + // fmt.Printf("\tNumGC = %v\n", m.NumGC) + + // fmt.Printf("Total available memory %v MiB\n", bToMb(memory.TotalMemory())) + // fmt.Printf("Currently using %.2f %% of memory\n", 100*float64(m.Alloc+m.Sys)/float64(memory.TotalMemory())) + + return m.Alloc + m.Sys, memory.TotalMemory() +} + +func bToMb(b uint64) uint64 { + return b / 1024 / 1024 +} diff --git a/emitter.go b/emitter.go index 3db337c..cb6997b 100644 --- a/emitter.go +++ b/emitter.go @@ -5,10 +5,13 @@ import ( "hash/fnv" "io" "log" + "runtime" "sync" "time" "github.com/buger/goreplay/byteutils" + + "github.com/pbnjay/memory" ) // Emitter represents an abject to manage plugins communication @@ -60,6 +63,7 @@ func (e *Emitter) Start(plugins *InOutPlugins, middlewareCmd string) { // Close closes all the goroutine and waits for it to finish. func (e *Emitter) Close() { + fmt.Println("Closing!") for _, p := range e.plugins.All { if cp, ok := p.(io.Closer); ok { cp.Close() @@ -69,6 +73,8 @@ func (e *Emitter) Close() { // wait for everything to stop e.Wait() } + + fmt.Println("Closed!") e.plugins.All = nil // avoid Close to make changes again } @@ -88,6 +94,7 @@ func CopyMulty(src PluginReader, writers ...PluginWriter) error { } return err } + if msg != nil && len(msg.Data) > 0 { if len(msg.Data) > int(Settings.CopyBufferSize) { msg.Data = msg.Data[:Settings.CopyBufferSize] @@ -175,3 +182,23 @@ func CopyMulty(src PluginReader, writers ...PluginWriter) error { } } } + +func memUsage() (uint64, uint64) { + var m runtime.MemStats + runtime.ReadMemStats(&m) + // fmt.Printf("Current memory usage: ") + // // For info on each, see: https://golang.org/pkg/runtime/#MemStats + // fmt.Printf("Alloc = %v MiB", bToMb(m.Alloc)) + // fmt.Printf("\tTotalAlloc = %v MiB", bToMb(m.TotalAlloc)) + // fmt.Printf("\tSys = %v MiB", bToMb(m.Sys)) + // fmt.Printf("\tNumGC = %v\n", m.NumGC) + + // fmt.Printf("Total available memory %v MiB\n", bToMb(memory.TotalMemory())) + // fmt.Printf("Currently using %.2f %% of memory\n", 100*float64(m.Alloc+m.Sys)/float64(memory.TotalMemory())) + + return m.Alloc + m.Sys, memory.TotalMemory() +} + +func bToMb(b uint64) uint64 { + return b / 1024 / 1024 +} diff --git a/input_raw.go b/input_raw.go index 7d73aff..29c4b2f 100644 --- a/input_raw.go +++ b/input_raw.go @@ -141,7 +141,7 @@ func (i *RAWInput) PluginRead() (*Message, error) { func (i *RAWInput) listen(address string) { var err error - i.listener, err = capture.NewListener(i.host, i.port, "", i.Engine, i.TrackResponse) + i.listener, err = capture.NewListener(i.host, i.port, "", i.Engine, i.TrackResponse, Settings.maxAllowedMem) if err != nil { log.Fatal(err) } diff --git a/settings.go b/settings.go index 928de3b..4deb5fd 100644 --- a/settings.go +++ b/settings.go @@ -4,8 +4,12 @@ import ( "flag" "fmt" "os" + "strconv" + "strings" "sync" "time" + + "github.com/buger/goreplay/size" ) // DEMO indicates that goreplay is running in demo mode @@ -59,6 +63,9 @@ type AppSettings struct { OutputHTTP MultiOption `json:"output-http"` PrettifyHTTP bool `json:"prettify-http"` + MaxAllowedMem string `json:"max-allowed-mem"` + maxAllowedMem int `json:"-"` + OutputHTTPConfig HTTPOutputConfig OutputBinary MultiOption `json:"output-binary"` @@ -142,6 +149,8 @@ func init() { flag.BoolVar(&Settings.Monitor, "input-raw-monitor", false, "enable RF monitor mode") flag.BoolVar(&Settings.Stats, "input-raw-stats", false, "enable stats generator on raw TCP messages") + flag.StringVar(&Settings.MaxAllowedMem, "max-allowed-mem", "", "Allow specify memory which can be consumed by the process. Example: --max-allowed-mem 10%, --max-allowed-mem 200m") + flag.StringVar(&Settings.Middleware, "middleware", "", "Used for modifying traffic using external command") flag.Var(&Settings.OutputHTTP, "output-http", "Forwards incoming requests to given http address.\n\t# Redirect all incoming requests to staging.com address \n\tgor --input-raw :80 --output-http http://staging.com") @@ -217,6 +226,21 @@ func checkSettings() { if Settings.CopyBufferSize < 1 { Settings.CopyBufferSize.Set("5mb") } + + if Settings.MaxAllowedMem != "" { + _, totalMemory := memUsage() + fmt.Println("Total available memory:", bToMb(totalMemory), "Mb") + if strings.HasSuffix(Settings.MaxAllowedMem, "%") { + pct, _ := strconv.Atoi(strings.TrimSuffix(Settings.MaxAllowedMem, "%")) + Settings.maxAllowedMem = int((float64(pct) / 100) * float64(totalMemory)) + } else { + var s size.Size + s.Set(Settings.MaxAllowedMem) + Settings.maxAllowedMem = int(s) + } + + fmt.Println("Limiting memory usage to:", bToMb(uint64(Settings.maxAllowedMem)), "Mb") + } } var previousDebugTime = time.Now()