3 files changed, 308 insertions, 98 deletions

diff --git a/filter/filter.go b/filter/filter.go
index c83fb13..2a774a3 100644
--- a/filter/filter.go
+++ b/filter/filter.go
@@ -5,114 +5,213 @@ import (
         "fmt"
         "log"
         "os"
-        "time"
+        "sync"
 
         "go.mongodb.org/mongo-driver/bson"
         "go.mongodb.org/mongo-driver/mongo"
         "go.mongodb.org/mongo-driver/mongo/options"
-        "go.mongodb.org/mongo-driver/mongo/readpref"
 )
 
-type Exec struct {
-        timestamp time.Time `bson:"timestamp"`
-        execArgs  []string  `bson:"execArgs"`
-}
-
-type Process struct {
-        timestamp      time.Time `bson:"start_timestamp"`
-        ppid           int       `bson:"ppid"`
-        parentTgid     int       `bson:"parentTgid"`
-        pid            int       `bson:"pid"`
-        tgid           int       `bson:"tgid"`
-        args           []string  `bson:"args"`
-        comm           string    `bson:"comm"`
-        cwd            string    `bson:"cwd"`
-        execve         []Exec    `bson:"execve"`
-        exit_code      int       `bson:"exit_code"`
-        exit_signal    int       `bson:"exit_signal"`
-        exit_timestamp time.Time `bson:"exit_timestamp"`
-}
+const (
+        oldDBName     = "test"
+        oldPidColName = "pids"
+)
 
-func (p Process) String() string {
-        var res string
-        res = ""
-        res += fmt.Sprintf("timestamp\t%v\n", p.timestamp)
-        res += fmt.Sprintf("ppid\t%d\nparentTgid\t%d\n", p.ppid, p.parentTgid)
-        res += fmt.Sprintf("pid\t%d\ntgid\t%d\nargs: ", p.pid, p.tgid)
-        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)
-        return res
+type treeNode struct {
+        Tgid     int
+        Threads  []Process
+        Children []int
 }
 
-// type Process struct {
-//      StartTimestamp time.Time  `bson:"start_timestamp"`
-//      Ppid           *int       `bson:"ppid"`
-//      ParentTgid     *int       `bson:"parentTgid"`
-//      Pid            int        `bson:"pid"`
-//      Tgid           int        `bson:"tgid"`
-//      Args           []string   `bson:"args"`
-//      Comm           *string    `bson:"comm"`
-//      Cwd            *string    `bson:"cwd"`
-//      Execve         []Exec     `bson:"execve"`
-//      ExitCode       *int       `bson:"exit_code"`
-//      ExitSignal     *int       `bson:"exit_signal"`
-//      ExitTimestamp  *time.Time `bson:"exit_timestamp"`
-// }
-
-// func (p Process) String() string {
-//      var res string
-//      res = ""
-//      res += fmt.Sprintf("timestamp\t%v\n", p.StartTimestamp)
-//      if p.Ppid != nil && p.ParentTgid != nil {
-//              res += fmt.Sprintf("ppid\t%d\nparentTgid\t%d\n", *(p.Ppid), *(p.ParentTgid))
-//      }
-//      res += fmt.Sprintf("pid\t%d\ntgid\t%d\nargs: ", p.Pid, p.Tgid)
-//      for i := 0; i < len(p.Args); i++ {
-//              res += fmt.Sprintf("%s ", p.Args[i])
-//      }
-//      if p.Comm != nil && p.Cwd != nil {
-//              res += fmt.Sprintf("\ncomm\t%s\ncwd\t%s\n", *(p.Comm), *(p.Cwd))
-//      }
-//      return res
-// }
-
 func main() {
+        // 连接到MongoDB
         client, err := mongo.Connect(context.TODO(), options.Client().ApplyURI("mongodb://localhost:27017"))
         if err != nil {
-                fmt.Fprintf(os.Stderr, "Err connecting mongodb: %v\n", err)
+                log.Fatal(err)
         }
         defer client.Disconnect(context.TODO())
 
-        // 检查连接
-        err = client.Ping(context.TODO(), readpref.Primary())
+        // 选择数据库和集合
+        db := client.Database(oldDBName)
+        collection := db.Collection(oldPidColName)
+
+        // 提取所有数据
+        var res []Process
+
+        cursor, err := collection.Find(context.Background(), bson.M{})
         if err != nil {
-                log.Fatal(err)
+                fmt.Fprintf(os.Stderr, "Err: %v\n", err)
+        }
+        if err := cursor.All(context.Background(), &res); err != nil {
+                fmt.Fprintf(os.Stderr, "Err All: %v\n", err)
         }
 
-        pidCol := client.Database("test").Collection("pids")
-        cur, err := pidCol.Find(context.TODO(), bson.M{}) // 查询所有文档
-        if err != nil {
-                log.Fatal(err)
+        var merged sync.Map
+        for _, process := range res {
+                tmp, ok := merged.Load(process.Pid)
+                if ok {
+                        // 证明重复了，要合并
+                        tmp := ProMerge(tmp.(Process), process)
+                        merged.Store(process.Pid, tmp)
+                } else {
+                        // 没有，直接插入
+                        merged.Store(process.Pid, process)
+                }
         }
-        defer cur.Close(context.TODO()) // 确保游标被关闭
 
-        var res []Process
-        for cur.Next(context.TODO()) {
-                var tmp Process
-                // 解码到Process结构体
-                if err := cur.Decode(&tmp); err != nil {
-                        log.Fatal(err)
+        var treeMap sync.Map
+        findTgid := make(map[int]int)
+        var stared int
+        merged.Range(func(key, val interface{}) bool {
+                tmp := val.(Process)
+                if tmp.Star {
+                        stared = tmp.Tgid
+                }
+                // 登记tgid
+                findTgid[tmp.Pid] = tmp.Tgid
+                nodeTmp, ok := treeMap.Load(tmp.Tgid)
+                if ok {
+                        // 直接记录
+                        node := nodeTmp.(treeNode)
+                        node.Threads = append(node.Threads, tmp)
+                        node.Children = append(node.Children, tmp.Children...)
+                        treeMap.Store(tmp.Tgid, node)
+                } else {
+                        node := treeNode{
+                                Tgid:     tmp.Tgid,
+                                Threads:  make([]Process, 0),
+                                Children: make([]int, 0),
+                        }
+                        node.Threads = append(node.Threads, tmp)
+                        node.Children = append(node.Children, tmp.Children...)
+                        treeMap.Store(tmp.Tgid, node)
                 }
-                res = append(res, tmp)
+                return true
+        })
+
+        // 从tgid==stared开始，构建树
+        var helloTree sync.Map // 在树上的tgid节点
+        var q Queue            // 记录每一个整理好的结构体，bfs
+        visited := make(map[int]bool)
+        visited[stared] = true
+        tmp, ok := treeMap.Load(stared)
+        if !ok {
+                return
         }
 
-        if err := cur.Err(); err != nil {
-                log.Fatal(err)
+        q.Enqueue(tmp)
+        helloTree.Store(stared, tmp)
+        for !q.IsEmpty() {
+                tmp, ok := q.Dequeue()
+                if !ok {
+                        continue
+                }
+                node := tmp.(treeNode)
+                for i := 0; i < len(node.Children); i++ {
+                        tgid := findTgid[node.Children[i]]
+                        _, exists := visited[tgid]
+                        if !exists {
+                                visited[tgid] = true
+                                tgidNode, ok := treeMap.Load(tgid)
+                                if !ok {
+                                        continue
+                                }
+                                helloTree.Store(tgid, tgidNode)
+                                q.Enqueue(tgidNode)
+                        }
+                }
+        }
+
+        // TODO:
+        // 1.修改数据结构，使之自身即存储树结构，插入数据库后前端拿出来就能用
+        // 2.还有其余优化要做，比如线程退出时间与进程推出时间
+
+        // count := 0
+        // helloTree.Range(func(key, val interface{}) bool {
+        //      count++
+        //      fmt.Printf("tgid: %d\n", val.(treeNode).Tgid)
+        //      return true
+        // })
+        // fmt.Printf("Star: %d, res: %d\n", stared, count)
+
+        // 接下来处理文件
+}
+
+func ProMerge(a, b Process) (res Process) {
+        // 合并过程中会遇到什么问题？
+        res.Star = false
+
+        if a.StartTimestamp.IsZero() {
+                res.StartTimestamp = b.StartTimestamp
+        } else if b.StartTimestamp.IsZero() {
+                res.StartTimestamp = a.StartTimestamp
+        } else if a.StartTimestamp.Before(b.StartTimestamp) {
+                res.StartTimestamp = a.StartTimestamp
+        } else {
+                res.StartTimestamp = b.StartTimestamp
+        }
+
+        res.Ppid = a.Ppid
+        if a.ParentTgid == 0 {
+                res.ParentTgid = b.ParentTgid
+        } else {
+                res.ParentTgid = a.ParentTgid
+        }
+
+        res.Pid = a.Pid
+        if a.Tgid == 0 {
+                res.Tgid = b.Tgid
+        } else {
+                res.Tgid = a.Tgid
+        }
+
+        if len(a.Args) == 0 {
+                res.Args = b.Args
+        } else {
+                res.Args = a.Args
         }
 
-        for i := 0; i < len(res); i++ {
-                fmt.Printf("------\n%v\n", res[i])
+        if a.Comm == "" {
+                res.Comm = b.Comm
+        } else {
+                res.Comm = a.Comm
         }
+
+        if a.RootFS == "" {
+                res.RootFS = b.RootFS
+        } else {
+                res.RootFS = a.RootFS
+        }
+
+        if a.Cwd == "" {
+                res.Cwd = b.Cwd
+        } else {
+                res.Cwd = a.Cwd
+        }
+
+        res.Execve = append(a.Execve, b.Execve...)
+        res.Children = append(a.Children, b.Children...)
+
+        var flag bool // 真a假b
+        if a.ExitTimestamp.IsZero() {
+                flag = false
+        } else if b.ExitTimestamp.IsZero() {
+                flag = true
+        } else if a.ExitTimestamp.Before(b.ExitTimestamp) {
+                flag = true
+        } else {
+                flag = false
+        }
+
+        if flag {
+                res.ExitCode = a.ExitCode
+                res.ExitSignal = a.ExitSignal
+                res.ExitTimestamp = a.ExitTimestamp
+        } else {
+                res.ExitCode = b.ExitCode
+                res.ExitSignal = b.ExitSignal
+                res.ExitTimestamp = b.ExitTimestamp
+        }
+
+        return res
 }
diff --git a/filter/global.go b/filter/global.go
new file mode 100644
index 0000000..45706d4
--- /dev/null
+++ b/filter/global.go
@@ -0,0 +1,92 @@
+package main
+
+import (
+        "fmt"
+        "time"
+)
+
+type Exec struct {
+        Timestamp time.Time `bson:"timestamp"`
+        ExecArgs  []string  `bson:"execArgs"`
+}
+
+type Process struct {
+        Star           bool      `bson:"star"`
+        StartTimestamp time.Time `bson:"start_timestamp"`
+        Ppid           int       `bson:"ppid"`
+        ParentTgid     int       `bson:"parentTgid"`
+        Pid            int       `bson:"pid"`
+        Tgid           int       `bson:"tgid"`
+        Args           []string  `bson:"args"`
+        Comm           string    `bson:"comm"`
+        RootFS         string    `bson:"rootfs"`
+        Cwd            string    `bson:"cwd"`
+        Children       []int     `bson:"children"`
+        Execve         []Exec    `bson:"execve"`
+        ExitCode       int       `bson:"exit_code"`
+        ExitSignal     int       `bson:"exit_signal"`
+        ExitTimestamp  time.Time `bson:"exit_timestamp"`
+}
+
+func (p Process) String() string {
+        var res string
+        res = ""
+        res += fmt.Sprintf("timestamp\t%v\n", p.StartTimestamp)
+        res += fmt.Sprintf("ppid\t%d\nparentTgid\t%d\n", p.Ppid, p.ParentTgid)
+        res += fmt.Sprintf("pid\t%d\ntgid\t%d\nargs: ", p.Pid, p.Tgid)
+        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)
+        if len(p.Execve) != 0 {
+                res += fmt.Sprintf("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 += fmt.Sprintf("children: ")
+        for i := 0; i < len(p.Children); i++ {
+                res += fmt.Sprintf("%d ", p.Children[i])
+        }
+        res += fmt.Sprintf("\n")
+        return res
+}
+
+// Queue 定义一个队列结构体
+type Queue struct {
+        items []interface{}
+}
+
+// NewQueue 创建一个新的队列
+func NewQueue() *Queue {
+        return &Queue{items: make([]interface{}, 0)}
+}
+
+// Enqueue 向队列中添加一个元素
+func (q *Queue) Enqueue(item interface{}) {
+        q.items = append(q.items, item)
+}
+
+// Dequeue 从队列中移除并返回队列前面的元素
+func (q *Queue) Dequeue() (interface{}, bool) {
+        if len(q.items) == 0 {
+                return nil, false
+        }
+        item := q.items[0]
+        q.items = q.items[1:]
+        return item, true
+}
+
+// Size 返回队列中的元素数量
+func (q *Queue) Size() int {
+        return len(q.items)
+}
+
+// IsEmpty 检查队列是否为空
+func (q *Queue) IsEmpty() bool {
+        return len(q.items) == 0
+}
diff --git a/listener/deal.go b/listener/deal.go
index 8f77431..8225224 100644
--- a/listener/deal.go
+++ b/listener/deal.go
@@ -3,6 +3,7 @@ package main
 import (
         "fmt"
         "os"
+        "strconv"
         "syscall"
         "time"
 
@@ -17,27 +18,45 @@ const (
 )
 
 var pidCol, fdCol, fileCol mongoClient
-var err error
+
+func initPidCol() (err error) {
+        // TODO: 这里是否需要补全一下进程信息？
+        dirs, err := os.ReadDir(fmt.Sprintf("/proc/%d/task", containerdPid))
+        if err != nil {
+                return err
+        }
+        for _, file := range dirs {
+                pid, _ := strconv.Atoi(file.Name())
+                process := Process{
+                        Ppid:       1,
+                        ParentTgid: 1,
+                        Pid:        pid,
+                        Tgid:       containerdPid,
+                        Cwd:        "/",
+                        Children:   make([]int, 0),
+                        Execve:     make([]Exec, 0),
+                        Args:       make([]string, 0),
+                }
+                if pid == containerdPid {
+                        process.Star = true
+                }
+                err = pidCol.InsertOne(process)
+        }
+        return nil
+}
 
 func deal() {
         defer wg.Done()
         var cooked Event
         var ok bool
+        var err error
 
         if err = pidCol.init(dbName, pidColName); err != nil {
                 fmt.Fprintf(os.Stderr, "Error while initing the mongodb: %v\n", err)
                 return
         }
-        err = pidCol.InsertOne(Process{
-                Ppid:     1,
-                Pid:      containerdPid,
-                Cwd:      "/",
-                Children: make([]int, 0),
-                Star:     true,
-        })
-        if err != nil {
-                fmt.Fprintf(os.Stderr, "Error while initing the mongodb: %v\n", err)
-                return
+        if err = initPidCol(); err != nil {
+                fmt.Fprintf(os.Stderr, "Err while initing pidcol: %v\n", err)
         }
 
         if err = fdCol.init(dbName, fdColName); err != nil {
@@ -96,7 +115,7 @@ func deletePid(cooked Event) {
 func dealNewPid(cooked Event) {
         // 自身是否已经记录
         var docRes []Process
-        err = pidCol.Finddoc(bson.M{"pid": cooked.pid}, &docRes)
+        err := pidCol.Finddoc(bson.M{"pid": cooked.pid}, &docRes)
         if err != nil {
                 fmt.Fprintf(os.Stderr, "Err finding: %v\n", err)
                 return
@@ -136,7 +155,7 @@ func dealNewPid(cooked Event) {
                 }
         }
 
-        err := pidCol.UpdateOne(bson.M{"pid": cooked.ppid}, bson.M{
+        err = pidCol.UpdateOne(bson.M{"pid": cooked.ppid}, bson.M{
                 "$push": bson.M{
                         "children": cooked.pid,
                 },
@@ -149,7 +168,7 @@ func dealNewPid(cooked Event) {
 func dealExecve(cooked Event) {
         var docRes []Process
         // 首先检查进程是否存在，如不存在则为之创建
-        err = pidCol.Finddoc(bson.M{"pid": cooked.pid}, &docRes)
+        err := pidCol.Finddoc(bson.M{"pid": cooked.pid}, &docRes)
         if err != nil {
                 return
         }