feat: add packet filtering support

netsim/packet/packet.go

@@ -182,3 +182,34 @@ func NewNetworkDeviceIOChannels() (chan *Packet, chan *Packet) {
 	output := make(chan *Packet, DefaultBufferChannel)
 	return input, output
 }
+
+// Target represents what to do with a [*Packet]
+// similarly to `iptables` target.
+type Target int
+
+const (
+	// ACCEPT lets the [*Packet] continue through the chain.
+	ACCEPT Target = iota
+
+	// DROP silently discards the [*Packet].
+	DROP
+)
+
+// Filter processes [*Packet] and determines its fate.
+//
+// The Filter method returns the [Target] and optionally
+// a list of new packets to inject.
+type Filter interface {
+	Filter(pkt *Packet) (Target, []*Packet)
+}
+
+// FilterFunc allows using a function as a [Filter].
+type FilterFunc func(pkt *Packet) (Target, []*Packet)
+
+// Ensure [FilterFunc] implements the [Filter] interface.
+var _ Filter = FilterFunc(nil)
+
+// Filter implements the [Filter] interface.
+func (fx FilterFunc) Filter(p *Packet) (Target, []*Packet) {
+	return fx(p)
+}

netsim/router/router.go

@@ -6,6 +6,7 @@ package router
 import (
 	"errors"
 	"net/netip"
+	"sync"
 
 	"github.com/rbmk-project/x/netsim/packet"
 )
@@ -15,18 +16,33 @@ type Router struct {
 	// devs tracks attached [packet.NetworkDevice].
 	devs []packet.NetworkDevice
 
+	// filtermu protects access to filters.
+	filtermu sync.RWMutex
+
+	// filters contains pre-routing packet filters.
+	filters []packet.Filter
+
 	// srt is the static routing table.
 	srt map[netip.Addr]packet.NetworkDevice
 }
 
 // New creates a new [*Router].
 func New() *Router {
 	return &Router{
-		devs: make([]packet.NetworkDevice, 0),
-		srt:  make(map[netip.Addr]packet.NetworkDevice),
+		devs:     make([]packet.NetworkDevice, 0),
+		filtermu: sync.RWMutex{},
+		filters:  make([]packet.Filter, 0),
+		srt:      make(map[netip.Addr]packet.NetworkDevice),
 	}
 }
 
+// AddFilter adds a packet filter to the router.
+func (r *Router) AddFilter(pf packet.Filter) {
+	r.filtermu.Lock()
+	r.filters = append(r.filters, pf)
+	r.filtermu.Unlock()
+}
+
 // Attach attaches a [packet.NetworkDevice] to the [*Router].
 func (r *Router) Attach(dev packet.NetworkDevice) {
 	r.devs = append(r.devs, dev)
@@ -47,11 +63,41 @@ func (r *Router) readLoop(dev packet.NetworkDevice) {
 		case <-dev.EOF():
 			return
 		case pkt := <-dev.Output():
-			r.route(pkt)
+			r.handle(pkt)
 		}
 	}
 }
 
+// handle handles a packet by applying filters and routing it.
+func (r *Router) handle(pkt *packet.Packet) error {
+	// Get a consistent view of filters
+	r.filtermu.RLock()
+	filters := make([]packet.Filter, len(r.filters))
+	copy(filters, r.filters)
+	r.filtermu.RUnlock()
+
+	// Apply filters
+	for _, pf := range filters {
+		target, inject := pf.Filter(pkt)
+
+		// Handle any packets to inject
+		for _, p := range inject {
+			_ = r.route(p)
+		}
+
+		// Stop processing if packet should be dropped
+		switch target {
+		case packet.DROP:
+			return nil
+		default:
+			// Continue processing
+		}
+	}
+
+	// Route the original packet if it wasn't dropped
+	return r.route(pkt)
+}
+
 var (
 	// errTTLExceeded is returned when a packet's TTL is exceeded.
 	errTTLExceeded = errors.New("TTL exceeded in transit")
@@ -77,7 +123,7 @@ func (r *Router) route(pkt *packet.Packet) error {
 		return errNoRouteToHost
 	}
 
-	// Forward packet (non-blocking)
+	// Forward packet (non-blocking).
 	select {
 	case nextHop.Input() <- pkt:
 		return nil