Merge pull request #5 from ssst0n3/v0.4.2

v0.4.2: fd watcher: walk driver

go.mod

@@ -5,7 +5,7 @@ go 1.21.6
 require (
 	github.com/ctrsploit/sploit-spec v0.4.3
 	github.com/fatih/color v1.15.0
-	github.com/slimtoolkit/slim v0.0.0-20240505211522-04aa51412d9f
+	github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826
 	github.com/ssst0n3/awesome_libs v0.6.7
 	github.com/urfave/cli/v2 v2.27.1
 	github.com/vishvananda/netlink v1.2.1-beta.2
@@ -18,6 +18,7 @@ require (
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	github.com/sirupsen/logrus v1.9.3 // indirect
+	github.com/stretchr/testify v1.8.4 // indirect
 	github.com/vishvananda/netns v0.0.0-20200728191858-db3c7e526aae // indirect
 	github.com/xrash/smetrics v0.0.0-20231213231151-1d8dd44e695e // indirect
 	golang.org/x/sys v0.15.0 // indirect

go.sum

@@ -26,6 +26,8 @@ github.com/mattn/go-isatty v0.0.17 h1:BTarxUcIeDqL27Mc+vyvdWYSL28zpIhv3RoTdsLMPn
 github.com/mattn/go-isatty v0.0.17/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
 github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
 github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
+github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 h1:RWengNIwukTxcDr9M+97sNutRR1RKhG96O6jWumTTnw=
+github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826/go.mod h1:TaXosZuwdSHYgviHp1DAtfrULt5eUgsSMsZf+YrPgl8=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
@@ -35,8 +37,6 @@ github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQD
 github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
 github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
 github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
-github.com/slimtoolkit/slim v0.0.0-20240505211522-04aa51412d9f h1:CzxuqUzdfDc7ycF4FJqL5x3Dn6Ay3+3CsDVeH0TYwRU=
-github.com/slimtoolkit/slim v0.0.0-20240505211522-04aa51412d9f/go.mod h1:55lB/fY9qINNkJDRnG2qWdU1/B99MazwBBMhz0sdSnc=
 github.com/ssst0n3/awesome_libs v0.6.7 h1:CxhRcWy/v1THpy4Jk8AJh4jMnfYE+5iyrkfV5HL5nYQ=
 github.com/ssst0n3/awesome_libs v0.6.7/go.mod h1:+JQKcgjs0TWWFszGXRzIs+8pZaL0qRf6HEhPY5n8cEk=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=

pkg/listener/fd/ebpf/watcher.go

@@ -5,7 +5,8 @@ import (
 )
 
 type Watcher struct {
-	msg chan Message
+	msg  chan Message
+	stop chan struct{}
 }
 
 func New(pid int) *Watcher {
@@ -16,33 +17,38 @@ func New(pid int) *Watcher {
 
 func (w Watcher) Init() {
 	w.register()
-	// receive IPC message
 	go w.receive()
 }
 
 func (w Watcher) Enable() (enabled bool) {
 	return
 }
 
+// register to kernel
 func (w Watcher) register() {
 	// register with os.GetPid()
 }
 
+// receive IPC message
 func (w Watcher) receive() {
 	for {
 		w.msg <- Message{}
 	}
 }
 
-func (w Watcher) Watch(stop <-chan struct{}, event chan<- event.Events) {
+func (w Watcher) Watch(event chan<- event.Events) {
 	select {
-	case <-stop:
+	case <-w.stop:
 		return
 	default:
 		w.do(event)
 	}
 }
 
+func (w Watcher) Close() {
+	w.stop <- struct{}{}
+}
+
 func (w Watcher) do(e chan<- event.Events) {
 	for {
 		msg := <-w.msg

pkg/listener/fd/event/event.go

@@ -13,7 +13,20 @@ type Event struct {
 	Type int
 	Pid  int
 	Fd   int
-	Stat stat.Stat
+	Stat *stat.Stat
 }
 
 type Events []Event
+
+func (e Events) Map(m map[int]*stat.Stat) {
+	for _, event := range e {
+		switch event.Type {
+		case Open, Change:
+			m[event.Fd] = event.Stat
+		case Close:
+			delete(m, event.Fd)
+		default:
+			panic("unhandled default case")
+		}
+	}
+}

pkg/listener/fd/listener.go

@@ -25,7 +25,7 @@ func (l *Listener) Handle() {
 		case event.ProcessExit:
 			worker := l.workers[e.Pid]
 			if worker != nil {
-				worker.Stop <- true
+				worker.Close()
 				delete(l.workers, e.Pid)
 			}
 		default:

pkg/listener/fd/stat/stat.go

@@ -13,15 +13,16 @@ type Stat struct {
 	SocketPath string
 	Leak       bool
 	Flags      int64
+	Changed    bool
 }
 
-func New(pid, fd int) (stat Stat, err error) {
+func New(pid, fd int) (stat *Stat, err error) {
 	fdPath := fmt.Sprintf("/proc/%d/fd/%d", pid, fd)
 	realPath, _ := os.Readlink(fdPath)
 	if realPath == "" {
 		realPath = "?"
 	}
-	stat = Stat{
+	stat = &Stat{
 		FdPath:   fdPath,
 		RealPath: realPath,
 	}
@@ -43,19 +44,32 @@ func New(pid, fd int) (stat Stat, err error) {
 	return
 }
 
-func (s Stat) String() (content string) {
+func (s *Stat) Change(c bool) {
+	s.Changed = c
+}
+
+func (s *Stat) String() (content string) {
+	fdPath := s.FdPath
+	if s.Changed {
+		fdPath = color.New(color.Underline).Sprint(fdPath)
+	}
 	leaked := ""
 	if s.Leak {
-		leaked = "leaked!"
+		leaked = color.RedString("leaked!")
 	}
 	flags := pkg.ParseFlags(s.Flags)
 	var socketPath string
 	if s.SocketPath != "" {
 		socketPath = " -> " + s.SocketPath
 	}
-	content = fmt.Sprintf("%s -> %s%s\t; %s\t%s", s.FdPath, s.RealPath, socketPath, leaked, flags)
-	if s.Leak {
-		content = color.RedString(content)
-	}
+	content = fmt.Sprintf("%s -> %s%s\t; %s\t%s", fdPath, s.RealPath, socketPath, leaked, flags)
 	return
 }
+
+func (s *Stat) Equals(s2 *Stat) bool {
+	return s.FdPath == s2.FdPath &&
+		s.RealPath == s2.RealPath &&
+		s.SocketPath == s2.SocketPath &&
+		s.Leak == s2.Leak &&
+		s.Flags == s2.Flags
+}

pkg/listener/fd/walk/watcher.go

@@ -11,46 +11,53 @@ import (
 )
 
 type Watcher struct {
-	store   *sync.Map
-	changed bool
-	event   event.Events
-	pid     int
-	max     int
+	pid   int
+	store sync.Map
+	event event.Events
+	max   int
+	stop  chan struct{}
 }
 
-func New(pid int, store *sync.Map) *Watcher {
-	return &Watcher{
+func New(pid int) (w *Watcher) {
+	w = &Watcher{
 		pid:   pid,
-		store: store,
+		store: sync.Map{},
+		event: event.Events{},
+		stop:  make(chan struct{}),
 	}
+	return
 }
 
-func (w Watcher) Enable() (enabled bool) {
+func (w *Watcher) Enable() (enabled bool) {
 	return true
 }
 
-func (w Watcher) Watch(stop <-chan struct{}, event chan<- event.Events) {
+func (w *Watcher) Watch(c chan<- event.Events) {
 	for {
 		select {
-		case <-stop:
+		case <-w.stop:
 			return
 		default:
-			_ = w.do(event)
-			if len(w.event) > 0 {
-				event <- w.event
+			e := event.Events{}
+			_ = w.do(&e)
+			if len(e) > 0 {
+				c <- e
 			}
 		}
 	}
 }
 
-func (w Watcher) do(e chan<- event.Events) (err error) {
-	w.changed = false
+func (w *Watcher) Close() {
+	w.stop <- struct{}{}
+}
+
+func (w *Watcher) do(e *event.Events) (err error) {
 	fds, err := fds(w.pid)
 	if err != nil {
 		return
 	}
-	w.close(fds)
-	w.openOrChange(fds)
+	w.closeFd(fds, e)
+	w.openOrChangeFd(fds, e)
 	return
 }
 
@@ -76,7 +83,10 @@ func fds(pid int) (fds []int, err error) {
 }
 
 // fds has been sorted
-func (w Watcher) close(fds []int) {
+func (w *Watcher) closeFd(fds []int, e *event.Events) {
+	if len(fds) == 0 {
+		return
+	}
 	m := max(w.max, fds[len(fds)-1])
 	var missing []int
 	for i := 0; i < len(fds)-1; i++ {
@@ -89,31 +99,35 @@ func (w Watcher) close(fds []int) {
 	}
 	for _, fd := range missing {
 		if _, ok := w.store.LoadAndDelete(fd); ok {
-			w.changed = true
-			w.event = append(w.event, event.Event{
+			*e = append(*e, event.Event{
 				Type: event.Close,
 				Pid:  w.pid,
 				Fd:   fd,
+				Stat: &stat.Stat{
+					FdPath:   fmt.Sprintf("/proc/%d/fd/%d", w.pid, fd),
+					RealPath: "[FD CLOSED]",
+					Changed:  true,
+				},
 			})
 		}
 	}
 	w.max = fds[len(fds)-1]
 	return
 }
 
-func (w Watcher) openOrChange(fds []int) {
+func (w *Watcher) openOrChangeFd(fds []int, e *event.Events) {
 	for _, fd := range fds {
 		s, _ := stat.New(w.pid, fd)
 		if old, ok := w.store.LoadOrStore(fd, s); !ok {
-			w.event = append(w.event, event.Event{
+			*e = append(*e, event.Event{
 				Type: event.Open,
 				Pid:  w.pid,
 				Fd:   fd,
 				Stat: s,
 			})
 		} else {
-			if old != s {
-				w.event = append(w.event, event.Event{
+			if !old.(*stat.Stat).Equals(s) {
+				*e = append(*e, event.Event{
 					Type: event.Change,
 					Pid:  w.pid,
 					Fd:   fd,