Rewrite network tree saving and cleaning procedures

process/database.go

@@ -5,17 +5,27 @@ import (
 	"sync"
 	"time"
 
+	processInfo "github.com/shirou/gopsutil/process"
+
 	"github.com/Safing/portbase/database"
+	"github.com/Safing/portbase/log"
 	"github.com/Safing/portmaster/profile"
 	"github.com/tevino/abool"
 )
 
+const (
+	processDatabaseNamespace = "network:tree"
+)
+
 var (
 	processes     = make(map[int]*Process)
 	processesLock sync.RWMutex
 
 	dbController     *database.Controller
 	dbControllerFlag = abool.NewBool(false)
+
+	deleteProcessesThreshold       = 15 * time.Minute
+	lastEstablishedUpdateThreshold = 30 * time.Second
 )
 
 // GetProcessFromStorage returns a process from the internal storage.
@@ -28,13 +38,13 @@ func GetProcessFromStorage(pid int) (*Process, bool) {
 }
 
 // All returns a copy of all process objects.
-func All() []*Process {
+func All() map[int]*Process {
 	processesLock.RLock()
 	defer processesLock.RUnlock()
 
-	all := make([]*Process, 0, len(processes))
+	all := make(map[int]*Process)
 	for _, proc := range processes {
-		all = append(all, proc)
+		all[proc.Pid] = proc
 	}
 
 	return all
@@ -46,7 +56,7 @@ func (p *Process) Save() {
 	defer p.Unlock()
 
 	if !p.KeyIsSet() {
-		p.SetKey(fmt.Sprintf("network:tree/%d", p.Pid))
+		p.SetKey(fmt.Sprintf("%s/%d", processDatabaseNamespace, p.Pid))
 		p.CreateMeta()
 	}
 
@@ -89,49 +99,90 @@ func (p *Process) Delete() {
 }
 
 // CleanProcessStorage cleans the storage from old processes.
-func CleanProcessStorage(thresholdDuration time.Duration) {
+func CleanProcessStorage(activeComms map[string]struct{}) {
+	activePIDs, err := getActivePIDs()
+	if err != nil {
+		log.Warningf("process: failed to get list of active PIDs: %s", err)
+		activePIDs = nil
+	}
 	processesCopy := All()
 
-	threshold := time.Now().Add(-thresholdDuration).Unix()
+	threshold := time.Now().Add(-deleteProcessesThreshold).Unix()
 	delete := false
 
 	// clean primary processes
 	for _, p := range processesCopy {
 		p.Lock()
-		if !p.Virtual && p.LastCommEstablished < threshold && p.CommCount == 0 {
-			delete = true
+		// check if internal
+		if p.Pid <= 0 {
+			p.Unlock()
+			continue
+		}
+
+		// has comms?
+		_, hasComms := activeComms[p.DatabaseKey()]
+
+		// virtual / active
+		virtual := p.Virtual
+		active := false
+		if activePIDs != nil {
+			_, active = activePIDs[p.Pid]
 		}
 		p.Unlock()
 
-		if delete {
-			p.Delete()
-			delete = false
+		if !virtual && !hasComms && !active && p.LastCommEstablished < threshold {
+			go p.Delete()
 		}
 	}
 
 	// clean virtual/failed processes
 	for _, p := range processesCopy {
 		p.Lock()
+		// check if internal
+		if p.Pid <= 0 {
+			p.Unlock()
+			continue
+		}
+
 		switch {
 		case p.Error != "":
 			if p.Meta().Created < threshold {
 				delete = true
 			}
 		case p.Virtual:
-			_, parentIsAlive := processes[p.ParentPid]
-			if !parentIsAlive {
+			_, parentIsActive := processesCopy[p.ParentPid]
+			active := true
+			if activePIDs != nil {
+				_, active = activePIDs[p.Pid]
+			}
+			if !parentIsActive || !active {
 				delete = true
 			}
 		}
 		p.Unlock()
 
 		if delete {
-			p.Delete()
+			log.Tracef("process.clean: deleted %s", p.DatabaseKey())
+			go p.Delete()
 			delete = false
 		}
 	}
 }
 
+func getActivePIDs() (map[int]struct{}, error) {
+	procs, err := processInfo.Processes()
+	if err != nil {
+		return nil, err
+	}
+
+	activePIDs := make(map[int]struct{})
+	for _, p := range procs {
+		activePIDs[int(p.Pid)] = struct{}{}
+	}
+
+	return activePIDs, nil
+}
+
 // SetDBController sets the database controller and allows the package to push database updates on a save. It must be set by the package that registers the "network" database.
 func SetDBController(controller *database.Controller) {
 	dbController = controller

process/process.go

@@ -50,7 +50,6 @@ type Process struct {
 
 	FirstCommEstablished int64
 	LastCommEstablished  int64
-	CommCount            uint
 
 	Virtual bool   // This process is either merged into another process or is not needed.
 	Error   string // If this is set, the process is invalid. This is used to cache failing or inexistent processes.
@@ -80,23 +79,53 @@ func (p *Process) AddCommunication() {
 	p.Lock()
 	defer p.Unlock()
 
-	p.CommCount++
+	// check if we should save
+	save := false
+	if p.LastCommEstablished < time.Now().Add(-3*time.Second).Unix() {
+		save = true
+	}
+
+	// update LastCommEstablished
 	p.LastCommEstablished = time.Now().Unix()
 	if p.FirstCommEstablished == 0 {
 		p.FirstCommEstablished = p.LastCommEstablished
 	}
-}
 
-// RemoveCommunication lowers the connection counter by one.
-func (p *Process) RemoveCommunication() {
-	p.Lock()
-	defer p.Unlock()
-
-	if p.CommCount > 0 {
-		p.CommCount--
+	if save {
+		go p.Save()
 	}
 }
 
+// var db = database.NewInterface(nil)
+
+// CountConnections returns the count of connections of a process
+// func (p *Process) CountConnections() int {
+// 	q, err := query.New(fmt.Sprintf("%s/%d/", processDatabaseNamespace, p.Pid)).
+// 		Where(query.Where("Pid", query.Exists, nil)).
+// 		Check()
+// 	if err != nil {
+// 		log.Warningf("process: failed to build query to get connection count of process: %s", err)
+// 		return -1
+// 	}
+//
+// 	it, err := db.Query(q)
+// 	if err != nil {
+// 		log.Warningf("process: failed to query db to get connection count of process: %s", err)
+// 		return -1
+// 	}
+//
+// 	cnt := 0
+// 	for _ = range it.Next {
+// 		cnt++
+// 	}
+// 	if it.Err() != nil {
+// 		log.Warningf("process: failed to query db to get connection count of process: %s", err)
+// 		return -1
+// 	}
+//
+// 	return cnt
+// }
+
 // GetOrFindPrimaryProcess returns the highest process in the tree that matches the given PID.
 func GetOrFindPrimaryProcess(ctx context.Context, pid int) (*Process, error) {
 	log.Tracer(ctx).Tracef("process: getting primary process for PID %d", pid)