Merge pull request #57 from safing/feature/network-state-improvements

Transition to table structs, use OnceAgain for fetching tables

network/state/exists.go

@@ -19,38 +19,33 @@ func Exists(pktInfo *packet.Info, now time.Time) (exists bool) {
 
 	switch {
 	case pktInfo.Version == packet.IPv4 && pktInfo.Protocol == packet.TCP:
-		tcp4Lock.Lock()
+		return tcp4Table.exists(pktInfo)
-		defer tcp4Lock.Unlock()
-		return existsTCP(tcp4Connections, pktInfo)
 
 	case pktInfo.Version == packet.IPv6 && pktInfo.Protocol == packet.TCP:
-		tcp6Lock.Lock()
+		return tcp6Table.exists(pktInfo)
-		defer tcp6Lock.Unlock()
-		return existsTCP(tcp6Connections, pktInfo)
 
 	case pktInfo.Version == packet.IPv4 && pktInfo.Protocol == packet.UDP:
-		udp4Lock.Lock()
+		return udp4Table.exists(pktInfo, now)
-		defer udp4Lock.Unlock()
-		return existsUDP(udp4Binds, udp4States, pktInfo, now)
 
 	case pktInfo.Version == packet.IPv6 && pktInfo.Protocol == packet.UDP:
-		udp6Lock.Lock()
+		return udp6Table.exists(pktInfo, now)
-		defer udp6Lock.Unlock()
-		return existsUDP(udp6Binds, udp6States, pktInfo, now)
 
 	default:
 		return false
 	}
 }
 
-func existsTCP(connections []*socket.ConnectionInfo, pktInfo *packet.Info) (exists bool) {
+func (table *tcpTable) exists(pktInfo *packet.Info) (exists bool) {
+	table.lock.RLock()
+	defer table.lock.RUnlock()
+
 	localIP := pktInfo.LocalIP()
 	localPort := pktInfo.LocalPort()
 	remoteIP := pktInfo.RemoteIP()
 	remotePort := pktInfo.RemotePort()
 
 	// search connections
-	for _, socketInfo := range connections {
+	for _, socketInfo := range table.connections {
 		if localPort == socketInfo.Local.Port &&
 			remotePort == socketInfo.Remote.Port &&
 			remoteIP.Equal(socketInfo.Remote.IP) &&
@@ -62,12 +57,9 @@ func existsTCP(connections []*socket.ConnectionInfo, pktInfo *packet.Info) (exis
 	return false
 }
 
-func existsUDP(
+func (table *udpTable) exists(pktInfo *packet.Info, now time.Time) (exists bool) {
-	binds []*socket.BindInfo,
+	table.lock.RLock()
-	udpStates map[string]map[string]*udpState,
+	defer table.lock.RUnlock()
-	pktInfo *packet.Info,
-	now time.Time,
-) (exists bool) {
 
 	localIP := pktInfo.LocalIP()
 	localPort := pktInfo.LocalPort()
@@ -77,11 +69,11 @@ func existsUDP(
 	connThreshhold := now.Add(-UDPConnectionTTL)
 
 	// search binds
-	for _, socketInfo := range binds {
+	for _, socketInfo := range table.binds {
 		if localPort == socketInfo.Local.Port &&
 			(socketInfo.Local.IP[0] == 0 || localIP.Equal(socketInfo.Local.IP)) {
 
-			udpConnState, ok := getUDPConnState(socketInfo, udpStates, socket.Address{
+			udpConnState, ok := table.getConnState(socketInfo, socket.Address{
 				IP:   remoteIP,
 				Port: remotePort,
 			})

network/state/info.go

@@ -25,27 +25,27 @@ type Info struct {
 func GetInfo() *Info {
 	info := &Info{}
 
-	tcp4Lock.Lock()
+	tcp4Table.updateTables()
-	updateTCP4Tables()
+	tcp4Table.lock.RLock()
-	info.TCP4Connections = tcp4Connections
+	info.TCP4Connections = tcp4Table.connections
-	info.TCP4Listeners = tcp4Listeners
+	info.TCP4Listeners = tcp4Table.listeners
-	tcp4Lock.Unlock()
+	tcp4Table.lock.RUnlock()
 
-	tcp6Lock.Lock()
+	tcp6Table.updateTables()
-	updateTCP6Tables()
+	tcp6Table.lock.RLock()
-	info.TCP6Connections = tcp6Connections
+	info.TCP6Connections = tcp6Table.connections
-	info.TCP6Listeners = tcp6Listeners
+	info.TCP6Listeners = tcp6Table.listeners
-	tcp6Lock.Unlock()
+	tcp6Table.lock.RUnlock()
 
-	udp4Lock.Lock()
+	udp4Table.updateTable()
-	updateUDP4Table()
+	udp4Table.lock.RLock()
-	info.UDP4Binds = udp4Binds
+	info.UDP4Binds = udp4Table.binds
-	udp4Lock.Unlock()
+	udp4Table.lock.RUnlock()
 
-	udp6Lock.Lock()
+	udp6Table.updateTable()
-	updateUDP6Table()
+	udp6Table.lock.RLock()
-	info.UDP6Binds = udp6Binds
+	info.UDP6Binds = udp6Table.binds
-	udp6Lock.Unlock()
+	udp6Table.lock.RUnlock()
 
 	info.UpdateMeta()
 	return info

network/state/lookup.go

@@ -2,7 +2,6 @@ package state
 
 import (
 	"errors"
-	"sync"
 	"time"
 
 	"github.com/safing/portmaster/network/netutils"
@@ -31,12 +30,8 @@ var (
 )
 
 var (
-	tcp4Lock sync.Mutex
+	baseWaitTime  = 3 * time.Millisecond
-	tcp6Lock sync.Mutex
+	lookupRetries = 7
-	udp4Lock sync.Mutex
-	udp6Lock sync.Mutex
-
-	baseWaitTime = 3 * time.Millisecond
 )
 
 // Lookup looks for the given connection in the system state tables and returns the PID of the associated process and whether the connection is inbound.
@@ -52,36 +47,23 @@ func Lookup(pktInfo *packet.Info) (pid int, inbound bool, err error) {
 
 	switch {
 	case pktInfo.Version == packet.IPv4 && pktInfo.Protocol == packet.TCP:
-		tcp4Lock.Lock()
+		return tcp4Table.lookup(pktInfo)
-		defer tcp4Lock.Unlock()
-		return searchTCP(tcp4Connections, tcp4Listeners, updateTCP4Tables, pktInfo)
 
 	case pktInfo.Version == packet.IPv6 && pktInfo.Protocol == packet.TCP:
-		tcp6Lock.Lock()
+		return tcp6Table.lookup(pktInfo)
-		defer tcp6Lock.Unlock()
-		return searchTCP(tcp6Connections, tcp6Listeners, updateTCP6Tables, pktInfo)
 
 	case pktInfo.Version == packet.IPv4 && pktInfo.Protocol == packet.UDP:
-		udp4Lock.Lock()
+		return udp4Table.lookup(pktInfo)
-		defer udp4Lock.Unlock()
-		return searchUDP(udp4Binds, udp4States, updateUDP4Table, pktInfo)
 
 	case pktInfo.Version == packet.IPv6 && pktInfo.Protocol == packet.UDP:
-		udp6Lock.Lock()
+		return udp6Table.lookup(pktInfo)
-		defer udp6Lock.Unlock()
-		return searchUDP(udp6Binds, udp6States, updateUDP6Table, pktInfo)
 
 	default:
 		return socket.UnidentifiedProcessID, false, errors.New("unsupported protocol for finding process")
 	}
 }
 
-func searchTCP(
+func (table *tcpTable) lookup(pktInfo *packet.Info) (
-	connections []*socket.ConnectionInfo,
-	listeners []*socket.BindInfo,
-	updateTables func() ([]*socket.ConnectionInfo, []*socket.BindInfo),
-	pktInfo *packet.Info,
-) (
 	pid int,
 	inbound bool,
 	err error,
@@ -91,45 +73,48 @@ func searchTCP(
 	localPort := pktInfo.LocalPort()
 
 	// search until we find something
-	for i := 0; i < 7; i++ {
+	for i := 0; i <= lookupRetries; i++ {
+		table.lock.RLock()
+
 		// always search listeners first
-		for _, socketInfo := range listeners {
+		for _, socketInfo := range table.listeners {
 			if localPort == socketInfo.Local.Port &&
 				(socketInfo.Local.IP[0] == 0 || localIP.Equal(socketInfo.Local.IP)) {
+				table.lock.RUnlock()
 				return checkBindPID(socketInfo, true)
 			}
 		}
 
 		// search connections
-		for _, socketInfo := range connections {
+		for _, socketInfo := range table.connections {
 			if localPort == socketInfo.Local.Port &&
 				localIP.Equal(socketInfo.Local.IP) {
+				table.lock.RUnlock()
 				return checkConnectionPID(socketInfo, false)
 			}
 		}
 
-		// we found nothing, we could have been too fast, give the kernel some time to think
+		table.lock.RUnlock()
-		// back off timer: with 3ms baseWaitTime: 3, 6, 9, 12, 15, 18, 21ms - 84ms in total
-		time.Sleep(time.Duration(i+1) * baseWaitTime)
 
-		// refetch lists
+		// every time, except for the last iteration
-		connections, listeners = updateTables()
+		if i < lookupRetries {
+			// we found nothing, we could have been too fast, give the kernel some time to think
+			// back off timer: with 3ms baseWaitTime: 3, 6, 9, 12, 15, 18, 21ms - 84ms in total
+			time.Sleep(time.Duration(i+1) * baseWaitTime)
+
+			// refetch lists
+			table.updateTables()
+		}
 	}
 
 	return socket.UnidentifiedProcessID, false, ErrConnectionNotFound
 }
 
-func searchUDP(
+func (table *udpTable) lookup(pktInfo *packet.Info) (
-	binds []*socket.BindInfo,
-	udpStates map[string]map[string]*udpState,
-	updateTable func() []*socket.BindInfo,
-	pktInfo *packet.Info,
-) (
 	pid int,
 	inbound bool,
 	err error,
 ) {
-
 	localIP := pktInfo.LocalIP()
 	localPort := pktInfo.LocalPort()
 
@@ -140,13 +125,16 @@ func searchUDP(
 	// binding to different addresses. This highly unusual for clients.
 
 	// search until we find something
-	for i := 0; i < 5; i++ {
+	for i := 0; i <= lookupRetries; i++ {
+		table.lock.RLock()
+
 		// search binds
-		for _, socketInfo := range binds {
+		for _, socketInfo := range table.binds {
 			if localPort == socketInfo.Local.Port &&
 				(socketInfo.Local.IP[0] == 0 || // zero IP
 					isInboundMulticast || // inbound broadcast, multicast
 					localIP.Equal(socketInfo.Local.IP)) {
+				table.lock.RUnlock()
 
 				// do not check direction if remoteIP/Port is not given
 				if pktInfo.RemotePort() == 0 {
@@ -154,17 +142,22 @@ func searchUDP(
 				}
 
 				// get direction and return
-				connInbound := getUDPDirection(socketInfo, udpStates, pktInfo)
+				connInbound := table.getDirection(socketInfo, pktInfo)
 				return checkBindPID(socketInfo, connInbound)
 			}
 		}
 
-		// we found nothing, we could have been too fast, give the kernel some time to think
+		table.lock.RUnlock()
-		// back off timer: with 3ms baseWaitTime: 3, 6, 9, 12, 15, 18, 21ms - 84ms in total
-		time.Sleep(time.Duration(i+1) * baseWaitTime)
 
-		// refetch lists
+		// every time, except for the last iteration
-		binds = updateTable()
+		if i < lookupRetries {
+			// we found nothing, we could have been too fast, give the kernel some time to think
+			// back off timer: with 3ms baseWaitTime: 3, 6, 9, 12, 15, 18, 21ms - 84ms in total
+			time.Sleep(time.Duration(i+1) * baseWaitTime)
+
+			// refetch lists
+			table.updateTable()
+		}
 	}
 
 	return socket.UnidentifiedProcessID, pktInfo.Inbound, ErrConnectionNotFound

network/state/tables.go

@@ -2,67 +2,35 @@ package state
 
 import (
 	"github.com/safing/portbase/log"
-	"github.com/safing/portmaster/network/socket"
 )
 
-var (
+func (table *tcpTable) updateTables() {
-	tcp4Connections []*socket.ConnectionInfo
+	table.fetchOnceAgain.Do(func() {
-	tcp4Listeners   []*socket.BindInfo
+		table.lock.Lock()
+		defer table.lock.Unlock()
 
-	tcp6Connections []*socket.ConnectionInfo
+		connections, listeners, err := table.fetchTable()
-	tcp6Listeners   []*socket.BindInfo
+		if err != nil {
+			log.Warningf("state: failed to get TCP%d socket table: %s", table.version, err)
+			return
+		}
 
-	udp4Binds []*socket.BindInfo
+		table.connections = connections
-
+		table.listeners = listeners
-	udp6Binds []*socket.BindInfo
+	})
-)
-
-func updateTCP4Tables() (connections []*socket.ConnectionInfo, listeners []*socket.BindInfo) {
-	var err error
-	connections, listeners, err = getTCP4Table()
-	if err != nil {
-		log.Warningf("state: failed to get TCP4 socket table: %s", err)
-		return
-	}
-
-	tcp4Connections = connections
-	tcp4Listeners = listeners
-	return
 }
 
-func updateTCP6Tables() (connections []*socket.ConnectionInfo, listeners []*socket.BindInfo) {
+func (table *udpTable) updateTable() {
-	var err error
+	table.fetchOnceAgain.Do(func() {
-	connections, listeners, err = getTCP6Table()
+		table.lock.Lock()
-	if err != nil {
+		defer table.lock.Unlock()
-		log.Warningf("state: failed to get TCP6 socket table: %s", err)
-		return
-	}
 
-	tcp6Connections = connections
+		binds, err := table.fetchTable()
-	tcp6Listeners = listeners
+		if err != nil {
-	return
+			log.Warningf("state: failed to get UDP%d socket table: %s", table.version, err)
-}
+			return
-
+		}
-func updateUDP4Table() (binds []*socket.BindInfo) {
+
-	var err error
+		table.binds = binds
-	binds, err = getUDP4Table()
+	})
-	if err != nil {
-		log.Warningf("state: failed to get UDP4 socket table: %s", err)
-		return
-	}
-
-	udp4Binds = binds
-	return
-}
-
-func updateUDP6Table() (binds []*socket.BindInfo) {
-	var err error
-	binds, err = getUDP6Table()
-	if err != nil {
-		log.Warningf("state: failed to get UDP6 socket table: %s", err)
-		return
-	}
-
-	udp6Binds = binds
-	return
 }

network/state/tcp.go

@@ -0,0 +1,31 @@
+package state
+
+import (
+	"sync"
+
+	"github.com/safing/portbase/utils"
+	"github.com/safing/portmaster/network/socket"
+)
+
+type tcpTable struct {
+	version int
+
+	connections []*socket.ConnectionInfo
+	listeners   []*socket.BindInfo
+	lock        sync.RWMutex
+
+	fetchOnceAgain utils.OnceAgain
+	fetchTable     func() (connections []*socket.ConnectionInfo, listeners []*socket.BindInfo, err error)
+}
+
+var (
+	tcp4Table = &tcpTable{
+		version:    4,
+		fetchTable: getTCP4Table,
+	}
+
+	tcp6Table = &tcpTable{
+		version:    6,
+		fetchTable: getTCP6Table,
+	}
+)

network/state/udp.go

@@ -2,12 +2,27 @@ package state
 
 import (
 	"context"
+	"sync"
 	"time"
 
+	"github.com/safing/portbase/utils"
 	"github.com/safing/portmaster/network/packet"
 	"github.com/safing/portmaster/network/socket"
 )
 
+type udpTable struct {
+	version int
+
+	binds []*socket.BindInfo
+	lock  sync.RWMutex
+
+	fetchOnceAgain utils.OnceAgain
+	fetchTable     func() (binds []*socket.BindInfo, err error)
+
+	states     map[string]map[string]*udpState
+	statesLock sync.Mutex
+}
+
 type udpState struct {
 	inbound  bool
 	lastSeen time.Time
@@ -25,12 +40,38 @@ const (
 )
 
 var (
-	udp4States = make(map[string]map[string]*udpState) // locked with udp4Lock
+	udp4Table = &udpTable{
-	udp6States = make(map[string]map[string]*udpState) // locked with udp6Lock
+		version:    4,
+		fetchTable: getUDP4Table,
+		states:     make(map[string]map[string]*udpState),
+	}
+
+	udp6Table = &udpTable{
+		version:    6,
+		fetchTable: getUDP6Table,
+		states:     make(map[string]map[string]*udpState),
+	}
 )
 
-func getUDPConnState(socketInfo *socket.BindInfo, udpStates map[string]map[string]*udpState, remoteAddress socket.Address) (udpConnState *udpState, ok bool) {
+// CleanUDPStates cleans the udp connection states which save connection directions.
-	bindMap, ok := udpStates[makeUDPStateKey(socketInfo.Local)]
+func CleanUDPStates(_ context.Context) {
+	now := time.Now().UTC()
+
+	udp4Table.updateTable()
+	udp4Table.cleanStates(now)
+
+	udp6Table.updateTable()
+	udp6Table.cleanStates(now)
+}
+
+func (table *udpTable) getConnState(
+	socketInfo *socket.BindInfo,
+	remoteAddress socket.Address,
+) (udpConnState *udpState, ok bool) {
+	table.statesLock.Lock()
+	defer table.statesLock.Unlock()
+
+	bindMap, ok := table.states[makeUDPStateKey(socketInfo.Local)]
 	if ok {
 		udpConnState, ok = bindMap[makeUDPStateKey(remoteAddress)]
 		return
@@ -39,13 +80,19 @@ func getUDPConnState(socketInfo *socket.BindInfo, udpStates map[string]map[strin
 	return nil, false
 }
 
-func getUDPDirection(socketInfo *socket.BindInfo, udpStates map[string]map[string]*udpState, pktInfo *packet.Info) (connDirection bool) {
+func (table *udpTable) getDirection(
+	socketInfo *socket.BindInfo,
+	pktInfo *packet.Info,
+) (connDirection bool) {
+	table.statesLock.Lock()
+	defer table.statesLock.Unlock()
+
 	localKey := makeUDPStateKey(socketInfo.Local)
 
-	bindMap, ok := udpStates[localKey]
+	bindMap, ok := table.states[localKey]
 	if !ok {
 		bindMap = make(map[string]*udpState)
-		udpStates[localKey] = bindMap
+		table.states[localKey] = bindMap
 	}
 
 	remoteKey := makeUDPStateKey(socket.Address{
@@ -65,38 +112,25 @@ func getUDPDirection(socketInfo *socket.BindInfo, udpStates map[string]map[strin
 	return udpConnState.inbound
 }
 
-// CleanUDPStates cleans the udp connection states which save connection directions.
+func (table *udpTable) cleanStates(now time.Time) {
-func CleanUDPStates(_ context.Context) {
-	now := time.Now().UTC()
 
-	udp4Lock.Lock()
-	updateUDP4Table()
-	cleanStates(udp4Binds, udp4States, now)
-	udp4Lock.Unlock()
-
-	udp6Lock.Lock()
-	updateUDP6Table()
-	cleanStates(udp6Binds, udp6States, now)
-	udp6Lock.Unlock()
-}
-
-func cleanStates(
-	binds []*socket.BindInfo,
-	udpStates map[string]map[string]*udpState,
-	now time.Time,
-) {
 	// compute thresholds
 	threshold := now.Add(-UDPConnStateTTL)
 	shortThreshhold := now.Add(-UDPConnStateShortenedTTL)
 
 	// make lookup map of all active keys
 	bindKeys := make(map[string]struct{})
-	for _, socketInfo := range binds {
+	table.lock.RLock()
+	for _, socketInfo := range table.binds {
 		bindKeys[makeUDPStateKey(socketInfo.Local)] = struct{}{}
 	}
+	table.lock.RUnlock()
+
+	table.statesLock.Lock()
+	defer table.statesLock.Unlock()
 
 	// clean the udp state storage
-	for localKey, bindMap := range udpStates {
+	for localKey, bindMap := range table.states {
 		if _, active := bindKeys[localKey]; active {
 			// clean old entries
 			for remoteKey, udpConnState := range bindMap {
@@ -114,7 +148,7 @@ func cleanStates(
 			}
 		} else {
 			// delete the whole thing
-			delete(udpStates, localKey)
+			delete(table.states, localKey)
 		}
 	}
 }