Simplify state handling

miguelangel-nubla · Jan 2, 2025 · eeb50c6 · eeb50c6
1 parent febc6d0
commit eeb50c6
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 52 deletions.
diff --git a/addr_collection.go b/addr_collection.go
@@ -92,14 +92,18 @@ func (c *AddrCollection) Copy() *AddrCollection {
 	return result
 }
 
-func (c *AddrCollection) FilterPrefix(prefix netip.Prefix) *AddrCollection {
+func (c *AddrCollection) FilterSubnets(subnets []netip.Prefix) *AddrCollection {
 	c.addressesMutex.RLock()
 	defer c.addressesMutex.RUnlock()
 
 	result := NewAddrCollection()
+
 	for _, addr := range c.addresses {
-		if prefix.Contains(addr.Addr.WithZone("")) {
-			result.Enlist(addr)
+		ip := addr.Addr.WithZone("")
+		for _, sub := range subnets {
+			if sub.Contains(ip) {
+				result.Enlist(addr)
+			}
 		}
 	}
 

diff --git a/state.go b/state.go
@@ -3,56 +3,40 @@ package ipv6disc
 import (
 	"fmt"
 	"net"
-	"net/netip"
 	"sort"
 	"strings"
 	"sync"
-	"time"
 )
 
 type State struct {
-	macs                    map[string]*AddrCollection
-	macsMutex               sync.RWMutex
-	addrDefaultLifetime     time.Duration
-	addrDefaultOnExpiration func(*Addr, AddrExpirationRemainingEvents)
+	macs      map[string]*AddrCollection
+	macsMutex sync.RWMutex
 }
 
 // accepts default TTL and onExpiration function
-func (s *State) Enlist(hw net.HardwareAddr, netipAddr netip.Addr, ttl time.Duration, onExpiration func(*Addr, AddrExpirationRemainingEvents)) (*Addr, bool) {
+func (s *State) Enlist(addr *Addr) (*Addr, bool) {
 	s.macsMutex.Lock()
 	defer s.macsMutex.Unlock()
 
-	mac := hw.String()
+	mac := addr.Hw.String()
 	_, exists := s.macs[mac]
 	if !exists {
 		s.macs[mac] = NewAddrCollection()
 	}
 
-	if ttl == 0 {
-		ttl = s.addrDefaultLifetime
-	}
-
-	if onExpiration == nil {
-		onExpiration = s.addrDefaultOnExpiration
-	}
-
-	newAddr := NewAddr(hw, netipAddr, ttl, onExpiration)
-
-	return s.macs[mac].Enlist(newAddr)
+	return s.macs[mac].Enlist(addr)
 }
 
-func (s *State) Filter(hws []net.HardwareAddr, prefixes []netip.Prefix) *AddrCollection {
+func (s *State) FilterMACs(hws []net.HardwareAddr) *AddrCollection {
 	results := NewAddrCollection()
 
 	s.macsMutex.Lock()
 	defer s.macsMutex.Unlock()
 
-	for _, prefix := range prefixes {
-		for _, hw := range hws {
-			collection, exists := s.macs[hw.String()]
-			if exists {
-				results.Join(collection.FilterPrefix(prefix))
-			}
+	for _, hw := range hws {
+		collection, exists := s.macs[hw.String()]
+		if exists {
+			results.Join(collection)
 		}
 	}
 
@@ -85,10 +69,8 @@ func (s *State) PrettyPrint(prefix string) string {
 	return result.String()
 }
 
-func NewState(lifetime time.Duration) *State {
+func NewState() *State {
 	return &State{
-		macs:                    make(map[string]*AddrCollection),
-		addrDefaultLifetime:     lifetime,
-		addrDefaultOnExpiration: func(addr *Addr, remainingEvents AddrExpirationRemainingEvents) {},
+		macs: make(map[string]*AddrCollection),
 	}
 }
diff --git a/worker.go b/worker.go
@@ -31,7 +31,7 @@ type Worker struct {
 
 func NewWorker(logger *zap.SugaredLogger, rediscover time.Duration, lifetime time.Duration) *Worker {
 	return &Worker{
-		State:      NewState(lifetime),
+		State:      NewState(),
 		logger:     logger,
 		rediscover: rediscover,
 		lifetime:   lifetime,
@@ -94,47 +94,48 @@ func (w *Worker) StartInterfaceAddr(iface net.Interface, addr netip.Addr) {
 	var ssdpConn *ssdp.Conn
 	var wsdConn *wsd.Conn
 
-	w.State.addrDefaultOnExpiration = func(addr *Addr, remainingEvents AddrExpirationRemainingEvents) {
+	addrOnExpiration := func(expiredAddr *Addr, remainingEvents AddrExpirationRemainingEvents) {
 		if remainingEvents == 0 {
 			w.logger.Infow("host expired",
-				zap.String("ipv6", netip.AddrFrom16(addr.As16()).String()),
-				zap.String("mac", addr.Hw.String()),
-				zap.String("iface", addr.Zone()),
+				zap.String("ipv6", netip.AddrFrom16(expiredAddr.As16()).String()),
+				zap.String("mac", expiredAddr.Hw.String()),
+				zap.String("iface", expiredAddr.Zone()),
 			)
 			return
 		}
 
 		w.logger.Debugw("host not seen for a while, pinging",
-			zap.String("ipv6", netip.AddrFrom16(addr.As16()).String()),
-			zap.String("mac", addr.Hw.String()),
-			zap.String("iface", addr.Zone()),
+			zap.String("ipv6", netip.AddrFrom16(expiredAddr.As16()).String()),
+			zap.String("mac", expiredAddr.Hw.String()),
+			zap.String("iface", expiredAddr.Zone()),
 			zap.Int("remainingEvents", int(remainingEvents)),
 		)
 
-		err := pingConn.SendPing(&addr.Addr)
+		err := pingConn.SendPing(&expiredAddr.Addr)
 		if err != nil {
 			w.logger.Errorw("ping failed",
-				zap.String("ipv6", netip.AddrFrom16(addr.As16()).String()),
-				zap.String("mac", addr.Hw.String()),
-				zap.String("iface", addr.Zone()),
+				zap.String("ipv6", netip.AddrFrom16(expiredAddr.As16()).String()),
+				zap.String("mac", expiredAddr.Hw.String()),
+				zap.String("iface", expiredAddr.Zone()),
 				zap.Error(err),
 			)
 		}
 	}
 
-	processNDP := func(netipAddr netip.Addr, netHardwareAddr net.HardwareAddr) {
-		addr, existing := w.State.Enlist(netHardwareAddr, netipAddr, 0, nil)
+	processNDP := func(receivedAddr netip.Addr, receivedNetHardwareAddr net.HardwareAddr) {
+		newAddr := NewAddr(receivedNetHardwareAddr, receivedAddr, w.lifetime, addrOnExpiration)
+		addr, existing := w.State.Enlist(newAddr)
 		if existing {
 			w.logger.Debugw("ttl refreshed",
-				zap.String("ipv6", netipAddr.String()),
-				zap.String("mac", netHardwareAddr.String()),
+				zap.String("ipv6", receivedAddr.String()),
+				zap.String("mac", receivedNetHardwareAddr.String()),
 			)
 		} else {
 			addr.Watch()
 			w.logger.Infow("host identified",
-				zap.String("ipv6", netip.AddrFrom16(netipAddr.As16()).String()),
-				zap.String("mac", netHardwareAddr.String()),
-				zap.String("iface", netipAddr.Zone()),
+				zap.String("ipv6", netip.AddrFrom16(receivedAddr.As16()).String()),
+				zap.String("mac", receivedNetHardwareAddr.String()),
+				zap.String("iface", receivedAddr.Zone()),
 			)
 		}
 	}