From 7da7ebf1837a2a4edf817bc2b651bdb0fb4e980a Mon Sep 17 00:00:00 2001
From: Daniel <dhaavi@users.noreply.github.com>
Date: Fri, 11 Jan 2019 22:19:45 +0100
Subject: [PATCH] Fix minor issues

---
 main.go                        |  6 ++----
 network/clean.go               | 29 +++++++++++++++++++----------
 network/connection.go          |  1 +
 network/link.go                |  1 +
 network/reference/protocols.go |  2 ++
 network/unknown.go             |  4 ++--
 6 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/main.go b/main.go
index b34e0f95..4c359384 100644
--- a/main.go
+++ b/main.go
@@ -71,16 +71,14 @@ func main() {
 			fmt.Println("=== END STACK ===")
 		}
 
-		go func() {
-			modules.Shutdown()
-			os.Exit(0)
-		}()
 		go func() {
 			time.Sleep(3 * time.Second)
 			fmt.Println("===== TAKING TOO LONG FOR SHUTDOWN - PRINTING STACK TRACES =====")
 			pprof.Lookup("goroutine").WriteTo(os.Stdout, 2)
 			os.Exit(1)
 		}()
+		modules.Shutdown()
+		os.Exit(0)
 
 	case <-modules.ShuttingDown():
 	}
diff --git a/network/clean.go b/network/clean.go
index 8aebf5d1..e7e671d2 100644
--- a/network/clean.go
+++ b/network/clean.go
@@ -9,9 +9,9 @@ import (
 )
 
 var (
-	cleanerTickDuration = 1 * time.Minute
-	deadLinksTimeout    = 5 * time.Minute
-	thresholdDuration   = 1 * time.Minute
+	cleanerTickDuration = 10 * time.Second
+	deadLinksTimeout    = 3 * time.Minute
+	thresholdDuration   = 3 * time.Minute
 )
 
 func cleaner() {
@@ -19,9 +19,9 @@ func cleaner() {
 		time.Sleep(cleanerTickDuration)
 
 		cleanLinks()
-		time.Sleep(10 * time.Second)
+		time.Sleep(2 * time.Second)
 		cleanConnections()
-		time.Sleep(10 * time.Second)
+		time.Sleep(2 * time.Second)
 		cleanProcesses()
 	}
 }
@@ -32,6 +32,9 @@ func cleanLinks() {
 	now := time.Now().Unix()
 	deleteOlderThan := time.Now().Add(-deadLinksTimeout).Unix()
 
+	// log.Tracef("network.clean: now=%d", now)
+	// log.Tracef("network.clean: deleteOlderThan=%d", deleteOlderThan)
+
 	linksLock.RLock()
 	defer linksLock.RUnlock()
 
@@ -39,11 +42,15 @@ func cleanLinks() {
 	for key, link := range links {
 
 		// delete dead links
-		link.Lock()
-		deleteThis := link.Ended > 0 && link.Ended < deleteOlderThan
-		link.Unlock()
-		if deleteThis {
-			go link.Delete()
+		if link.Ended > 0 {
+			link.Lock()
+			deleteThis := link.Ended < deleteOlderThan
+			link.Unlock()
+			if deleteThis {
+				// log.Tracef("network.clean: deleted %s", link.DatabaseKey())
+				go link.Delete()
+			}
+
 			continue
 		}
 
@@ -59,6 +66,7 @@ func cleanLinks() {
 		// mark end time
 		if !found {
 			link.Ended = now
+			// log.Tracef("network.clean: marked %s as ended.", link.DatabaseKey())
 			go link.Save()
 		}
 
@@ -73,6 +81,7 @@ func cleanConnections() {
 	for _, conn := range connections {
 		conn.Lock()
 		if conn.FirstLinkEstablished < threshold && conn.LinkCount == 0 {
+			// log.Tracef("network.clean: deleted %s", conn.DatabaseKey())
 			go conn.Delete()
 		}
 		conn.Unlock()
diff --git a/network/connection.go b/network/connection.go
index 6369f938..418350ea 100644
--- a/network/connection.go
+++ b/network/connection.go
@@ -268,6 +268,7 @@ func (conn *Connection) Delete() {
 	conn.Meta().Delete()
 	go dbController.PushUpdate(conn)
 	conn.process.RemoveConnection()
+	go conn.process.Save()
 }
 
 // AddLink applies the connection to the link and increases sets counter and timestamps.
diff --git a/network/link.go b/network/link.go
index 118ed7dc..a8deae6b 100644
--- a/network/link.go
+++ b/network/link.go
@@ -263,6 +263,7 @@ func (link *Link) Delete() {
 	link.Meta().Delete()
 	go dbController.PushUpdate(link)
 	link.connection.RemoveLink()
+	go link.connection.Save()
 }
 
 // GetLink fetches a Link from the database from the default namespace for this object
diff --git a/network/reference/protocols.go b/network/reference/protocols.go
index ba1dbd13..326ea885 100644
--- a/network/reference/protocols.go
+++ b/network/reference/protocols.go
@@ -9,6 +9,7 @@ var (
 		6:   "TCP",
 		17:  "UDP",
 		27:  "RDP",
+		58:  "ICMPv6",
 		33:  "DCCP",
 		136: "UDPLite",
 	}
@@ -20,6 +21,7 @@ var (
 		"UDP":     17,
 		"RDP":     27,
 		"DCCP":    33,
+		"ICMPv6":  58,
 		"UDPLite": 136,
 	}
 )
diff --git a/network/unknown.go b/network/unknown.go
index 060fbd7f..fc31aca8 100644
--- a/network/unknown.go
+++ b/network/unknown.go
@@ -9,7 +9,7 @@ const (
 
 var (
 	UnknownDirectConnection = &Connection{
-		Domain:    "D",
+		Domain:    "PI",
 		Direction: Outbound,
 		Verdict:   DROP,
 		Reason:    ReasonUnknownProcess,
@@ -17,7 +17,7 @@ var (
 	}
 
 	UnknownIncomingConnection = &Connection{
-		Domain:    "I",
+		Domain:    "II",
 		Direction: Inbound,
 		Verdict:   DROP,
 		Reason:    ReasonUnknownProcess,