DebTik/internal/firewall/ipset.go

package firewall

import (
	"crypto/sha256"
	"fmt"
	"net"
	"os/exec"
	"strings"
)

func listToken(name string) string { h := sha256.Sum256([]byte(name)); return fmt.Sprintf("%x", h[:8]) }
func GetIPSetName(name string) string {
	name = strings.TrimPrefix(name, "!")
	if strings.HasPrefix(name, "geo:") {
		return GetGeoIPSetName(name)
	}
	if strings.HasPrefix(name, "geo_") {
		return GetGeoIPSetName(name)
	}
	return "dt_al_" + listToken(name)
}
func GetDynamicIPSetName(name string) string { return "dt_ad_" + listToken(name) }
func staticIPSetName(name string) string     { return "dt_as_" + listToken(name) }
func GetGeoIPSetName(code string) string {
	code = strings.TrimPrefix(code, "!")
	code = strings.TrimPrefix(code, "geo:")
	code = strings.TrimPrefix(code, "geo_")
	return "mk_geo_" + strings.ToUpper(code)
}
func SyncAddressLists(entries []AddressListEntry, referenced ...string) error {
	lists := map[string][]string{}
	// Retain empty lists represented by disabled entries, so disabling the last
	// static entry clears its addresses without deleting a dynamic collection.
	for _, e := range entries {
		if err := validateListName(e.Name); err != nil {
			return err
		}
		if _, ok := lists[e.Name]; !ok {
			lists[e.Name] = nil
		}
		if !e.Disabled {
			if e.Address == "" {
				return fmt.Errorf("address is required")
			}
			lists[e.Name] = append(lists[e.Name], e.Address)
		}
	}
	for _, name := range referenced {
		name = strings.TrimPrefix(name, "!")
		if name != "" && !strings.HasPrefix(name, "geo:") && !strings.HasPrefix(name, "geo_") {
			if err := validateListName(name); err != nil {
				return err
			}
			if _, ok := lists[name]; !ok {
				lists[name] = nil
			}
		}
	}
	desired := map[string]bool{}
	for name, addrs := range lists {
		families := map[string][]string{"inet": nil, "inet6": nil}
		for _, address := range addrs {
			ip := net.ParseIP(address)
			if ip == nil {
				ip, _, _ = net.ParseCIDR(address)
			}
			if ip == nil {
				return fmt.Errorf("invalid address: %s", address)
			}
			family := "inet6"
			if ip.To4() != nil {
				family = "inet"
			}
			families[family] = append(families[family], address)
		}
		for family, addresses := range families {
			suffix := ""
			if family == "inet6" {
				suffix = "_v6"
			}
			static, dynamic, union := staticIPSetName(name)+suffix, GetDynamicIPSetName(name)+suffix, GetIPSetName(name)+suffix
			desired[static] = true
			desired[dynamic] = true
			desired[union] = true
			if err := replaceSet(static, family, addresses); err != nil {
				return err
			}
			if err := runIPSet("create", dynamic, "hash:net", "family", family, "timeout", "0", "maxelem", "262144", "-exist"); err != nil {
				return err
			}
			if err := runIPSet("create", union, "list:set", "size", "2", "-exist"); err != nil {
				return err
			}
			for _, child := range []string{static, dynamic} {
				if err := runIPSet("add", union, child, "-exist"); err != nil {
					return err
				}
			}
		}
	}
	// Only collect our own sets and only if the kernel reports no references.
	out, err := exec.Command("ipset", "list", "-name").Output()
	if err != nil {
		return err
	}
	for _, prefix := range []string{"dt_al_", "dt_as_", "dt_ad_"} {
		for _, name := range strings.Fields(string(out)) {
			if strings.HasPrefix(name, prefix) && !desired[name] {
				_ = exec.Command("ipset", "destroy", name).Run()
			}
		}
	}
	return nil
}

func collectUnusedAddressLists(entries []AddressListEntry, refs []string) {
	desired := map[string]bool{}
	for _, e := range entries {
		refs = append(refs, e.Name)
	}
	for _, name := range refs {
		name = strings.TrimPrefix(name, "!")
		if name == "" || strings.HasPrefix(name, "geo:") || strings.HasPrefix(name, "geo_") {
			continue
		}
		for _, suffix := range []string{"", "_v6"} {
			desired[GetIPSetName(name)+suffix] = true
			desired[GetDynamicIPSetName(name)+suffix] = true
			desired[staticIPSetName(name)+suffix] = true
		}
	}
	out, err := exec.Command("ipset", "list", "-name").Output()
	if err != nil {
		return
	}
	for _, prefix := range []string{"dt_al_", "dt_as_", "dt_ad_"} {
		for _, name := range strings.Fields(string(out)) {
			if strings.HasPrefix(name, prefix) && !desired[name] {
				if exec.Command("ipset", "destroy", name).Run() == nil {
					setApplyMu.Lock()
					delete(setDigests, name)
					setApplyMu.Unlock()
				}
			}
		}
	}
}