Try t use coroutine, but starvationstarvation

5 files changed, 57 insertions, 85 deletions

diff --git a/src/deal.go b/src/deal.go
index 871b7ff..483d4d2 100644
--- a/src/deal.go
+++ b/src/deal.go
@@ -2,20 +2,20 @@ package main
 
 import (
         "fmt"
-        "sync"
         "syscall"
 
         "go.mongodb.org/mongo-driver/bson"
 )
 
 const (
-        dbName     string = "test"
-        pidColName string = "pids"
-        fdColName  string = "fds"
+        dbName      string = "test"
+        pidColName  string = "pids"
+        fdColName   string = "fds"
+        fileColName string = "files"
 )
 
-var mongoMutex sync.Mutex
-var pidCol, fdCol mongoClient
+// var mongoMutex sync.Mutex
+var pidCol, fdCol, fileCol mongoClient
 
 var docRes []bson.M
 var err error
@@ -44,10 +44,14 @@ func deal() {
                 fmt.Printf("Error while initing the mongodb: %v\n", err)
                 return
         }
+        if err = fileCol.init(dbName, fileColName); err != nil {
+                fmt.Printf("Error while initing the mongodb: %v\n", err)
+        }
 
         fmt.Printf("Containerd: %d\n", containerdPid)
         defer pidCol.Disconnect()
         defer fdCol.Disconnect()
+        defer fileCol.Disconnect()
 
         for {
                 cooked, ok = <-cookedChan
@@ -57,11 +61,11 @@ func deal() {
 
                 switch cooked.tag {
                 case NEWPID:
-                        dealNewPid(cooked)
+                        go dealNewPid(cooked)
                 case EXECVE:
-                        dealExecve(cooked)
+                        go dealExecve(cooked)
                 case PIDEXIT:
-                        deletePid(cooked)
+                        go deletePid(cooked)
                 case FILEOPEN:
                         fileOpen(cooked)
                 case FILEWRITE:
@@ -73,9 +77,7 @@ func deal() {
 }
 
 func deletePid(cooked Event) {
-        // TODO: 是否还需要延时？
-        // time.Sleep(1 * time.Second)
-        mongoMutex.Lock()
+        pidCol.GetLock()
         // 先从老爹那里销户
         pidCol.UpdateOne(bson.M{"pid": cooked.ppid}, bson.M{
                 "$pull": bson.M{
@@ -85,7 +87,8 @@ func deletePid(cooked Event) {
 
         // 孩子们需要收容
         // 不必到children里一个个找，直接看ppid即可
-        pidCol.UpdateMany(bson.M{"ppid": cooked.pid}, bson.M{"ppid": 1})
+        // pidCol.UpdateMany(bson.M{"ppid": cooked.pid}, bson.M{"ppid": 1})
+        // 在这套逻辑里，孩子是不需要收容的，因为我们根本就不看ppid来工作
 
         // 可以去死了
         pidCol.UpdateOne(bson.M{"pid": cooked.pid}, bson.M{
@@ -95,8 +98,7 @@ func deletePid(cooked Event) {
                         "exit_signal":    cooked.exit_signal,
                 },
         })
-        mongoMutex.Unlock()
-        fmt.Printf("Exit: %v\t%6d\t%6d\n", cooked.timestamp, cooked.pid, cooked.exit_code)
+        pidCol.Mutex.Unlock()
 }
 
 func dealNewPid(cooked Event) {
@@ -113,7 +115,8 @@ func dealNewPid(cooked Event) {
                 fmt.Printf("Err finding: %v\n", err)
                 return
         }
-        mongoMutex.Lock()
+
+        pidCol.GetLock()
         if len(docRes) != 0 {
                 // 进程原本就存在，换言之别的消息先到了
                 // 所有先行抵达的消息必须保留execve/children字段
@@ -146,7 +149,7 @@ func dealNewPid(cooked Event) {
                         "children": cooked.pid,
                 },
         })
-        mongoMutex.Unlock()
+        pidCol.Mutex.Unlock()
 }
 
 func dealExecve(cooked Event) {
@@ -162,7 +165,8 @@ func dealExecve(cooked Event) {
         if err != nil {
                 return
         }
-        mongoMutex.Lock()
+
+        pidCol.GetLock()
         if len(docRes) == 1 {
                 // 自身已在，直接记录
                 pidCol.UpdateOne(bson.M{"pid": cooked.pid}, bson.M{
@@ -187,12 +191,10 @@ func dealExecve(cooked Event) {
                         },
                 })
         }
-        mongoMutex.Unlock()
+        pidCol.Mutex.Unlock()
 }
 
 func fileOpen(cooked Event) {
-        // fmt.Printf("Open: %6d\t%6d\t%s\n", cooked.ppid, cooked.pid, cooked.pathName)
-
         // 权限检查过了，不必再查
         fdCol.InsertOne(bson.M{
                 "timestamp": cooked.timestamp,
@@ -213,8 +215,6 @@ func fileOpen(cooked Event) {
 }
 
 func fileClose(cooked Event) {
-        // fmt.Printf("Close: %6d\t%6d\t%s\n", cooked.ppid, cooked.pid, cooked.pathName)
-        // 直接看文件表有无记录
         res, err := fdCol.Finddoc(bson.M{
                 "pid":             cooked.pid,
                 "fd":              cooked.syscallParam[0],
@@ -234,8 +234,6 @@ func fileClose(cooked Event) {
 }
 
 func fileWrite(cooked Event) {
-        // fmt.Printf("Write: %6d\t%6d\t%s\n", cooked.ppid, cooked.pid, cooked.pathName)
-        // 直接看文件表有无记录
         res, err := fdCol.Finddoc(bson.M{
                 "pid":             cooked.pid,
                 "fd":              cooked.syscallParam[0],
diff --git a/src/global.go b/src/global.go
index a266b1b..b6635c9 100644
--- a/src/global.go
+++ b/src/global.go
@@ -1,11 +1,8 @@
 package main
 
 import (
-        "fmt"
         "sync"
         "time"
-
-        "go.mongodb.org/mongo-driver/bson/primitive"
 )
 
 type eventType int
@@ -29,50 +26,18 @@ func (et eventType) String() string {
 }
 
 type Event struct {
-        tag          eventType
-        timestamp    time.Time
-        pid, ppid    int
-        syscall      int
-        syscallParam [4]uint64
-        pathName     string
-        argc         int
-        argv         []string
-        cwd          string
-        exit_code    uint64
-        exit_signal  int
-}
-
-func (event Event) String() string {
-        var res string
-        res = fmt.Sprintf("tag: %v\ntimestamp: %v\nppid: %d\npid: %d\n", event.tag, event.timestamp.Local(), event.ppid, event.pid)
-        res += fmt.Sprintf("syscall: %s\nexit_code: %d\nargs: \n", syscallTable[event.syscall], event.exit_code)
-        for i := 0; i < len(event.argv); i++ {
-                res += fmt.Sprintf("\t\"%s\"\n", event.argv[i])
-        }
-        res += "syscallParam: "
-        for i := 0; i < len(event.syscallParam); i++ {
-                res += fmt.Sprintf("\t\"%d\"\n", event.syscallParam[i])
-        }
-        res += "pathName: \"" + event.pathName + "\"\n------\n"
-        return res
-}
-
-type pidExec struct {
-        timestamp time.Time `bson:"timestamp"`
-        execArgs  []string  `bson:"execArgs"`
-}
-
-type pid struct {
-        ID              primitive.ObjectID `bson:"_id,ometempty"`
-        start_timestamp time.Time          `bson:"start_timestamp"`
-        ppid            int                `bson:"ppid"`
-        pid             int                `bson:"pid"`
-        cwd             string             `bson:"cwd"`
-        args            []string           `bson:"args"`
-        execve          []pidExec          `bson:"execve"`
-        children        []int              `bson:"children"`
-        exit_timestamp  time.Time          `bson:"exit_timestamp"`
-        exit_code       uint64             `bson:"exit_code"`
+        tag              eventType
+        timestamp        time.Time
+        pid, tgid        int
+        ppid, parentTgid int
+        syscall          int
+        syscallParam     [4]uint64
+        pathName         string
+        argc             int
+        argv             []string
+        cwd              string
+        exit_code        uint64
+        exit_signal      int
 }
 
 var wg sync.WaitGroup        // 掌管协程
diff --git a/src/godo.go b/src/godo.go
index 2ba32d6..77e677c 100644
--- a/src/godo.go
+++ b/src/godo.go
@@ -44,7 +44,6 @@ func main() {
         var auditCmd *exec.Cmd
 
         pidSyscall := []string{"execve"}
-        // pidSyscall := []string{"fork", "vfork", "clone", "execve", "exit", "exit_group"}
         // 设置监听规则
         for i := 0; i < len(pidSyscall); i++ {
                 auditCmd = exec.Command("auditctl", "-a", "exit,always", "-F", "arch=b64", "-S", pidSyscall[i])
@@ -52,7 +51,7 @@ func main() {
         }
 
         // 监听文件的消息
-        fileSyscall := []string{"open"}
+        fileSyscall := []string{"open", "write", "close"}
         // fileSyscall := []string{"open", "write", "creat", "unlink", "opendir", "mkdir", "rmdir", "chmod", "fchmod", "chown", "fchown", "lchown", "flock"}
         for i := 0; i < len(fileSyscall); i++ {
                 auditCmd = exec.Command("auditctl", "-a", "exit,always", "-F", "arch=b64", "-S", fileSyscall[i])
@@ -118,16 +117,14 @@ func procWatch() error {
                         case netlink.PROC_EVENT_FORK:
                                 data := procEvent.Data.(netlink.ProcEventFork)
                                 cooked := Event{
-                                        tag:       NEWPID,
-                                        ppid:      int(data.ParentTgid),
-                                        pid:       int(data.ChildPid),
-                                        timestamp: time.Now(),
+                                        tag:        NEWPID,
+                                        timestamp:  time.Now(),
+                                        pid:        int(data.ChildPid),
+                                        tgid:       int(data.ChildTgid),
+                                        ppid:       int(data.ParentPid),
+                                        parentTgid: int(data.ParentTgid),
                                 }
                                 checkProc(&cooked)
-                                if data.ChildPid != data.ChildTgid {
-                                        cooked.ppid = int(data.ChildTgid)
-                                        cooked.pid = int(data.ChildPid)
-                                }
                                 cookedChan <- cooked
                         case netlink.PROC_EVENT_EXIT:
                                 data := procEvent.Data.(netlink.ProcEventExit)
@@ -146,7 +143,7 @@ func procWatch() error {
 }
 
 func checkProc(pCooked *Event) {
-        fileName := fmt.Sprintf("/proc/%d/cmdline", pCooked.pid)
+        fileName := fmt.Sprintf("/proc/%d/task/%d/cmdline", pCooked.tgid, pCooked.pid)
         fd, err := os.Open(fileName)
         if err != nil {
                 fmt.Printf("Err: %v\n", err)
@@ -162,7 +159,7 @@ func checkProc(pCooked *Event) {
         pCooked.argc = len(pCooked.argv)
         fd.Close()
 
-        fileName = fmt.Sprintf("/proc/%d/cwd", pCooked.pid)
+        fileName = fmt.Sprintf("/proc/%d/task/%d/cwd", pCooked.tgid, pCooked.pid)
         pCooked.cwd, err = os.Readlink(fileName)
         if err != nil {
                 fmt.Printf("Err readlink %s: %v\n", fileName, err)
diff --git a/src/mongo.go b/src/mongo.go
index 54f9533..3a23131 100644
--- a/src/mongo.go
+++ b/src/mongo.go
@@ -2,6 +2,8 @@ package main
 
 import (
         "context"
+        "fmt"
+        "sync"
         "time"
 
         "go.mongodb.org/mongo-driver/bson"
@@ -13,6 +15,8 @@ type mongoClient struct {
         dbName, colName string
         client          *mongo.Client
         col             *mongo.Collection
+        Mutex           sync.Mutex
+        // Attention: 这把锁是否有必要？
 }
 
 func (mc *mongoClient) init(dbName, colName string) error {
@@ -89,3 +93,12 @@ func (mc *mongoClient) Disconnect() error {
         mc.colName = ""
         return nil
 }
+
+func (mc *mongoClient) GetLock() {
+        for i := 0; i < 20000; {
+                if !mc.Mutex.TryLock() {
+                        i++
+                }
+        }
+        fmt.Printf("Die...\n")
+}
diff --git a/src/organize.go b/src/organize.go
index 8deba53..238509f 100644
--- a/src/organize.go
+++ b/src/organize.go
@@ -47,7 +47,6 @@ func orgnaze() {
                         break
                 }
                 rawEvent = raw.(libaudit.RawAuditMessage)
-                // fmt.Printf("type=%v msg=%s\n", rawEvent.Type, rawEvent.Data)
 
                 switch rawEvent.Type {
                 case auparse.AUDIT_SYSCALL: