DebTik/internal/firewall/ipset_atomic.go

package firewall

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

var setApplyMu sync.Mutex
var setDigests = map[string][32]byte{}

func runIPSet(args ...string) error {
	out, err := exec.Command("ipset", args...).CombinedOutput()
	if err != nil {
		return fmt.Errorf("ipset %v: %s: %w", args, out, err)
	}
	return nil
}
func replaceSet(name, family string, addresses []string) error {
	setApplyMu.Lock()
	defer setApplyMu.Unlock()
	if len(name) > 31 {
		return fmt.Errorf("address list name too long")
	}
	unique := map[string]bool{}
	var clean []string
	for _, address := range addresses {
		ip := net.ParseIP(address)
		if ip != nil {
			address = ip.String()
		}
		if ip == nil {
			var network *net.IPNet
			var err error
			ip, network, err = net.ParseCIDR(address)
			if err == nil {
				address = network.String()
			}
		}
		if ip == nil || (ip.To4() != nil) != (family == "inet") {
			return fmt.Errorf("invalid %s address %q", family, address)
		}
		if !unique[address] {
			unique[address] = true
			clean = append(clean, address)
		}
	}
	digest := sha256.Sum256([]byte(family + "\n" + strings.Join(clean, "\n")))
	header, inspectErr := exec.Command("ipset", "list", name, "-terse").CombinedOutput()
	exists := inspectErr == nil
	if inspectErr != nil && !strings.Contains(string(header), "does not exist") {
		return fmt.Errorf("cannot inspect ipset %s: %s: %w", name, header, inspectErr)
	}
	if exists && (!strings.Contains(string(header), "Type: hash:net\n") || !strings.Contains(string(header), "family "+family+" ")) {
		return fmt.Errorf("set %s has an incompatible type/family; existing data was preserved", name)
	}
	if exists && setDigests[name] == digest && strings.Contains(string(header), "Number of entries: "+strconv.Itoa(len(clean))+"\n") {
		return nil
	}

	h := sha256.Sum256([]byte(name))
	temp := fmt.Sprintf("dt_tmp_%x", h[:8])
	// A previous version may have left a temporary set with different limits.
	_ = exec.Command("ipset", "destroy", temp).Run()
	var restore strings.Builder
	fmt.Fprintf(&restore, "create %s hash:net family %s maxelem 1048576\n", temp, family)
	for _, address := range clean {
		fmt.Fprintf(&restore, "add %s %s -exist\n", temp, address)
	}
	if !exists {
		fmt.Fprintf(&restore, "create %s hash:net family %s maxelem 1048576\n", name, family)
	}
	fmt.Fprintf(&restore, "swap %s %s\ndestroy %s\n", temp, name, temp)
	cmd := exec.Command("ipset", "restore")
	cmd.Stdin = strings.NewReader(restore.String())
	if out, err := cmd.CombinedOutput(); err != nil {
		_ = exec.Command("ipset", "destroy", temp).Run()
		return fmt.Errorf("ipset bulk restore %s: %s: %w", name, out, err)
	}
	setDigests[name] = digest
	return nil
}