From 92d41961e041aadea896f7c829c25a524484fc2f Mon Sep 17 00:00:00 2001
From: Daniel <dhaavi@users.noreply.github.com>
Date: Mon, 20 Apr 2020 13:57:40 +0200
Subject: [PATCH] Fix and improve network database ops

---
 network/clean.go      |  3 +--
 network/connection.go | 46 ++++++++++++++++++-------------------------
 network/database.go   | 10 +++++-----
 3 files changed, 25 insertions(+), 34 deletions(-)

diff --git a/network/clean.go b/network/clean.go
index 2c8cff02..ec51b611 100644
--- a/network/clean.go
+++ b/network/clean.go
@@ -57,8 +57,7 @@ func cleanConnections() (activePIDs map[int]struct{}) {
 					// Step 2: mark end
 					activePIDs[conn.process.Pid] = struct{}{}
 					conn.Ended = now
-					// "save"
-					dbController.PushUpdate(conn)
+					conn.Save()
 				}
 			case conn.Ended < deleteOlderThan:
 				// Step 3: delete
diff --git a/network/connection.go b/network/connection.go
index 88b36e14..b9bef333 100644
--- a/network/connection.go
+++ b/network/connection.go
@@ -230,39 +230,31 @@ func (conn *Connection) SaveWhenFinished() {
 
 // Save saves the connection in the storage and propagates the change through the database system.
 func (conn *Connection) Save() {
-	if conn.ID == "" {
+	conn.UpdateMeta()
 
-		// dns request
-		if !conn.KeyIsSet() {
+	if !conn.KeyIsSet() {
+		if conn.ID == "" {
+			// dns request
+
+			// set key
 			conn.SetKey(fmt.Sprintf("network:tree/%d/%s", conn.process.Pid, conn.Scope))
-			conn.UpdateMeta()
-		}
-		// save to internal state
-		// check if it already exists
-		mapKey := strconv.Itoa(conn.process.Pid) + "/" + conn.Scope
-		dnsConnsLock.Lock()
-		_, ok := dnsConns[mapKey]
-		if !ok {
+			mapKey := strconv.Itoa(conn.process.Pid) + "/" + conn.Scope
+
+			// save
+			dnsConnsLock.Lock()
 			dnsConns[mapKey] = conn
-		}
-		dnsConnsLock.Unlock()
+			dnsConnsLock.Unlock()
+		} else {
+			// network connection
 
-	} else {
-
-		// connection
-		if !conn.KeyIsSet() {
+			// set key
 			conn.SetKey(fmt.Sprintf("network:tree/%d/%s/%s", conn.process.Pid, conn.Scope, conn.ID))
-			conn.UpdateMeta()
-		}
-		// save to internal state
-		// check if it already exists
-		connsLock.Lock()
-		_, ok := conns[conn.ID]
-		if !ok {
-			conns[conn.ID] = conn
-		}
-		connsLock.Unlock()
 
+			// save
+			connsLock.Lock()
+			conns[conn.ID] = conn
+			connsLock.Unlock()
+		}
 	}
 
 	// notify database controller
diff --git a/network/database.go b/network/database.go
index d7dca398..073dcbc0 100644
--- a/network/database.go
+++ b/network/database.go
@@ -77,7 +77,7 @@ func (s *StorageInterface) processQuery(q *query.Query, it *iterator.Iterator) {
 	if slashes <= 1 {
 		// processes
 		for _, proc := range process.All() {
-			if strings.HasPrefix(proc.DatabaseKey(), q.DatabaseKeyPrefix()) {
+			if q.Matches(proc) {
 				it.Next <- proc
 			}
 		}
@@ -86,9 +86,9 @@ func (s *StorageInterface) processQuery(q *query.Query, it *iterator.Iterator) {
 	if slashes <= 2 {
 		// dns scopes only
 		dnsConnsLock.RLock()
-		for _, dnsConns := range dnsConns {
-			if strings.HasPrefix(dnsConns.DatabaseKey(), q.DatabaseKeyPrefix()) {
-				it.Next <- dnsConns
+		for _, dnsConn := range dnsConns {
+			if q.Matches(dnsConn) {
+				it.Next <- dnsConn
 			}
 		}
 		dnsConnsLock.RUnlock()
@@ -98,7 +98,7 @@ func (s *StorageInterface) processQuery(q *query.Query, it *iterator.Iterator) {
 		// connections
 		connsLock.RLock()
 		for _, conn := range conns {
-			if strings.HasPrefix(conn.DatabaseKey(), q.DatabaseKeyPrefix()) {
+			if q.Matches(conn) {
 				it.Next <- conn
 			}
 		}