5 files changed, 153 insertions, 56 deletions

diff --git a/filter/filter.go b/filter/filter.go
index b2341ec..98c326c 100644
--- a/filter/filter.go
+++ b/filter/filter.go
@@ -5,6 +5,7 @@ import (
         "fmt"
         "log"
         "os"
+        "path"
         "sort"
 
         "go.mongodb.org/mongo-driver/bson"
@@ -28,7 +29,7 @@ var findTgid map[int]int
 var helloTree map[int]*tgidNode
 
 // 文件信息
-var files []*File
+var files []File
 
 func main() {
         // 连接到MongoDB
@@ -110,8 +111,8 @@ func main() {
         }
 
         newFileCol := newDB.Collection(newFileColName)
-        for _, pFile := range files {
-                newFileCol.InsertOne(context.Background(), *pFile)
+        for _, file := range files {
+                newFileCol.InsertOne(context.Background(), file)
         }
 }
 
@@ -194,41 +195,48 @@ func ProMerge(a, b Process) (res Process) {
         return res
 }
 
-func filtPids(pRawPidData *[]Process) {
+func mergeProcess(pRawPidData *[]Process) (merged []Process) {
         rawPidData := *pRawPidData
-        // 合并由多线程导致的重复记录
-        merged := make(map[int]Process) // pid --> Process
+        // 合并由多线程导致的重复记录，顺便按照pid升序
+        index := make(map[int]int)
         for _, process := range rawPidData {
-                tmp, exists := merged[process.Pid]
+                i, exists := index[process.Pid]
                 if exists {
-                        // 证���重���了，要合并
-                        merged[process.Pid] = ProMerge(tmp, process)
+                        // 已�����，合并
+                        merged[i] = ProMerge(merged[i], process)
                 } else {
-                        // 没有，直接插入
-                        merged[process.Pid] = process
+                        // 不存在，直接添加
+                        merged = append(merged, process)
+                        index[process.Pid] = len(merged) - 1
                 }
         }
+        sort.Slice(merged, func(i, j int) bool {
+                return merged[i].Pid < merged[j].Pid
+        })
+        return merged
+}
 
+func getTgidNodes(merged []Process) (tgidMap map[int]*tgidNode, starTgid int, rootfsPids []int) {
         // 合并出来的进程整理为tgidNode
-        // var tgidMap map[int]*tgidNode // tgid --> tgidNode
-        tgidMap := make(map[int]*tgidNode)
+        tgidMap = make(map[int]*tgidNode)
         findTgid = make(map[int]int) // pid --> tgid
-        var stared int
-        stared = -1
+        // var starTgid, rootFsPid int
+        starTgid = -1
+        // rootfsPid = -1
+        rootfsPids = make([]int, 0)
         for _, val := range merged {
                 if val.Star {
-                        stared = val.Tgid
+                        starTgid = val.Tgid
+                } else if val.RootFS != "" {
+                        rootfsPids = append(rootfsPids, val.Pid)
                 }
                 // 登记tgid
                 findTgid[val.Pid] = val.Tgid
-                // nodeval, ok := tgidMap.Load(val.Tgid)
                 nodeval, exists := tgidMap[val.Tgid]
                 if exists {
                         // 直接记录
-                        // node := nodeval.(tgidNode)
                         nodeval.Threads = append(nodeval.Threads, val)
                         nodeval.FindPid[val.Pid] = len(nodeval.Threads) - 1
-                        // tgidMap.Store(val.Tgid, node)
                 } else {
                         node := tgidNode{
                                 Tgid:      val.Tgid,
@@ -238,21 +246,19 @@ func filtPids(pRawPidData *[]Process) {
                         }
                         node.Threads = append(node.Threads, val)
                         node.FindPid[val.Pid] = 0
-                        // tgidMap.Store(val.Tgid, node)
                         tgidMap[val.Tgid] = &node
                 }
         }
+        return tgidMap, starTgid, rootfsPids
+}
 
-        // 从tgid==stared开始，构建树
+func buildTree(tgidMap map[int]*tgidNode, starTgid int) {
+        // 从tgid==starTgid开始，构建树
         helloTree = make(map[int]*tgidNode) // 在树上的tgid节点，tgid --> *tgidNode
         var q Queue                         // 记录每一个整理好的结构体，bfs
         visited := make(map[int]bool)       // 哪些tgid已经访问过
 
-        // tmp, ok := tgidMap.Load(stared)
-        // if !ok {
-        //      return
-        // }
-        tmp, exists := tgidMap[stared]
+        tmp, exists := tgidMap[starTgid]
         if !exists {
                 return
         }
@@ -262,7 +268,7 @@ func filtPids(pRawPidData *[]Process) {
         // 因而所有添加子代tgid的行为只针对helloTree
         // q不添加，直接把新的tgid对应的tgidNode入队就是了
         q.Enqueue(tmp)
-        visited[stared] = true
+        visited[starTgid] = true
         for !q.IsEmpty() {
                 tmp, ok := q.Dequeue()
                 if !ok {
@@ -276,7 +282,6 @@ func filtPids(pRawPidData *[]Process) {
                                 _, exists := visited[tgid]
                                 if !exists {
                                         // 子代里有没见过的tgid
-                                        // tgidNode, ok := tgidMap.Load(tgid)
                                         tgidNode, exists := tgidMap[tgid]
                                         if !exists {
                                                 continue
@@ -288,31 +293,79 @@ func filtPids(pRawPidData *[]Process) {
                         }
                 }
         }
+}
 
-        // TODO:
-        // 1.√修改数据结构，使之自身即存储树结构，插入数据库后前端拿出来就能用
-        // 2.还有其余优化要做，比如线程退出时间与进程推出时间，关系到后续的文件修理
-        // 3.根文件系统，问题很重大
+func optimazePid(starTgid int, rootfsPids []int) {
+        getDockerRootFs := make(map[string]string) // dockerId --> rootfs
+        // 首先处理一下记录有pivot_root信息的进程，防止pivot先于fork
+        for _, rootfsPid := range rootfsPids {
+                rootfsTgid := findTgid[rootfsPid]
+                i := helloTree[rootfsTgid].FindPid[rootfsPid]
+                rootfsProcess := &(helloTree[rootfsTgid].Threads[i])
+                if rootfsProcess.RootFS == "cwd" {
+                        rootfsProcess.RootFS = rootfsProcess.Cwd
+                }
+                getDockerRootFs[rootfsProcess.DockerId] = rootfsProcess.RootFS
+        }
 
         count := 0
         for _, val := range helloTree {
-                count++
-                fmt.Printf("==============================\ntgid: %6d, size: %6d, children: ", val.Tgid, len(val.Threads))
-                for _, child := range val.ChildTgid {
-                        fmt.Printf("%7d", child)
-                }
-                fmt.Printf("\n")
-                for _, process := range val.Threads {
-                        fmt.Printf("%v\n", process)
+                // 处理一下pid结束时间，顺便找找爹
+                // 结束时间是因为很多线程结束时间没获取到，默认按照进程退出时间处理
+                // Ppid是因为进程产生之初收到的信息写的爹一定是亲爹
+                // 但是产生线程时候该进程很可能已作为孤儿被收养，导致线程里关于爹的记录是继父
+                for i := 0; i < len(val.Threads); i++ {
+                        if i != 0 {
+                                if val.Threads[i].Tgid < val.Threads[0].Tgid {
+                                        val.Threads[i].ParentTgid = val.Threads[0].ParentTgid
+                                        val.Threads[i].Ppid = val.Threads[0].Ppid
+                                }
+                                if val.Threads[i].ExitTimestamp.IsZero() {
+                                        val.Threads[i].ExitCode = val.Threads[0].ExitCode
+                                        val.Threads[i].ExitTimestamp = val.Threads[0].ExitTimestamp
+                                        val.Threads[i].ExitSignal = val.Threads[0].ExitSignal
+                                }
+                        }
+
+                        dockerId := val.Threads[i].DockerId
+                        if dockerId != "" {
+                                rootfs, exists := getDockerRootFs[dockerId]
+                                if !exists {
+                                        fmt.Fprintf(os.Stderr, "Err: the docker rootfs of pid %d is not known!\n", val.Threads[i].Pid)
+                                        continue
+                                }
+                                val.Threads[i].RootFS = rootfs
+                        }
                 }
-                fmt.Printf("\n\n\n")
+
+                count++
+                fmt.Printf("%v\n", *val)
         }
-        fmt.Printf("Star: %d, res: %d\n", stared, count)
+        fmt.Printf("Star: %d, res: %d\n", starTgid, count)
+}
+
+func filtPids(pRawPidData *[]Process) {
+        /* ATTENTION: 把map/slice直接传参是危险的
+         * 传递的是指针，不会引起大的复制开销，
+         * 但是map/slice在callee func内被修改**可能**导致内存更改
+         * 而这样的内存更改对caller function来说是不可见的，看到的还是原来的东西
+         * 这里由于参数几乎都是只读不写，因而用一下
+         */
+
+        // 合并由多线程导致的重复记录，顺便按照pid升序
+        // 多线程已经取消了，但保险起见还是留着
+        merged := mergeProcess(pRawPidData)
+        // 将Process按照tgid合并
+        tgidMap, starTgid, rootfsPids := getTgidNodes(merged)
+        // 建树，helloTree
+        buildTree(tgidMap, starTgid)
+        // 对树上的进程做一些优化处理
+        optimazePid(starTgid, rootfsPids)
 }
 
 func filtFiles(pRawFileData *[]File) {
         rawFileData := *pRawFileData
-        files = make([]*File, 0)
+        files = make([]File, 0)
 
         // 所有文件按照特定顺序排
         sort.Slice(rawFileData, func(i, j int) bool {
@@ -342,6 +395,9 @@ func filtFiles(pRawFileData *[]File) {
         })
 
         for _, file := range rawFileData {
+                if file.FileName == "/root/test/1/../.hello.c.swp" {
+                        fmt.Printf("Test\n")
+                }
                 tgid := findTgid[file.Pid]
                 pTgidNode, exists := helloTree[tgid]
                 if !exists {
@@ -354,6 +410,7 @@ func filtFiles(pRawFileData *[]File) {
                         }
                         file.CloseTimestamp = pTgidNode.Threads[index].ExitTimestamp
                 }
-                files = append(files, &file)
+                file.FileName = path.Clean(file.FileName)
+                files = append(files, file)
         }
 }
diff --git a/filter/global.go b/filter/global.go
index 37af52b..bade895 100644
--- a/filter/global.go
+++ b/filter/global.go
@@ -22,6 +22,7 @@ type Process struct {
         RootFS         string    `bson:"rootfs"`
         Cwd            string    `bson:"cwd"`
         Children       []int     `bson:"children"`
+        DockerId       string    `bson:"docker_id"`
         Execve         []Exec    `bson:"execve"`
         ExitCode       int       `bson:"exit_code"`
         ExitSignal     int       `bson:"exit_signal"`
@@ -44,26 +45,40 @@ func (p Process) String() string {
         for i := 0; i < len(p.Args); i++ {
                 res += fmt.Sprintf("%s ", p.Args[i])
         }
-        res += fmt.Sprintf("\ncomm\t%s\ncwd\t%s\n", p.Comm, p.Cwd)
+        res += fmt.Sprintf("\ncomm\t%s\ncwd\t%s\nrootfs\t%s\ndocker_id\t%s\n", p.Comm, p.Cwd, p.RootFS, p.DockerId)
         if len(p.Execve) != 0 {
-                res += fmt.Sprintf("exec:\n")
+                res += "exec:\n"
                 for i := 0; i < len(p.Execve); i++ {
                         res += fmt.Sprintf("\ttimestamp: %v\n\texecArgs:\t", p.Execve[i].Timestamp)
                         for j := 0; j < len(p.Execve[i].ExecArgs); j++ {
                                 res += fmt.Sprintf("%s ", p.Execve[i].ExecArgs[j])
                         }
-                        res += fmt.Sprintf("\n")
+                        res += "\n"
                 }
         }
-        res += fmt.Sprintf("children: ")
+        res += "children: "
         for i := 0; i < len(p.Children); i++ {
                 res += fmt.Sprintf("%d ", p.Children[i])
         }
-        res += fmt.Sprintf("\n")
+        res += "\n"
         res += fmt.Sprintf("exit_timestamp:\t%v\nexit_code:\t%d\nexit_signal:\t%d\n", p.ExitTimestamp, p.ExitCode, p.ExitSignal)
         return res
 }
 
+func (node tgidNode) String() string {
+        var res string
+        res += fmt.Sprintf("==============================\ntgid: %6d, size: %6d, children: ", node.Tgid, len(node.Threads))
+        for _, child := range node.ChildTgid {
+                res += fmt.Sprintf("%7d", child)
+        }
+        res += "\n"
+        for _, process := range node.Threads {
+                res += fmt.Sprintf("%v\n", process)
+        }
+        res += "\n"
+        return res
+}
+
 type File struct {
         OpenTimestamp  time.Time   `bson:"timestamp"`
         FileName       string      `bson:"fileName"`
diff --git a/listener/deal.go b/listener/deal.go
index 8225224..70c2827 100644
--- a/listener/deal.go
+++ b/listener/deal.go
@@ -79,19 +79,19 @@ func deal() {
 
                 switch cooked.tag {
                 case NEWPID:
-                        go dealNewPid(cooked)
+                        dealNewPid(cooked)
                 case EXECVE:
-                        go dealExecve(cooked)
+                        dealExecve(cooked)
                 case PIDEXIT:
-                        go deletePid(cooked)
+                        deletePid(cooked)
                 case FILEOPEN:
-                        go fileOpen(cooked)
+                        fileOpen(cooked)
                 case FILEWRITE:
-                        go fileWrite(cooked)
+                        fileWrite(cooked)
                 case FILECLOSE:
-                        go fileClose(cooked)
+                        fileClose(cooked)
                 case PIVOTROOT:
-                        go pivotRoot(cooked)
+                        pivotRoot(cooked)
                 }
         }
 }
@@ -131,6 +131,7 @@ func dealNewPid(cooked Event) {
                 docRes[0].Cwd = cooked.cwd
                 docRes[0].Comm = cooked.comm
                 docRes[0].Args = cooked.argv
+                docRes[0].DockerId = cooked.cgroup
 
                 err := pidCol.ReplaceOne(bson.M{"pid": cooked.pid}, docRes[0])
                 if err != nil {
@@ -149,6 +150,7 @@ func dealNewPid(cooked Event) {
                         Cwd:            cooked.cwd,
                         Execve:         make([]Exec, 0),
                         Children:       make([]int, 0),
+                        DockerId:       cooked.cgroup,
                 })
                 if err != nil {
                         fmt.Fprintf(os.Stderr, "Err inserting: %v\n", err)
diff --git a/listener/global.go b/listener/global.go
index 11b18bf..b782284 100644
--- a/listener/global.go
+++ b/listener/global.go
@@ -37,6 +37,7 @@ type Event struct {
         argv             []string
         comm             string
         cwd              string
+        cgroup           string
         exit_code        int
         exit_signal      int
         srcPath          string
@@ -67,6 +68,7 @@ type Process struct {
         RootFS         string    `bson:"rootfs"`
         Cwd            string    `bson:"cwd"`
         Children       []int     `bson:"children"`
+        DockerId       string    `bson:"docker_id"`
         Execve         []Exec    `bson:"execve"`
         ExitCode       int       `bson:"exit_code"`
         ExitSignal     int       `bson:"exit_signal"`
diff --git a/listener/godo.go b/listener/godo.go
index 87e9446..8d82231 100644
--- a/listener/godo.go
+++ b/listener/godo.go
@@ -8,6 +8,7 @@ import (
         "netlink"
         "os"
         "os/exec"
+        "regexp"
         "strings"
         "syscall"
         "time"
@@ -176,4 +177,24 @@ func checkProc(pCooked *Event) {
                 fmt.Fprintf(os.Stderr, "Err: %v\n", err)
                 pCooked.cwd = ""
         }
+
+        fd, err = os.Open(fileName + "cgroup")
+        if err != nil {
+                fmt.Fprintf(os.Stderr, "Err: %v\n", err)
+                // cgroup记空，即没赶上
+                return
+        }
+        scanner = bufio.NewScanner(fd)
+        cgroupRegex := regexp.MustCompile(`/docker/([0-9a-f]+)$`)
+        scanner.Split(bufio.ScanLines)
+        for scanner.Scan() {
+                line := scanner.Text()
+                if cgroupRegex.MatchString(line) {
+                        match := cgroupRegex.FindStringSubmatch(line)
+                        pCooked.cgroup = match[1]
+                        return
+                }
+        }
+        fd.Close()
+        pCooked.cgroup = ""
 }