feat: ip rule subscribe

Author: cheina97

rule.go

@@ -80,3 +80,9 @@ type RuleUIDRange struct {
 	Start uint32
 	End   uint32
 }
+
+// RuleUpdate is sent when a rule changes - type is RTM_NEWRULE or RTM_DELRULE.
+type RuleUpdate struct {
+	Type uint16
+	Rule
+}

rule_linux.go

@@ -5,8 +5,10 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"syscall"
 
 	"github.com/vishvananda/netlink/nl"
+	"github.com/vishvananda/netns"
 	"golang.org/x/sys/unix"
 )
 
@@ -119,8 +121,8 @@ func ruleHandle(rule *Rule, req *nl.NetlinkRequest) error {
 		req.AddData(nl.NewRtAttr(nl.FRA_FLOW, b))
 	}
 	if rule.TunID > 0 {
-		b := make([]byte, 4)
-		native.PutUint32(b, uint32(rule.TunID))
+		b := make([]byte, 8)
+		networkOrder.PutUint64(b, uint64(rule.TunID))
 		req.AddData(nl.NewRtAttr(nl.FRA_TUN_ID, b))
 	}
 	if rule.Table >= 256 {
@@ -154,9 +156,7 @@ func ruleHandle(rule *Rule, req *nl.NetlinkRequest) error {
 	}
 
 	if rule.IPProto > 0 {
-		b := make([]byte, 4)
-		native.PutUint32(b, uint32(rule.IPProto))
-		req.AddData(nl.NewRtAttr(nl.FRA_IP_PROTO, b))
+		req.AddData(nl.NewRtAttr(nl.FRA_IP_PROTO, nl.Uint8Attr(uint8(rule.IPProto))))
 	}
 
 	if rule.Dport != nil {
@@ -227,73 +227,11 @@ func (h *Handle) RuleListFiltered(family int, filter *Rule, filterMask uint64) (
 
 	var res = make([]Rule, 0)
 	for i := range msgs {
-		msg := nl.DeserializeRtMsg(msgs[i])
-		attrs, err := nl.ParseRouteAttr(msgs[i][msg.Len():])
+		rule, err := deserializeRule(msgs[i])
 		if err != nil {
 			return nil, err
 		}
 
-		rule := NewRule()
-		rule.Priority = 0 // The default priority from kernel
-
-		rule.Invert = msg.Flags&FibRuleInvert > 0
-		rule.Family = int(msg.Family)
-		rule.Tos = uint(msg.Tos)
-
-		for j := range attrs {
-			switch attrs[j].Attr.Type {
-			case unix.RTA_TABLE:
-				rule.Table = int(native.Uint32(attrs[j].Value[0:4]))
-			case nl.FRA_SRC:
-				rule.Src = &net.IPNet{
-					IP:   attrs[j].Value,
-					Mask: net.CIDRMask(int(msg.Src_len), 8*len(attrs[j].Value)),
-				}
-			case nl.FRA_DST:
-				rule.Dst = &net.IPNet{
-					IP:   attrs[j].Value,
-					Mask: net.CIDRMask(int(msg.Dst_len), 8*len(attrs[j].Value)),
-				}
-			case nl.FRA_FWMARK:
-				rule.Mark = native.Uint32(attrs[j].Value[0:4])
-			case nl.FRA_FWMASK:
-				mask := native.Uint32(attrs[j].Value[0:4])
-				rule.Mask = &mask
-			case nl.FRA_TUN_ID:
-				rule.TunID = uint(native.Uint64(attrs[j].Value[0:8]))
-			case nl.FRA_IIFNAME:
-				rule.IifName = string(attrs[j].Value[:len(attrs[j].Value)-1])
-			case nl.FRA_OIFNAME:
-				rule.OifName = string(attrs[j].Value[:len(attrs[j].Value)-1])
-			case nl.FRA_SUPPRESS_PREFIXLEN:
-				i := native.Uint32(attrs[j].Value[0:4])
-				if i != 0xffffffff {
-					rule.SuppressPrefixlen = int(i)
-				}
-			case nl.FRA_SUPPRESS_IFGROUP:
-				i := native.Uint32(attrs[j].Value[0:4])
-				if i != 0xffffffff {
-					rule.SuppressIfgroup = int(i)
-				}
-			case nl.FRA_FLOW:
-				rule.Flow = int(native.Uint32(attrs[j].Value[0:4]))
-			case nl.FRA_GOTO:
-				rule.Goto = int(native.Uint32(attrs[j].Value[0:4]))
-			case nl.FRA_PRIORITY:
-				rule.Priority = int(native.Uint32(attrs[j].Value[0:4]))
-			case nl.FRA_IP_PROTO:
-				rule.IPProto = int(native.Uint32(attrs[j].Value[0:4]))
-			case nl.FRA_DPORT_RANGE:
-				rule.Dport = NewRulePortRange(native.Uint16(attrs[j].Value[0:2]), native.Uint16(attrs[j].Value[2:4]))
-			case nl.FRA_SPORT_RANGE:
-				rule.Sport = NewRulePortRange(native.Uint16(attrs[j].Value[0:2]), native.Uint16(attrs[j].Value[2:4]))
-			case nl.FRA_UID_RANGE:
-				rule.UIDRange = NewRuleUIDRange(native.Uint32(attrs[j].Value[0:4]), native.Uint32(attrs[j].Value[4:8]))
-			case nl.FRA_PROTOCOL:
-				rule.Protocol = uint8(attrs[j].Value[0])
-			}
-		}
-
 		if filter != nil {
 			switch {
 			case filterMask&RT_FILTER_SRC != 0 &&
@@ -316,12 +254,210 @@ func (h *Handle) RuleListFiltered(family int, filter *Rule, filterMask uint64) (
 			}
 		}
 
-		res = append(res, *rule)
+		res = append(res, rule)
 	}
 
 	return res, executeErr
 }
 
+func deserializeRule(m []byte) (Rule, error) {
+	msg := nl.DeserializeRtMsg(m)
+	attrs, err := nl.ParseRouteAttr(m[msg.Len():])
+	if err != nil {
+		return Rule{}, err
+	}
+
+	rule := NewRule()
+
+	rule.Invert = msg.Flags&FibRuleInvert > 0
+	rule.Tos = uint(msg.Tos)
+	// The kernel returns the table id in rtmsg.table for values < 256 and as
+	// the RTA_TABLE attribute for larger values. Start with the header value.
+	rule.Table = int(msg.Table)
+
+	for j := range attrs {
+		switch attrs[j].Attr.Type {
+		case unix.RTA_TABLE:
+			rule.Table = int(native.Uint32(attrs[j].Value[0:4]))
+		case nl.FRA_SRC:
+			rule.Src = &net.IPNet{
+				IP:   attrs[j].Value,
+				Mask: net.CIDRMask(int(msg.Src_len), 8*len(attrs[j].Value)),
+			}
+		case nl.FRA_DST:
+			rule.Dst = &net.IPNet{
+				IP:   attrs[j].Value,
+				Mask: net.CIDRMask(int(msg.Dst_len), 8*len(attrs[j].Value)),
+			}
+		case nl.FRA_FWMARK:
+			rule.Mark = native.Uint32(attrs[j].Value[0:4])
+		case nl.FRA_FWMASK:
+			mask := native.Uint32(attrs[j].Value[0:4])
+			rule.Mask = &mask
+		case nl.FRA_TUN_ID:
+			rule.TunID = uint(networkOrder.Uint64(attrs[j].Value[0:8]))
+		case nl.FRA_IIFNAME:
+			rule.IifName = string(attrs[j].Value[:len(attrs[j].Value)-1])
+		case nl.FRA_OIFNAME:
+			rule.OifName = string(attrs[j].Value[:len(attrs[j].Value)-1])
+		case nl.FRA_SUPPRESS_PREFIXLEN:
+			i := native.Uint32(attrs[j].Value[0:4])
+			if i != 0xffffffff {
+				rule.SuppressPrefixlen = int(i)
+			}
+		case nl.FRA_SUPPRESS_IFGROUP:
+			i := native.Uint32(attrs[j].Value[0:4])
+			if i != 0xffffffff {
+				rule.SuppressIfgroup = int(i)
+			}
+		case nl.FRA_FLOW:
+			rule.Flow = int(native.Uint32(attrs[j].Value[0:4]))
+		case nl.FRA_GOTO:
+			rule.Goto = int(native.Uint32(attrs[j].Value[0:4]))
+		case nl.FRA_PRIORITY:
+			rule.Priority = int(native.Uint32(attrs[j].Value[0:4]))
+		case nl.FRA_IP_PROTO:
+			rule.IPProto = int(attrs[j].Value[0])
+		case nl.FRA_DPORT_RANGE:
+			rule.Dport = NewRulePortRange(native.Uint16(attrs[j].Value[0:2]), native.Uint16(attrs[j].Value[2:4]))
+		case nl.FRA_SPORT_RANGE:
+			rule.Sport = NewRulePortRange(native.Uint16(attrs[j].Value[0:2]), native.Uint16(attrs[j].Value[2:4]))
+		case nl.FRA_UID_RANGE:
+			rule.UIDRange = NewRuleUIDRange(native.Uint32(attrs[j].Value[0:4]), native.Uint32(attrs[j].Value[4:8]))
+		case nl.FRA_PROTOCOL:
+			rule.Protocol = uint8(attrs[j].Value[0])
+		}
+	}
+
+	// Some kernels omit FRA_PRIORITY when the rule priority is 0. Since 0 is a
+	// valid priority (e.g. the local-table rule), default to 0 if no attribute
+	// was present.
+	if rule.Priority < 0 {
+		rule.Priority = 0
+	}
+
+	rule.Family = int(msg.Family)
+	rule.Type = msg.Type
+
+	return *rule, nil
+}
+
+// RuleSubscribe takes a chan down which notifications will be sent
+// when rules are added or deleted. The 'done' chan must be closed to stop the
+// subscription and release the underlying socket/goroutine.
+func RuleSubscribe(ch chan<- RuleUpdate, done <-chan struct{}) error {
+	return ruleSubscribeAt(netns.None(), netns.None(), ch, done, nil, false, 0, nil, false)
+}
+
+// RuleSubscribeAt works like RuleSubscribe plus it allows the caller
+// to choose the network namespace in which to subscribe (ns).
+// The 'done' chan must be closed to stop the subscription and release resources.
+func RuleSubscribeAt(ns netns.NsHandle, ch chan<- RuleUpdate, done <-chan struct{}) error {
+	return ruleSubscribeAt(ns, netns.None(), ch, done, nil, false, 0, nil, false)
+}
+
+// RuleSubscribeOptions contains a set of options to use with
+// RuleSubscribeWithOptions.
+type RuleSubscribeOptions struct {
+	Namespace              *netns.NsHandle
+	ErrorCallback          func(error)
+	ListExisting           bool
+	ReceiveBufferSize      int
+	ReceiveBufferForceSize bool
+	ReceiveTimeout         *unix.Timeval
+}
+
+// RuleSubscribeWithOptions works like RuleSubscribe but enables
+// additional options to modify the behavior.
+func RuleSubscribeWithOptions(ch chan<- RuleUpdate, done <-chan struct{}, options RuleSubscribeOptions) error {
+	if options.Namespace == nil {
+		none := netns.None()
+		options.Namespace = &none
+	}
+	return ruleSubscribeAt(*options.Namespace, netns.None(), ch, done, options.ErrorCallback, options.ListExisting,
+		options.ReceiveBufferSize, options.ReceiveTimeout, options.ReceiveBufferForceSize)
+}
+
+func ruleSubscribeAt(newNs, curNs netns.NsHandle, ch chan<- RuleUpdate, done <-chan struct{}, cberr func(error), listExisting bool,
+	rcvbuf int, rcvTimeout *unix.Timeval, rcvBufForce bool) error {
+	s, err := nl.SubscribeAt(newNs, curNs, unix.NETLINK_ROUTE, unix.RTNLGRP_IPV4_RULE, unix.RTNLGRP_IPV6_RULE)
+	if err != nil {
+		return err
+	}
+	if rcvTimeout != nil {
+		if err := s.SetReceiveTimeout(rcvTimeout); err != nil {
+			return err
+		}
+	}
+	if rcvbuf != 0 {
+		err = s.SetReceiveBufferSize(rcvbuf, rcvBufForce)
+		if err != nil {
+			return err
+		}
+	}
+	if done != nil {
+		go func() {
+			<-done
+			s.Close()
+		}()
+	}
+	if listExisting {
+		req := pkgHandle().newNetlinkRequest(unix.RTM_GETRULE, unix.NLM_F_DUMP)
+		infmsg := nl.NewIfInfomsg(unix.AF_UNSPEC)
+		req.AddData(infmsg)
+		if err := s.Send(req); err != nil {
+			s.Close()
+			return err
+		}
+	}
+	go func() {
+		defer close(ch)
+		for {
+			msgs, from, err := s.Receive()
+			if err != nil {
+				if cberr != nil {
+					cberr(fmt.Errorf("Receive failed: %v", err))
+				}
+				return
+			}
+			if from.Pid != nl.PidKernel {
+				if cberr != nil {
+					cberr(fmt.Errorf("Wrong sender portid %d, expected %d", from.Pid, nl.PidKernel))
+				}
+				continue
+			}
+			for _, m := range msgs {
+				if m.Header.Flags&unix.NLM_F_DUMP_INTR != 0 && cberr != nil {
+					cberr(ErrDumpInterrupted)
+				}
+				if m.Header.Type == unix.NLMSG_DONE {
+					continue
+				}
+				if m.Header.Type == unix.NLMSG_ERROR {
+					error := int32(native.Uint32(m.Data[0:4]))
+					if error == 0 {
+						continue
+					}
+					if cberr != nil {
+						cberr(fmt.Errorf("error message: %v", syscall.Errno(-error)))
+					}
+					continue
+				}
+				rule, err := deserializeRule(m.Data)
+				if err != nil {
+					if cberr != nil {
+						cberr(err)
+					}
+					continue
+				}
+				ch <- RuleUpdate{Type: m.Header.Type, Rule: rule}
+			}
+		}
+	}()
+
+	return nil
+}
+
 func (pr *RulePortRange) toRtAttrData() []byte {
 	b := [][]byte{make([]byte, 2), make([]byte, 2)}
 	native.PutUint16(b[0], pr.Start)