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")))
	if setDigests[name] == digest {
		out, err := exec.Command("ipset", "list", name, "-terse").Output()
		if err == nil && strings.Contains(string(out), "Number of entries: "+strconv.Itoa(len(clean))+"\n") {
			return nil
		}
	}
	h := sha256.Sum256([]byte(name))
	temp := fmt.Sprintf("dt_tmp_%x", h[:8])
	var restore strings.Builder
	fmt.Fprintf(&restore, "create %s hash:net family %s maxelem 1048576 -exist\nflush %s\n", temp, family, temp)
	for _, address := range clean {
		fmt.Fprintf(&restore, "add %s %s -exist\n", temp, address)
	}
	fmt.Fprintf(&restore, "create %s hash:net family %s maxelem 1048576 -exist\nswap %s %s\ndestroy %s\n", name, family, 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
}