diff --git a/crit/utils.go b/crit/utils.go index f76c3a7df..9ce0adb7e 100644 --- a/crit/utils.go +++ b/crit/utils.go @@ -232,3 +232,17 @@ func getUnixSkFilePath(dir string, file *fdinfo.FileEntry, fID uint32) (string, return "unix[?]", nil } + +// FindPs performs a short-circuiting depth-first search to find +// a process with a given PID in a process tree. +func (ps *PsTree) FindPs(pid uint32) *PsTree { + if ps.PID == pid { + return ps + } + for _, child := range ps.Children { + if process := child.FindPs(pid); process != nil { + return process + } + } + return nil +} diff --git a/crit/utils_test.go b/crit/utils_test.go new file mode 100644 index 000000000..6174e447b --- /dev/null +++ b/crit/utils_test.go @@ -0,0 +1,45 @@ +package crit + +import ( + "testing" +) + +func TestFindPs(t *testing.T) { + root := &PsTree{ + PID: 1, + Children: []*PsTree{ + { + PID: 2, + Children: []*PsTree{ + { + PID: 3, + }, + }, + }, + { + PID: 4, + Children: []*PsTree{ + { + PID: 5, + }, + }, + }, + }, + } + + // Test Case 1: Find an existing process with a valid PID + ps := root.FindPs(3) + if ps == nil { + t.Errorf("FindPs(3) returned nil, expected a valid process") + } + if ps != nil && ps.PID != 3 { + t.Errorf("FindPs(3) returned a process with PID %d, expected 3", ps.PID) + } + + // Test Case 2: Find a non-existing process with an invalid PID + nonExistentPID := uint32(999) + notFoundProcess := root.FindPs(nonExistentPID) + if notFoundProcess != nil { + t.Errorf("FindPs(%d) returned a process, expected nil", nonExistentPID) + } +}