adguardhome nftables nftset support, instead of ipset

ipset_linux.go

//go:build linux

package ipset

import (
	"context"
	"fmt"
	"log/slog"
	"net"
	"strings"
	"sync"

	"github.com/AdguardTeam/golibs/container"
	"github.com/AdguardTeam/golibs/errors"
	"github.com/google/nftables"
)

// How to test on a real Linux machine:
//
//  1. Run "sudo ipset create example_set hash:ip family ipv4".
//
//  2. Run "sudo ipset list example_set".  The Members field should be empty.
//
//  3. Add the line "example.com/example_set" to your AdGuardHome.yaml.
//
//  4. Start AdGuardHome.
//
//  5. Make requests to example.com and its subdomains.
//
//  6. Run "sudo ipset list example_set".  The Members field should contain the
//     resolved IP addresses.

// newManager returns a new Linux ipset manager.
func newManager(ctx context.Context, conf *Config) (set Manager, err error) {
	defer func() { err = errors.Annotate(err, "ipset: %w") }()

	c, err := nftables.New(nftables.AsLasting())
        if err != nil {
	   return nil, fmt.Errorf("dialing netfilter: %w", err)
        }

	m := &manager{
		mu: &sync.Mutex{},

		nameToIpset:    make(map[string]*nftables.Set),
		domainToIpsets: make(map[string][]*nftables.Set),

		logger: conf.Logger,

		conn: c,

		addedIPs: container.NewMapSet[ipInIpsetEntry](),
	}

	err = m.parseIpsetConfig(ctx, conf.Lines)
	if err != nil {
		return nil, fmt.Errorf("getting ipsets: %w", err)
	}

	m.logger.DebugContext(ctx, "initialized")

	return m, nil
}

// manager is the Linux Netfilter ipset manager.
type manager struct {
	nameToIpset    map[string]*nftables.Set //"4#inet#fw4#setmefree4":set
	domainToIpsets map[string][]*nftables.Set

	logger *slog.Logger


	// 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 *container.MapSet[ipInIpsetEntry]

	conn *nftables.Conn
}

// ipInIpsetEntry is the type for entries in [manager.addIPs].
type ipInIpsetEntry struct {
	ipsetName string
	// TODO(schzen):  Use netip.Addr.
	ipArr [net.IPv6len]byte
}

// parseIpsetConfigLine parses one ipset configuration line: "www.gogle.com,www.aple.com/4#inet#fw4#s4,6#inet#fw4#s6", same format to dnsmasq
func parseIpsetConfigLine(confStr string) (hosts, ipsetNames []string, err error) {
	confStr = strings.TrimSpace(confStr)
	hostsAndNames := strings.Split(confStr, "/")
	if len(hostsAndNames) != 2 {
		return nil, nil, fmt.Errorf("invalid value %q: expected one slash", confStr)
	}

	hosts = strings.Split(hostsAndNames[0], ",")
	ipsetNames = strings.Split(hostsAndNames[1], ",")

	if len(ipsetNames) == 0 {
		return nil, nil, nil
	}

	for i := range ipsetNames {
		ipsetNames[i] = strings.TrimSpace(ipsetNames[i])
		if len(ipsetNames[i]) == 0 {
			return nil, nil, fmt.Errorf("invalid value %q: empty ipset name", confStr)
		}
	}

	for i := range hosts {
		hosts[i] = strings.ToLower(strings.TrimSpace(hosts[i]))
	}

	return hosts, ipsetNames, nil
}

// parseIpsetConfig parses the ipset configuration and stores ipsets.  It
// returns an error if the configuration can't be used.
func (m *manager) parseIpsetConfig(ctx context.Context, ipsetConf []string) (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.

	for i, confStr := range ipsetConf {
		var hosts, ipsetNames []string
		hosts, ipsetNames, err = parseIpsetConfigLine(confStr)
		if err != nil {
			return fmt.Errorf("config line at idx %d(%s): %w", i, confStr, err)
		}

		var ipsets []*nftables.Set
		for _, n := range ipsetNames {
		       vfts := strings.Split(n, "#") // "4#inet#fw4#s4"
//		       if (len(vfts) != 4) || (vfts[0]!="4"||vfts[0]!="6") || (vfts[1]!="inet") {
		       if (len(vfts) != 4) {
                          return fmt.Errorf("parsing ipsets from config line at idx %d(l=%s,n=%s): wrong format. %s", i,confStr,n,vfts)
		       }

		       p, ok := m.nameToIpset[n]
                       if !ok {
                          tbl := &nftables.Table{Family: nftables.TableFamilyINet, Name: vfts[2],}
			  p,err = m.conn.GetSetByName(tbl,vfts[3])
			  if err != nil {
			     return fmt.Errorf("getting ipsets from config line at idx %d(l=%s,n=%s): %w", i,confStr, n, err)
		          }
			  m.nameToIpset[n] = p
			  m.logger.DebugContext(ctx, "ipset parse", "l",confStr,"hh",hosts,"ss",ipsetNames, "s",n, "p",p.Name)
			}
			if (vfts[0]=="4" && p.KeyType!=nftables.TypeIPAddr) || (vfts[0]=="6" && p.KeyType!=nftables.TypeIP6Addr) {
			    return fmt.Errorf("got ipsets from config line at idx %d(l=%s,n=%s): wrong type ipset", i,confStr, n)
			}
			ipsets = append(ipsets,p) //todo: verify set is ipset and if family is ok
		}

		for _, host := range hosts {
			m.domainToIpsets[host] = append(m.domainToIpsets[host], ipsets...) //todo: dedup
		}
	}

	return nil
}

// lookupHost find the ipsets for the host, taking subdomain wildcards into
// account.
func (m *manager) lookupHost(host string) (sets []*nftables.Set) {
	// Search for matching ipset hosts starting with most specific domain.
	// We could use a trie here but the simple, inefficient solution isn't
	// that expensive: ~10 ns for TLD + SLD vs. ~140 ns for 10 subdomains on
	// an AMD Ryzen 7 PRO 4750U CPU; ~120 ns vs. ~ 1500 ns on a Raspberry
	// Pi's ARMv7 rev 4 CPU.
	for i := 0; ; i++ {
		host = host[i:]
		sets = m.domainToIpsets[host]
		if sets != nil {
			return sets
		}

		i = strings.Index(host, ".")
		if i == -1 {
			break
		}
	}

	// Check the root catch-all one.
	return m.domainToIpsets[""]
}

// addIPs adds the IP addresses for the host to the ipset.  set must be same
// family as set's family.
func (m *manager) addIPs(host string, set *nftables.Set, ips []net.IP) (n int, err error) {
	if len(ips) == 0 {
		return 0, nil
	}

	var entries []nftables.SetElement
	var newAddedEntries []ipInIpsetEntry
	for _, ip := range ips {
                V:="4"
		if set.KeyType==nftables.TypeIP6Addr {
		   V="6"
		}
		F:="inet" //todo: set.Table.TableFamily
		e := ipInIpsetEntry{
			ipsetName: V+"#"+F+"#"+set.Table.Name+"#"+set.Name,
		}
		copy(e.ipArr[:], ip.To16())

		if m.addedIPs.Has(e) {
			continue
		}

		entries = append(entries, nftables.SetElement{Key: []byte(ip.To4())} ) //To16 todo
		newAddedEntries = append(newAddedEntries, e)
	}

	n = len(entries)
	if n == 0 {
		return 0, nil
	}

	err = m.conn.SetAddElements(set, entries)
	if err != nil {
		return 0, fmt.Errorf("adding %q%s to %q %q: %w", host, ips, set.Name, set.KeyType, err)
	}
	err = m.conn.Flush()
	if err != nil {
                return 0, fmt.Errorf("flushing %q%s to %q %q: %w", host, ips, set.Name, set.KeyType, 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.HasTimeout {
			m.addedIPs.Add(e)
		}
	}

	return n, nil
}

// addToSets adds the IP addresses to the corresponding ipset.
func (m *manager) addToSets(
	ctx context.Context,
	host string,
	ip4s []net.IP,
	ip6s []net.IP,
	sets []*nftables.Set,
) (n int, err error) {
	for _, set := range sets {
		var nn int
		switch set.KeyType {
		case nftables.TypeIPAddr:
			nn, err = m.addIPs(host, set, ip4s)
			if err != nil {
				return n, err
			}
		case nftables.TypeIP6Addr:
			nn, err = m.addIPs(host, set, ip6s)
			if err != nil {
				return n, err
			}
		default:
			return n, fmt.Errorf("%q %q unexpected set type", set.Name, set.KeyType)
		}

		m.logger.DebugContext(
			ctx,
			"added ips to set",
			"ips_num", nn,
			"ip4s",ip4s,
			"ip6s",ip6s,
			"set_name", set.Name,
			"set_type", set.KeyType,
		)

		n += nn
	}

	return n, nil
}

// Add implements the [Manager] interface for *manager.
func (m *manager) Add(ctx context.Context, host string, ip4s, ip6s []net.IP) (n int, err error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	sets := m.lookupHost(host)
	if len(sets) == 0 {
		return 0, nil
	}

	m.logger.DebugContext(ctx, "found sets", "set_num", len(sets))

	return m.addToSets(ctx, host, ip4s, ip6s, sets)
}

// Close implements the [Manager] interface for *manager.
func (m *manager) Close() (err error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	var errs []error

	// Close both and collect errors so that the errors from closing one
	// don't interfere with closing the other.
	err = m.conn.CloseLasting()
	if err != nil {
		errs = append(errs, err)
	}

	return errors.Annotate(errors.Join(errs...), "closing ipsets: %w")
}

ipset:
    - www.gogle.com,www.aple.net/4#inet#fw4#setfree4,6#inet#fw4#setwild6

same format to dnsmasq, Version#Family#Table#Set