Use DNS request connections to attribute DNS requests

network/clean.go

@@ -118,6 +118,9 @@ func cleanConnections() (activePIDs map[int]struct{}) {
 			conn.Unlock()
 		}
 
+		// rerouted dns requests
+		cleanDNSRequestConnections()
+
 		return nil
 	})
 

network/connection.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"net"
-	"strings"
 	"sync"
 	"time"
 
@@ -293,18 +292,22 @@ func NewConnectionFromDNSRequest(ctx context.Context, fqdn string, cnames []stri
 	}
 
 	// Check if the dns request connection was reported with process info.
-	dnsRequestConnID := pi.CreateConnectionID()
-	// Cut the destination, as the dns request may have been redirected and we
-	// don't know the original destination.
-	dnsRequestConnIDPrefix, ok := strings.CutSuffix(dnsRequestConnID, "<nil>-0")
-	if !ok {
-		log.Tracer(ctx).Warningf("network: unexpected connection ID for finding dns requests connection: %s", dnsRequestConnID)
-	}
-	// Find matching dns request connection.
-	dnsRequestConn, ok := conns.findByPrefix(dnsRequestConnIDPrefix)
-	if ok && dnsRequestConn.PID != process.UndefinedProcessID {
-		log.Tracer(ctx).Debugf("network: found matching dns request connection %s", dnsRequestConn)
+	var proc *process.Process
+	dnsRequestConn, ok := GetDNSRequestConnection(pi)
+	switch {
+	case !ok:
+		// No dns request connection found.
+	case dnsRequestConn.PID < 0:
+		// Process is not identified or is special.
+	case dnsRequestConn.Ended > 0 && dnsRequestConn.Ended < time.Now().Unix()-3:
+		// Connection has already ended (too long ago).
+		log.Tracer(ctx).Debugf("network: found ended dns request connection %s for dns request for %s", dnsRequestConn, fqdn)
+	default:
+		log.Tracer(ctx).Debugf("network: found matching dns request connection %s", dnsRequestConn.String())
+		// Inherit PID.
 		pi.PID = dnsRequestConn.PID
+		// Inherit process struct itself, as the PID may already be re-used.
+		proc = dnsRequestConn.process
 	}
 
 	// Find process by remote IP/Port.
@@ -316,7 +319,9 @@ func NewConnectionFromDNSRequest(ctx context.Context, fqdn string, cnames []stri
 	}
 
 	// Get process and profile with PID.
-	proc, _ := process.GetProcessWithProfile(ctx, pi.PID)
+	if proc == nil {
+		proc, _ = process.GetProcessWithProfile(ctx, pi.PID)
+	}
 
 	timestamp := time.Now().Unix()
 	dnsConn := &Connection{
@@ -1017,6 +1022,8 @@ func (conn *Connection) SetInspectorData(newInspectorData map[uint8]interface{})
 // String returns a string representation of conn.
 func (conn *Connection) String() string {
 	switch {
+	case conn.process == nil || conn.Entity == nil:
+		return conn.ID
 	case conn.Inbound:
 		return fmt.Sprintf("%s <- %s", conn.process, conn.Entity.IP)
 	case conn.Entity.Domain != "":

network/connection_store.go

@@ -40,7 +40,7 @@ func (cs *connectionStore) get(id string) (*Connection, bool) {
 
 // findByPrefix returns the first connection where the key matches the given prefix.
 // If the prefix matches multiple entries, the result is not deterministic.
-func (cs *connectionStore) findByPrefix(prefix string) (*Connection, bool) {
+func (cs *connectionStore) findByPrefix(prefix string) (*Connection, bool) { //nolint:unused
 	cs.rw.RLock()
 	defer cs.rw.RUnlock()
 

network/dns.go

@@ -12,12 +12,16 @@ import (
 
 	"github.com/safing/portbase/log"
 	"github.com/safing/portmaster/nameserver/nsutil"
+	"github.com/safing/portmaster/network/packet"
 	"github.com/safing/portmaster/process"
 	"github.com/safing/portmaster/resolver"
 )
 
 var (
-	openDNSRequests     = make(map[string]*Connection) // key: <pid>/fqdn
+	dnsRequestConnections     = make(map[string]*Connection) // key: <protocol>-<local ip>-<local port>
+	dnsRequestConnectionsLock sync.RWMutex
+
+	openDNSRequests     = make(map[string]*Connection) // key: <pid>/<fqdn>
 	openDNSRequestsLock sync.Mutex
 
 	supportedDomainToIPRecordTypes = []uint16{
@@ -38,6 +42,82 @@ const (
 	openDNSRequestLimit = 3 * time.Second
 )
 
+func getDNSRequestConnectionKey(packetInfo *packet.Info) (id string, ok bool) {
+	// We only support protocols with ports.
+	if packetInfo.SrcPort == 0 {
+		return "", false
+	}
+
+	return fmt.Sprintf("%d-%s-%d", packetInfo.Protocol, packetInfo.Src, packetInfo.SrcPort), true
+}
+
+// SaveDNSRequestConnection saves a dns request connection for later retrieval.
+func SaveDNSRequestConnection(conn *Connection, pkt packet.Packet) {
+	// Check connection.
+	if conn.PID == process.UndefinedProcessID {
+		log.Tracer(pkt.Ctx()).Tracef("network: not saving dns request connection because the PID is undefined")
+		return
+	}
+
+	// Create key.
+	key, ok := getDNSRequestConnectionKey(pkt.Info())
+	if !ok {
+		log.Tracer(pkt.Ctx()).Debugf("network: not saving dns request connection %s because the protocol is not supported", pkt)
+		return
+	}
+
+	// Add or update DNS request connection.
+	log.Tracer(pkt.Ctx()).Tracef("network: saving %s with PID %d as dns request connection for fast DNS request attribution", pkt, conn.PID)
+	dnsRequestConnectionsLock.Lock()
+	defer dnsRequestConnectionsLock.Unlock()
+	dnsRequestConnections[key] = conn
+}
+
+// GetDNSRequestConnection returns a saved dns request connection.
+func GetDNSRequestConnection(packetInfo *packet.Info) (conn *Connection, ok bool) {
+	// Make key.
+	key, ok := getDNSRequestConnectionKey(packetInfo)
+	if !ok {
+		return nil, false
+	}
+
+	// Get and return
+	dnsRequestConnectionsLock.RLock()
+	defer dnsRequestConnectionsLock.RUnlock()
+
+	conn, ok = dnsRequestConnections[key]
+	return
+}
+
+// deleteDNSRequestConnection removes a connection from the dns request connections.
+func deleteDNSRequestConnection(packetInfo *packet.Info) { //nolint:unused,deadcode
+	dnsRequestConnectionsLock.Lock()
+	defer dnsRequestConnectionsLock.Unlock()
+
+	key, ok := getDNSRequestConnectionKey(packetInfo)
+	if ok {
+		delete(dnsRequestConnections, key)
+	}
+}
+
+// cleanDNSRequestConnections deletes old DNS request connections.
+func cleanDNSRequestConnections() {
+	deleteOlderThan := time.Now().Unix() - 3
+
+	dnsRequestConnectionsLock.Lock()
+	defer dnsRequestConnectionsLock.Unlock()
+
+	for key, conn := range dnsRequestConnections {
+		conn.Lock()
+
+		if conn.Ended > 0 && conn.Ended < deleteOlderThan {
+			delete(dnsRequestConnections, key)
+		}
+
+		conn.Unlock()
+	}
+}
+
 // IsSupportDNSRecordType returns whether the given DSN RR type is supported
 // by the network package, as in the requests are specially handled and can be
 // "merged" into the resulting connection.