1 files changed, 102 insertions, 71 deletions

diff --git a/src/deal.go b/src/deal.go
index 118d914..783dab8 100644
--- a/src/deal.go
+++ b/src/deal.go
@@ -1,51 +1,48 @@
 package main
 
 import (
-        "context"
         "fmt"
+        "sync"
         "time"
 
         "go.mongodb.org/mongo-driver/bson"
-        "go.mongodb.org/mongo-driver/mongo"
-        "go.mongodb.org/mongo-driver/mongo/options"
 )
 
 const (
-        dbName  string = "test"
-        colName string = "pids"
+        dbName     string = "test"
+        pidColName string = "pids"
 )
 
+var mongoMutex sync.Mutex
+var pidCol mongoClient
+
 func deal() {
         defer wg.Done()
         var cooked Event
         var ok bool
 
         var err error
-        var mongo *mongo.Client
         var res []bson.M
 
-        mongo, err = connect()
-        if err != nil {
-                fmt.Printf("Err connecting the mongodb: %v\n", err)
+        if err = pidCol.Connect(dbName, pidColName); err != nil {
+                fmt.Printf("Error connecting the mongodb: %v\n", err)
         }
-        pidCol := mongo.Database(dbName).Collection(colName)
-
-        err = pidCol.Drop(context.Background())
-        if err != nil {
-                fmt.Printf("Err drop: %v\n", err)
+        if err = pidCol.Drop(); err != nil {
+                fmt.Printf("Error drop the mongodb: %v\n", err)
         }
 
-        _, err = pidCol.InsertOne(context.Background(), bson.M{
-                "ppid": 1,
-                "pid":  containerdPid,
-                "cwd":  "/",
+        err = pidCol.InsertOne(bson.M{
+                "ppid":     1,
+                "pid":      containerdPid,
+                "cwd":      "/",
+                "children": bson.M{},
         })
         if err != nil {
                 fmt.Printf("Err containerd: %v", err)
                 return
         }
-
         fmt.Printf("Containerd: %d\n", containerdPid)
+        defer pidCol.Disconnect()
 
         for {
                 cooked, ok = <-cookedChan
@@ -54,81 +51,115 @@ func deal() {
                 }
 
                 switch syscallTable[cooked.syscall] {
-                case "fork", "vfork", "clone":
+                case "clone":
                         // 有无父进程在观察中
-                        res, err = findDocuments(mongo, "test", "pids", bson.M{"pid": cooked.ppid})
+                        res, err = pidCol.Finddoc(bson.M{"pid": cooked.ppid})
                         if err != nil || len(res) != 1 {
                                 break
                         }
 
                         // 自身是否已经记录
-                        res, err = findDocuments(mongo, "test", "pids", bson.M{"pid": cooked.pid})
+                        res, err = pidCol.Finddoc(bson.M{"pid": cooked.pid})
                         if err != nil {
                                 fmt.Printf("Err finding: %v\n", err)
                                 break
-                        } else if len(res) != 0 {
-                                fmt.Printf("Err inserting pid %v: already in db: %v\n", cooked.pid, res)
-                                break
-                        }
-
-                        doc := []bson.A{}
-                        for _, str := range cooked.argv {
-                                doc = append(doc, bson.A{str})
                         }
-                        _, err := pidCol.InsertOne(context.Background(), bson.M{
-                                "timestamp": cooked.timestamp,
-                                "ppid":      cooked.ppid,
-                                "pid":       cooked.pid,
-                                "cwd":       cooked.cwd,
-                                "args":      doc,
-                                "children":  []bson.M{},
-                        })
-                        if err != nil {
-                                fmt.Printf("Err insert: %v\n", err)
+                        mongoMutex.Lock()
+                        if len(res) != 0 {
+                                // 进程原本就存在，换言之别的消息先到了
+                                // 所有先行抵达的消息必须保留execve/children字段
+                                // 此处不再更新
+                                // 以防把原有信息更没了
+                                pidCol.UpdateOne(bson.M{"pid": cooked.pid}, bson.M{
+                                        "start_timestamp": cooked.timestamp,
+                                        "ppid":            cooked.ppid,
+                                        "pid":             cooked.pid,
+                                        "cwd":             cooked.cwd,
+                                        // "execve":          []bson.M{},
+                                        "args": cooked.argv,
+                                        // "children":        []bson.M{},
+                                })
+                        } else {
+                                // 这进程本是新修的
+                                pidCol.InsertOne(bson.M{
+                                        "start_timestamp": cooked.timestamp,
+                                        "ppid":            cooked.ppid,
+                                        "pid":             cooked.pid,
+                                        "cwd":             cooked.cwd,
+                                        "execve":          []bson.M{},
+                                        "args":            cooked.argv,
+                                        "children":        []bson.M{},
+                                })
                         }
 
-                        _, err = pidCol.UpdateOne(context.Background(), bson.M{"pid": cooked.pid}, bson.M{
+                        pidCol.UpdateOne(bson.M{"pid": cooked.ppid}, bson.M{
                                 "$push": bson.M{
                                         "children": cooked.pid,
                                 },
                         })
+                        mongoMutex.Unlock()
+                case "execve":
+                        // 父进程在不在？不在扔
+                        res, err = pidCol.Finddoc(bson.M{"pid": cooked.ppid})
+                        if err != nil || len(res) != 1 {
+                                break
+                        }
+
+                        // 首先检查进程是否存在，如不存在则为之创建
+                        res, err = pidCol.Finddoc(bson.M{"pid": cooked.pid})
                         if err != nil {
-                                fmt.Printf("Err insert: %v\n", err)
+                                break
+                        }
+                        mongoMutex.Lock()
+                        if len(res) == 1 {
+                                // 自身已在，直接记录
+                                pidCol.UpdateOne(bson.M{"pid": cooked.pid}, bson.M{
+                                        "$push": bson.M{
+                                                "execve": bson.M{
+                                                        "timestamp": cooked.timestamp,
+                                                        "args":      cooked.argv,
+                                                },
+                                        },
+                                })
+                        } else {
+                                // 先fork抵达，插入
+                                pidCol.InsertOne(bson.M{
+                                        "children": []bson.M{},
+                                        "execve": []bson.M{
+                                                {
+                                                        "timestamp": cooked.timestamp,
+                                                        "execve":    cooked.argv,
+                                                },
+                                        },
+                                })
                         }
+                        mongoMutex.Unlock()
                 case "exit", "exit_group":
-                        // TODO: 记得补全退出逻辑
-                        // 上哪找exit code呢？
+                        go deletePid(cooked)
                 }
         }
 }
 
-func connect() (*mongo.Client, error) {
-        client, err := mongo.NewClient(options.Client().ApplyURI("mongodb://localhost:27017"))
-
-        if err != nil {
-                return nil, err
-        }
-
-        ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)
-        err = client.Connect(ctx)
-
-        if err != nil {
-                return nil, err
-        }
-
-        return client, nil
-}
-
-func findDocuments(client *mongo.Client, dbName, colName string, filter bson.M) ([]bson.M, error) {
-        collection := client.Database(dbName).Collection(colName)
-
-        cur, err := collection.Find(context.Background(), filter)
-        if err != nil {
-                return nil, err
-        }
+func deletePid(cooked Event) {
+        time.Sleep(1 * time.Second)
+        mongoMutex.Lock()
+        // 先从老爹那里销户
+        pidCol.UpdateOne(bson.M{"pid": cooked.ppid}, bson.M{
+                "$pull": bson.M{
+                        "children": cooked.pid,
+                },
+        })
 
-        var results []bson.M
-        err = cur.All(context.Background(), &results)
+        // 孩子们需要收容
+        // 不必到children里一个个找，直接看ppid即可
+        pidCol.UpdateMany(bson.M{"ppid": cooked.pid}, bson.M{"ppid": 1})
 
-        return results, err
+        // 可以去死了
+        pidCol.UpdateOne(bson.M{"pid": cooked.pid}, bson.M{
+                "$set": bson.M{
+                        "exit_timestamp": cooked.timestamp,
+                        "exit_code":      cooked.exit_code,
+                },
+        })
+        mongoMutex.Unlock()
 }