ipset: add persistent entries

internal/ipset/ipset_linux.go

@@ -3,6 +3,7 @@
 package ipset
 
 import (
+	"bytes"
 	"fmt"
 	"net"
 	"strings"
@@ -13,6 +14,7 @@ import (
 	"github.com/digineo/go-ipset/v2"
 	"github.com/mdlayher/netlink"
 	"github.com/ti-mo/netfilter"
+	"golang.org/x/exp/slices"
 	"golang.org/x/sys/unix"
 )
 
@@ -38,19 +40,58 @@ func newManager(ipsetConf []string) (set Manager, err error) {
 
 // defaultDial is the default netfilter dialing function.
 func defaultDial(pf netfilter.ProtoFamily, conf *netlink.Config) (conn ipsetConn, err error) {
-	conn, err = ipset.Dial(pf, conf)
+	c, err := ipset.Dial(pf, conf)
 	if err != nil {
 		return nil, err
 	}
 
-	return conn, nil
+	return &queryConn{c}, nil
+}
+
+type queryConn struct {
+	*ipset.Conn
+}
+
+func (qc *queryConn) ListAll() (sets []props, err error) {
+	msg, err := netfilter.MarshalNetlink(
+		netfilter.Header{
+			Family:      netfilter.ProtoIPv4,
+			SubsystemID: netfilter.NFSubsysIPSet,
+			MessageType: netfilter.MessageType(ipset.CmdList),
+			Flags:       netlink.Request | netlink.Dump,
+		},
+		[]netfilter.Attribute{{
+			Type: uint16(ipset.AttrProtocol),
+			Data: []byte{ipset.Protocol},
+		}},
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	ms, err := qc.Conn.Conn.Query(msg)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, s := range ms {
+		p := props{}
+		err = p.unmarshalMessage(s)
+		if err != nil {
+			return nil, err
+		}
+
+		sets = append(sets, p)
+	}
+
+	return sets, nil
 }
 
 // ipsetConn is the ipset conn interface.
 type ipsetConn interface {
 	Add(name string, entries ...*ipset.Entry) (err error)
 	Close() (err error)
-	Header(name string) (p *ipset.HeaderPolicy, err error)
+	ListAll() (sets []props, err error)
 }
 
 // dialer creates an ipsetConn.
@@ -60,6 +101,57 @@ type dialer func(pf netfilter.ProtoFamily, conf *netlink.Config) (conn ipsetConn
 type props struct {
 	name   string
 	family netfilter.ProtoFamily
+	temp   bool
+}
+
+func (p *props) unmarshalMessage(msg netlink.Message) (err error) {
+	_, attrs, err := netfilter.UnmarshalNetlink(msg)
+	if err != nil {
+		return err
+	}
+
+	for _, a := range attrs {
+		p.parseAttribute(a)
+	}
+
+	return nil
+}
+
+func (p *props) parseAttribute(a netfilter.Attribute) {
+	switch ipset.AttributeType(a.Type) {
+	case ipset.AttrData:
+		p.parseAttrData(a)
+	case ipset.AttrSetName:
+		p.name = string(bytes.TrimSuffix(a.Data, []byte{0}))
+	case ipset.AttrFamily:
+		p.family = netfilter.ProtoFamily(a.Data[0])
+	default:
+		// Go on.
+	}
+}
+
+func (p *props) parseAttrData(a netfilter.Attribute) {
+	for _, a := range a.Children {
+		switch ipset.AttributeType(a.Type) {
+		case ipset.AttrTimeout:
+			t := a.Uint32()
+			p.temp = t != 0
+		default:
+			// Go on.
+		}
+	}
+}
+
+// unit is a convenient alias for struct{}.
+type unit = struct{}
+
+// ipsInIpset is the type of a set of IP-address-to-ipset mappings.
+type ipsInIpset map[ipInIpsetEntry]unit
+
+// ipInIpsetEntry is the type for entries in an ipsInIpset set.
+type ipInIpsetEntry struct {
+	ipsetName string
+	ipArr     [net.IPv6len]byte
 }
 
 // manager is the Linux Netfilter ipset manager.
@@ -72,6 +164,13 @@ type manager struct {
 	// mu protects all properties below.
 	mu *sync.Mutex
 
+	// TODO(a.garipov): Currently, the ipset list is static, and we don't
+	// read the IPs already in sets, so we can assume that all incoming IPs
+	// are either added to all corresponding ipsets or not.  When that stops
+	// being the case, for example if we add dynamic reconfiguration of
+	// ipsets, this map will need to become a per-ipset-name one.
+	addedIPs ipsInIpset
+
 	ipv4Conn ipsetConn
 	ipv6Conn ipsetConn
 }
@@ -125,50 +224,22 @@ func parseIpsetConfig(confStr string) (hosts, ipsetNames []string, err error) {
 	return hosts, ipsetNames, nil
 }
 
-// ipsetProps returns the properties of an ipset with the given name.
-func (m *manager) ipsetProps(name string) (set props, err error) {
+// ipsets returns currently known ipsets.
+func (m *manager) ipsets(names []string) (sets []props, err error) {
 	// The family doesn't seem to matter when we use a header query, so
 	// query only the IPv4 one.
 	//
 	// TODO(a.garipov): Find out if this is a bug or a feature.
-	var res *ipset.HeaderPolicy
-	res, err = m.ipv4Conn.Header(name)
+	all, err := m.ipv4Conn.ListAll()
 	if err != nil {
-		return set, err
+		return nil, err
 	}
 
-	if res == nil || res.Family == nil {
-		return set, errors.Error("empty response or no family data")
-	}
-
-	family := netfilter.ProtoFamily(res.Family.Value)
-	if family != netfilter.ProtoIPv4 && family != netfilter.ProtoIPv6 {
-		return set, fmt.Errorf("unexpected ipset family %d", family)
-	}
-
-	return props{
-		name:   name,
-		family: family,
-	}, nil
-}
-
-// ipsets returns currently known ipsets.
-func (m *manager) ipsets(names []string) (sets []props, err error) {
-	for _, name := range names {
-		set, ok := m.nameToIpset[name]
-		if ok {
-			sets = append(sets, set)
-
-			continue
+	for _, p := range all {
+		if slices.Contains(names, p.name) {
+			m.nameToIpset[p.name] = p
+			sets = append(sets, p)
 		}
-
-		set, err = m.ipsetProps(name)
-		if err != nil {
-			return nil, fmt.Errorf("querying ipset %q: %w", name, err)
-		}
-
-		m.nameToIpset[name] = set
-		sets = append(sets, set)
 	}
 
 	return sets, nil
@@ -186,6 +257,8 @@ func newManagerWithDialer(ipsetConf []string, dial dialer) (mgr Manager, err err
 		domainToIpsets: make(map[string][]props),
 
 		dial: dial,
+
+		addedIPs: make(ipsInIpset),
 	}
 
 	err = m.dialNetfilter(&netlink.Config{})
@@ -259,8 +332,19 @@ func (m *manager) addIPs(host string, set props, ips []net.IP) (n int, err error
 	}
 
 	var entries []*ipset.Entry
+	var newAddedEntries []ipInIpsetEntry
 	for _, ip := range ips {
+		e := ipInIpsetEntry{
+			ipsetName: set.name,
+		}
+		copy(e.ipArr[:], ip.To16())
+
+		if _, added := m.addedIPs[e]; added {
+			continue
+		}
+
 		entries = append(entries, ipset.NewEntry(ipset.EntryIP(ip)))
+		newAddedEntries = append(newAddedEntries, e)
 	}
 
 	n = len(entries)
@@ -283,6 +367,15 @@ func (m *manager) addIPs(host string, set props, ips []net.IP) (n int, err error
 		return 0, fmt.Errorf("adding %q%s to ipset %q: %w", host, ips, set.name, err)
 	}
 
+	// Only add these to the cache once we're sure that all of them were
+	// actually sent to the ipset.
+	for _, e := range newAddedEntries {
+		s := m.nameToIpset[e.ipsetName]
+		if !s.temp {
+			m.addedIPs[e] = unit{}
+		}
+	}
+
 	return n, nil
 }
 

internal/ipset/ipset_linux_internal_test.go

@@ -21,6 +21,7 @@ type fakeConn struct {
 	ipv4Entries *[]*ipset.Entry
 	ipv6Header  *ipset.HeaderPolicy
 	ipv6Entries *[]*ipset.Entry
+	sets        []props
 }
 
 // Add implements the [ipsetConn] interface for *fakeConn.
@@ -43,15 +44,9 @@ func (c *fakeConn) Close() (err error) {
 	return nil
 }
 
-// Header implements the [ipsetConn] interface for *fakeConn.
-func (c *fakeConn) Header(name string) (p *ipset.HeaderPolicy, err error) {
-	if strings.Contains(name, "ipv4") {
-		return c.ipv4Header, nil
-	} else if strings.Contains(name, "ipv6") {
-		return c.ipv6Header, nil
-	}
-
-	return nil, errors.Error("test: ipset not found")
+// ListAll implements the [ipsetConn] interface for *fakeConn.
+func (c *fakeConn) ListAll() (sets []props, err error) {
+	return c.sets, nil
 }
 
 func TestManager_Add(t *testing.T) {
@@ -76,6 +71,13 @@ func TestManager_Add(t *testing.T) {
 				Family: ipset.NewUInt8Box(uint8(netfilter.ProtoIPv6)),
 			},
 			ipv6Entries: &ipv6Entries,
+			sets: []props{{
+				name:   "ipv4set",
+				family: netfilter.ProtoIPv4,
+			}, {
+				name:   "ipv6set",
+				family: netfilter.ProtoIPv6,
+			}},
 		}, nil
 	}