Implement routeIDLookup for managing local and remote route IDs

client/internal/peer/route.go

@@ -0,0 +1,90 @@
+package peer
+
+import (
+	"net/netip"
+	"sync"
+
+	log "github.com/sirupsen/logrus"
+
+	nftypes "github.com/netbirdio/netbird/client/internal/netflow/types"
+)
+
+type routeIDLookup struct {
+	localMap    sync.Map
+	remoteMap   sync.Map
+	resolvedIPs sync.Map
+}
+
+func (r *routeIDLookup) AddLocalRouteID(resourceID string, route netip.Prefix) {
+	_, exists := r.localMap.LoadOrStore(route, resourceID)
+	if exists {
+		log.Tracef("resourceID %s already exists in local map", resourceID)
+	}
+}
+
+func (r *routeIDLookup) RemoveLocalRouteID(route netip.Prefix) {
+	r.localMap.Delete(route)
+}
+
+func (r *routeIDLookup) AddRemoteRouteID(resourceID string, route netip.Prefix) {
+	_, exists := r.remoteMap.LoadOrStore(route, resourceID)
+	if exists {
+		log.Tracef("resourceID %s already exists in remote map", resourceID)
+	}
+}
+
+func (r *routeIDLookup) RemoveRemoteRouteID(route netip.Prefix) {
+	r.remoteMap.Delete(route)
+}
+
+func (r *routeIDLookup) AddResolvedIP(resourceID string, route netip.Prefix) {
+	r.resolvedIPs.Store(route, resourceID)
+}
+
+func (r *routeIDLookup) RemoveResolvedIP(route netip.Prefix) {
+	r.resolvedIPs.Delete(route)
+}
+
+func (r *routeIDLookup) Lookup(src, dst netip.Addr, direction nftypes.Direction) (srcResourceID, dstResourceID string) {
+
+	// TODO: check resolved ip's first
+
+	switch direction {
+	case nftypes.Ingress:
+		if srcResourceID == "" || dstResourceID == "" {
+			r.remoteMap.Range(func(key, value interface{}) bool {
+				if key.(netip.Prefix).Contains(src) {
+					srcResourceID = value.(string)
+
+				} else if key.(netip.Prefix).Contains(dst) {
+					dstResourceID = value.(string)
+				}
+
+				if srcResourceID != "" && dstResourceID != "" {
+					return false
+				}
+
+				return true
+			})
+		}
+	case nftypes.Egress:
+		if srcResourceID == "" || dstResourceID == "" {
+			r.localMap.Range(func(key, value interface{}) bool {
+				if key.(netip.Prefix).Contains(src) {
+					srcResourceID = value.(string)
+
+				} else if key.(netip.Prefix).Contains(dst) {
+					dstResourceID = value.(string)
+				}
+
+				if srcResourceID != "" && dstResourceID != "" {
+					return false
+				}
+
+				return true
+			})
+		}
+	}
+
+	return srcResourceID, dstResourceID
+}

client/internal/peer/status.go

@@ -178,8 +178,7 @@ type Status struct {
 
 	ingressGwMgr *ingressgw.Manager
 
-	resIdMux sync.Mutex
-	resIdMap map[netip.Prefix]string
+	routeIDLookup routeIDLookup
 }
 
 // NewRecorder returns a new Status instance
@@ -331,9 +330,8 @@ func (d *Status) AddPeerStateRoute(peer string, route string, resourceId string)
 	if err != nil {
 		log.Errorf("failed to parse prefix %s: %v", route, err)
 	} else {
-		d.resIdMux.Lock()
-		d.resIdMap[pref] = resourceId
-		d.resIdMux.Unlock()
+
+		d.routeIDLookup.AddRemoteRouteID(resourceId, pref)
 	}
 
 	// todo: consider to make sense of this notification or not
@@ -357,9 +355,7 @@ func (d *Status) RemovePeerStateRoute(peer string, route string) error {
 	if err != nil {
 		log.Errorf("failed to parse prefix %s: %v", route, err)
 	} else {
-		d.resIdMux.Lock()
-		delete(d.resIdMap, pref)
-		d.resIdMux.Unlock()
+		d.routeIDLookup.RemoveRemoteRouteID(pref)
 	}
 
 	// todo: consider to make sense of this notification or not
@@ -374,24 +370,7 @@ func (d *Status) CheckRoutes(src, dst netip.Addr, direction nftypes.Direction) (
 		return
 	}
 
-	d.mux.Lock()
-	d.resIdMux.Lock()
-	defer d.resIdMux.Unlock()
-	defer d.mux.Unlock()
-
-	for route, resId := range d.resIdMap {
-		if route.Contains(src) {
-			srcResId = resId
-		} else if route.Contains(dst) {
-			dstResId = resId
-		}
-
-		if srcResId != "" && dstResId != "" {
-			break
-		}
-	}
-
-	return
+	return d.routeIDLookup.Lookup(src, dst, direction)
 }
 
 func (d *Status) UpdatePeerICEState(receivedState State) error {
@@ -607,6 +586,50 @@ func (d *Status) UpdateLocalPeerState(localPeerState LocalPeerState) {
 	d.notifyAddressChanged()
 }
 
+// AddLocalPeerStateRoute adds a route to the local peer state
+func (d *Status) AddLocalPeerStateRoute(route, resourceId string) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	pref, err := netip.ParsePrefix(route)
+	if err != nil {
+		log.Errorf("failed to parse prefix %s: %v", route, err)
+		return
+	}
+
+	if d.localPeer.Routes == nil {
+		d.localPeer.Routes = map[string]struct{}{}
+	}
+
+	d.localPeer.Routes[route] = struct{}{}
+
+	d.routeIDLookup.AddLocalRouteID(resourceId, pref)
+}
+
+// RemoveLocalPeerStateRoute removes a route from the local peer state
+func (d *Status) RemoveLocalPeerStateRoute(route, resourceId string) {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	pref, err := netip.ParsePrefix(route)
+	if err != nil {
+		log.Errorf("failed to parse prefix %s: %v", route, err)
+		return
+	}
+
+	delete(d.localPeer.Routes, route)
+
+	d.routeIDLookup.RemoveLocalRouteID(pref)
+}
+
+// CleanLocalPeerStateRoutes cleans all routes from the local peer state
+func (d *Status) CleanLocalPeerStateRoutes() {
+	d.mux.Lock()
+	defer d.mux.Unlock()
+
+	d.localPeer.Routes = map[string]struct{}{}
+}
+
 // CleanLocalPeerState cleans local peer status
 func (d *Status) CleanLocalPeerState() {
 	d.mux.Lock()
@@ -700,11 +723,9 @@ func (d *Status) UpdateResolvedDomainsStates(originalDomain domain.Domain, resol
 		ParentDomain: originalDomain,
 	}
 
-	d.resIdMux.Lock()
 	for _, prefix := range prefixes {
-		d.resIdMap[prefix] = resourceId
+		d.routeIDLookup.AddResolvedIP(resourceId, prefix)
 	}
-	d.resIdMux.Unlock()
 }
 
 func (d *Status) DeleteResolvedDomainsStates(domain domain.Domain) {
@@ -716,11 +737,9 @@ func (d *Status) DeleteResolvedDomainsStates(domain domain.Domain) {
 		if v.ParentDomain == domain {
 			delete(d.resolvedDomainsStates, k)
 
-			d.resIdMux.Lock()
 			for _, prefix := range v.Prefixes {
-				delete(d.resIdMap, prefix)
+				d.routeIDLookup.RemoveResolvedIP(prefix)
 			}
-			d.resIdMux.Unlock()
 		}
 	}
 }