map sizes can be specified through flags

3 files changed, 37 insertions, 23 deletions

diff --git a/internal/flags/flags.go b/internal/flags/flags.go
index 674f22f..1ba9d60 100644
--- a/internal/flags/flags.go
+++ b/internal/flags/flags.go
@@ -2,17 +2,22 @@ package flags
 
 import (
 	"flag"
+	"fmt"
 	"unsafe"
 
 	bpf "github.com/aquasecurity/libbpfgo"
 )
 
 type Flags struct {
-	UidFilter int
+	UidFilter        int
+	FdEventMapSize   int
+	OpenEventMapSize int
 }
 
 func New() (flags Flags) {
 	flag.IntVar(&flags.UidFilter, "uid", 0, "Filter for processes with UID")
+	flag.IntVar(&flags.FdEventMapSize, "fdMapSize", 4096, "BPF FD event map size")
+	flag.IntVar(&flags.OpenEventMapSize, "openMapSize", 1024, "BPF open event map size")
 	flag.Parse()
 	return flags
 }
@@ -32,3 +37,30 @@ func (flags Flags) SetBPF(bpfModule *bpf.Module) error {
 	key := uint32(1)
 	return flagsMap.Update(unsafe.Pointer(&key), unsafe.Pointer(&flagsValues))
 }
+
+func (flags Flags) ResizeBPFMaps(bpfModule *bpf.Module) error {
+	if err := resizeBPFMap(bpfModule, "open_event_map", uint32(flags.OpenEventMapSize)); err != nil {
+		return err
+	}
+	if err := resizeBPFMap(bpfModule, "fd_event_map", uint32(flags.FdEventMapSize)); err != nil {
+		return err
+	}
+	return nil
+}
+
+func resizeBPFMap(module *bpf.Module, name string, size uint32) error {
+	m, err := module.GetMap("open_event_map")
+	if err != nil {
+		return err
+	}
+
+	if err = m.SetMaxEntries(size); err != nil {
+		return err
+	}
+
+	if actual := m.MaxEntries(); actual != size {
+		return fmt.Errorf("map resize failed, expected %v, actual %v", size, actual)
+	}
+
+	return nil
+}
diff --git a/internal/ioriotng.go b/internal/ioriotng.go
index 890ac73..d840ea5 100644
--- a/internal/ioriotng.go
+++ b/internal/ioriotng.go
@@ -23,23 +23,6 @@ type BpfMapper interface {
 	String() string
 }
 
-func resizeMap(module *bpf.Module, name string, size uint32) error {
-	m, err := module.GetMap("open_event_map")
-	if err != nil {
-		return err
-	}
-
-	if err = m.SetMaxEntries(size); err != nil {
-		return err
-	}
-
-	if actual := m.MaxEntries(); actual != size {
-		return fmt.Errorf("map resize failed, expected %v, actual %v", size, actual)
-	}
-
-	return nil
-}
-
 func Run(flags flags.Flags) {
 	// To consider for implementation!
 	log.Println(debugfs.TracepointsWithFd())
@@ -50,10 +33,7 @@ func Run(flags flags.Flags) {
 	}
 	defer bpfModule.Close()
 
-	if err = resizeMap(bpfModule, "open_event_map", 8192); err != nil {
-		log.Fatal(err)
-	}
-	if err = resizeMap(bpfModule, "fd_event_map", 8192); err != nil {
+	if err := flags.ResizeBPFMaps(bpfModule); err != nil {
 		log.Fatal(err)
 	}
 
diff --git a/internal/types/types.go b/internal/types/types.go
index c9e2101..6e8cd3b 100644
--- a/internal/types/types.go
+++ b/internal/types/types.go
@@ -5,6 +5,7 @@ import "fmt"
 type OpenEvent struct {
 	FD       int32
 	TID      uint32
+	Time     uint64
 	Filename [256]byte // TODO, use same value as in ioriot.bpf.h
 	Comm     [16]byte
 }
@@ -12,7 +13,8 @@ type OpenEvent struct {
 func (e OpenEvent) String() string {
 	filename := e.Filename[:]
 	comm := e.Comm[:]
-	return fmt.Sprintf("tid:%d fd:%d filename:%s, comm:%s", e.TID, e.FD, string(filename), string(comm))
+	return fmt.Sprintf("%v tid:%d fd:%d filename:%s, comm:%s",
+		e.Time, e.TID, e.FD, string(filename), string(comm))
 }
 
 type FdEvent struct {