Move LMS scoring from nameserver to firewall

network/connection.go

@@ -87,10 +87,10 @@ func NewConnectionFromDNSRequest(ctx context.Context, fqdn string, cnames []stri
 	timestamp := time.Now().Unix()
 	dnsConn := &Connection{
 		Scope: fqdn,
-		Entity: (&intel.Entity{
+		Entity: &intel.Entity{
 			Domain: fqdn,
 			CNAME:  cnames,
-		}),
+		},
 		process: proc,
 		Started: timestamp,
 		Ended:   timestamp,
@@ -123,20 +123,20 @@ func NewConnectionFromFirstPacket(pkt packet.Packet) *Connection {
 		default: // netutils.Invalid
 			scope = IncomingInvalid
 		}
-		entity = (&intel.Entity{
+		entity = &intel.Entity{
 			IP:       pkt.Info().Src,
 			Protocol: uint8(pkt.Info().Protocol),
 			Port:     pkt.Info().SrcPort,
-		})
+		}
 
 	} else {
 
 		// outbound connection
-		entity = (&intel.Entity{
+		entity = &intel.Entity{
 			IP:       pkt.Info().Dst,
 			Protocol: uint8(pkt.Info().Protocol),
 			Port:     pkt.Info().DstPort,
-		})
+		}
 
 		// check if we can find a domain for that IP
 		ipinfo, err := resolver.GetIPInfo(pkt.Info().Dst.String())

network/dns.go

@@ -23,15 +23,22 @@ var (
 	unidentifiedProcessScopePrefix = strconv.Itoa(process.UnidentifiedProcessID) + "/"
 )
 
+func getDNSRequestCacheKey(pid int, fqdn string) string {
+	return strconv.Itoa(pid) + "/" + fqdn
+}
+
 func removeOpenDNSRequest(pid int, fqdn string) {
 	openDNSRequestsLock.Lock()
 	defer openDNSRequestsLock.Unlock()
 
-	key := strconv.Itoa(pid) + "/" + fqdn
+	key := getDNSRequestCacheKey(pid, fqdn)
 	_, ok := openDNSRequests[key]
 	if ok {
 		delete(openDNSRequests, key)
-	} else if pid != process.UnidentifiedProcessID {
+		return
+	}
+
+	if pid != process.UnidentifiedProcessID {
 		// check if there is an open dns request from an unidentified process
 		delete(openDNSRequests, unidentifiedProcessScopePrefix+fqdn)
 	}
@@ -42,26 +49,24 @@ func SaveOpenDNSRequest(conn *Connection) {
 	openDNSRequestsLock.Lock()
 	defer openDNSRequestsLock.Unlock()
 
-	key := strconv.Itoa(conn.process.Pid) + "/" + conn.Scope
-
-	existingConn, ok := openDNSRequests[key]
-	if ok {
+	key := getDNSRequestCacheKey(conn.process.Pid, conn.Scope)
+	if existingConn, ok := openDNSRequests[key]; ok {
 		existingConn.Lock()
 		defer existingConn.Unlock()
-
 		existingConn.Ended = conn.Started
-	} else {
-		openDNSRequests[key] = conn
+		return
 	}
+
+	openDNSRequests[key] = conn
 }
 
 func openDNSRequestWriter(ctx context.Context) error {
 	ticker := time.NewTicker(writeOpenDNSRequestsTickDuration)
+	defer ticker.Stop()
 
 	for {
 		select {
 		case <-ctx.Done():
-			ticker.Stop()
 			return nil
 		case <-ticker.C:
 			writeOpenDNSRequestsToDB()