DebTik/internal/sdwan/engine.go

package sdwan

import (
	"fmt"
	"log"
	"math"
	"net"
	"os/exec"
	"strconv"
	"strings"
	"sync"

	"github.com/vishvananda/netlink"
)

// KernelEngine manages routing tables, ip rules, virtual devices, and netfilter mangle rules for SD-WAN.
type KernelEngine struct {
	mu             sync.Mutex
	virtualDevices map[string]bool
}

func NewKernelEngine() *KernelEngine {
	return &KernelEngine{
		virtualDevices: make(map[string]bool),
	}
}

// Sync applies the complete routing, virtual devices, and iptables rules for all active SD-WAN groups.
func (e *KernelEngine) Sync(groups []*SDWANGroup) error {
	e.mu.Lock()
	defer e.mu.Unlock()

	// 1. Clean existing SD-WAN mangle chain & rules
	e.cleanupMangle()

	// 2. Clear existing SD-WAN ip rules
	e.cleanupIPRules()

	hasActiveGroups := false
	for _, g := range groups {
		if !g.Disabled && len(g.Members) > 0 {
			hasActiveGroups = true
			break
		}
	}

	// 3. Synchronize virtual dummy interfaces for active groups (e.g. "sdwan0")
	activeDevices := make(map[string]bool)
	for _, g := range groups {
		if !g.Disabled && len(g.Members) > 0 && g.Name != "" {
			activeDevices[g.Name] = true
			e.ensureVirtualDevice(g.Name)
		}
	}
	for dev := range e.virtualDevices {
		if !activeDevices[dev] {
			e.removeVirtualDevice(dev)
		}
	}
	e.virtualDevices = activeDevices

	if !hasActiveGroups {
		return nil
	}

	// 4. Create fresh DEBTIK_SDWAN mangle chain
	_ = exec.Command("iptables", "-t", "mangle", "-N", "DEBTIK_SDWAN").Run()
	_ = exec.Command("iptables", "-t", "mangle", "-F", "DEBTIK_SDWAN").Run()

	// Ensure jumps from PREROUTING and OUTPUT
	e.ensureMangleJump("PREROUTING")
	e.ensureMangleJump("OUTPUT")

	// Bypass private RFC1918, loopback, link-local, and multicast
	privateNets := []string{"10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16", "127.0.0.0/8", "169.254.0.0/16", "224.0.0.0/4"}
	for _, net := range privateNets {
		_ = exec.Command("iptables", "-t", "mangle", "-A", "DEBTIK_SDWAN", "-d", net, "-j", "RETURN").Run()
	}

	// Restore connection mark: if a flow has already been classified, keep its mark (Session Stickiness!)
	_ = exec.Command("iptables", "-t", "mangle", "-A", "DEBTIK_SDWAN", "-j", "CONNMARK", "--restore-mark").Run()
	_ = exec.Command("iptables", "-t", "mangle", "-A", "DEBTIK_SDWAN", "-m", "mark", "!", "--mark", "0", "-j", "RETURN").Run()

	markBase := 0x5d00
	tableBase := 1100
	memberCounter := 1

	for _, g := range groups {
		if g.Disabled || len(g.Members) == 0 {
			g.Running = false
			continue
		}

		// Filter active candidates: prefer healthy/degraded over down
		candidates := make([]*SDWANMember, 0)
		for _, m := range g.Members {
			if !m.Disabled {
				candidates = append(candidates, m)
			}
		}

		if len(candidates) == 0 {
			g.Running = false
			continue
		}

		// Determine online candidates
		onlineCandidates := make([]*SDWANMember, 0)
		for _, m := range candidates {
			if m.Status != StatusDown {
				onlineCandidates = append(onlineCandidates, m)
			}
		}

		// If all candidates are down, fallback to using all enabled candidates so traffic isn't completely blackholed
		activePool := onlineCandidates
		if len(activePool) == 0 {
			activePool = candidates
		}

		// Assign table ID and marks
		for _, m := range g.Members {
			m.TableID = tableBase + memberCounter
			m.MarkInt = uint32(markBase + memberCounter)
			m.MarkHex = fmt.Sprintf("0x%x", m.MarkInt)
			memberCounter++

			// Apply kernel routing table
			e.applyMemberRoute(m)

			// Apply policy routing rule (ip rule)
			e.applyMemberRule(m)

			// Auto-NAT if configured
			if g.AutoNAT && !m.Disabled {
				e.ensureNAT(m.Interface)
			}
		}

		// Apply Mangle classification rules based on group mode
		switch g.Mode {
		case ModeWeighted:
			e.applyWeightedRules(activePool)
		case ModePCC:
			e.applyPCCRules(activePool)
		case ModeLowestLatency:
			e.applyBestQualityRules(activePool, true)
		case ModeLowestLoss:
			e.applyBestQualityRules(activePool, false)
		case ModeFailover:
			e.applyFailoverRules(activePool)
		default:
			e.applyWeightedRules(activePool)
		}

		g.Running = true
	}

	return nil
}

func (e *KernelEngine) applyMemberRoute(m *SDWANMember) {
	if m.Interface == "" {
		return
	}
	tableStr := strconv.Itoa(m.TableID)

	// Flush table first
	_ = exec.Command("ip", "route", "flush", "table", tableStr).Run()

	args := []string{"route", "replace", "default"}
	if m.Gateway != "" {
		args = append(args, "via", m.Gateway)
	}
	args = append(args, "dev", m.Interface, "table", tableStr)

	if out, err := exec.Command("ip", args...).CombinedOutput(); err != nil {
		log.Printf("[SD-WAN] Warning: failed adding route for %s (table %s): %s (%v)", m.Interface, tableStr, strings.TrimSpace(string(out)), err)
	}
}

func (e *KernelEngine) applyMemberRule(m *SDWANMember) {
	tableStr := strconv.Itoa(m.TableID)
	prefStr := strconv.Itoa(2000 + m.TableID%1000)

	// Delete matching rule if already present
	_ = exec.Command("ip", "rule", "del", "fwmark", m.MarkHex, "table", tableStr).Run()

	// Add ip rule
	_ = exec.Command("ip", "rule", "add", "fwmark", m.MarkHex, "table", tableStr, "pref", prefStr).Run()
}

func (e *KernelEngine) applyWeightedRules(pool []*SDWANMember) {
	if len(pool) == 0 {
		return
	}
	if len(pool) == 1 {
		e.addMarkRule(pool[0].MarkHex, 1.0)
		return
	}

	totalWeight := 0
	for _, m := range pool {
		w := m.Weight
		if w <= 0 {
			w = 1
		}
		totalWeight += w
	}

	remainingWeight := totalWeight
	for i, m := range pool {
		if i == len(pool)-1 {
			// Last member gets 100% of the remainder
			e.addMarkRule(m.MarkHex, 1.0)
			break
		}

		w := m.Weight
		if w <= 0 {
			w = 1
		}
		prob := float64(w) / float64(remainingWeight)
		// Clamp probability between 0.0001 and 0.9999
		prob = math.Max(0.0001, math.Min(0.9999, prob))

		e.addMarkRule(m.MarkHex, prob)
		remainingWeight -= w
	}
}

func (e *KernelEngine) applyPCCRules(pool []*SDWANMember) {
	n := len(pool)
	if n == 0 {
		return
	}
	if n == 1 {
		e.addMarkRule(pool[0].MarkHex, 1.0)
		return
	}

	// Using nth packet mode for new connection distribution
	for i, m := range pool {
		args := []string{
			"-t", "mangle", "-A", "DEBTIK_SDWAN",
			"-m", "conntrack", "--ctstate", "NEW",
			"-m", "statistic", "--mode", "nth",
			"--every", strconv.Itoa(n - i),
			"--packet", "0",
			"-j", "MARK", "--set-mark", m.MarkHex,
		}
		_ = exec.Command("iptables", args...).Run()

		// Save mark into conntrack
		_ = exec.Command("iptables", "-t", "mangle", "-A", "DEBTIK_SDWAN",
			"-m", "conntrack", "--ctstate", "NEW",
			"-m", "mark", "--mark", m.MarkHex,
			"-j", "CONNMARK", "--save-mark").Run()

		_ = exec.Command("iptables", "-t", "mangle", "-A", "DEBTIK_SDWAN",
			"-m", "conntrack", "--ctstate", "NEW",
			"-m", "mark", "--mark", m.MarkHex,
			"-j", "RETURN").Run()
	}
}

func (e *KernelEngine) applyBestQualityRules(pool []*SDWANMember, latencyMode bool) {
	if len(pool) == 0 {
		return
	}
	best := pool[0]
	for _, m := range pool[1:] {
		if latencyMode {
			if m.LiveLatency > 0 && (best.LiveLatency <= 0 || m.LiveLatency < best.LiveLatency) {
				best = m
			}
		} else {
			if m.PacketLoss < best.PacketLoss {
				best = m
			}
		}
	}
	e.addMarkRule(best.MarkHex, 1.0)
}

func (e *KernelEngine) applyFailoverRules(pool []*SDWANMember) {
	if len(pool) == 0 {
		return
	}
	// Sort by priority (1 is highest)
	primary := pool[0]
	for _, m := range pool[1:] {
		prio := m.Priority
		if prio <= 0 {
			prio = 100
		}
		bestPrio := primary.Priority
		if bestPrio <= 0 {
			bestPrio = 100
		}
		if prio < bestPrio {
			primary = m
		}
	}
	e.addMarkRule(primary.MarkHex, 1.0)
}

func (e *KernelEngine) addMarkRule(markHex string, prob float64) {
	if prob >= 1.0 {
		_ = exec.Command("iptables", "-t", "mangle", "-A", "DEBTIK_SDWAN",
			"-m", "conntrack", "--ctstate", "NEW",
			"-j", "MARK", "--set-mark", markHex).Run()
	} else {
		probStr := fmt.Sprintf("%.4f", prob)
		_ = exec.Command("iptables", "-t", "mangle", "-A", "DEBTIK_SDWAN",
			"-m", "conntrack", "--ctstate", "NEW",
			"-m", "statistic", "--mode", "random", "--probability", probStr,
			"-j", "MARK", "--set-mark", markHex).Run()
	}

	_ = exec.Command("iptables", "-t", "mangle", "-A", "DEBTIK_SDWAN",
		"-m", "conntrack", "--ctstate", "NEW",
		"-m", "mark", "--mark", markHex,
		"-j", "CONNMARK", "--save-mark").Run()

	_ = exec.Command("iptables", "-t", "mangle", "-A", "DEBTIK_SDWAN",
		"-m", "conntrack", "--ctstate", "NEW",
		"-m", "mark", "--mark", markHex,
		"-j", "RETURN").Run()
}

func (e *KernelEngine) ensureMangleJump(chain string) {
	// Check if jump exists
	check := exec.Command("iptables", "-t", "mangle", "-C", chain, "-j", "DEBTIK_SDWAN")
	if check.Run() != nil {
		_ = exec.Command("iptables", "-t", "mangle", "-I", chain, "1", "-j", "DEBTIK_SDWAN").Run()
	}
}

func (e *KernelEngine) ensureNAT(iface string) {
	if iface == "" {
		return
	}
	check := exec.Command("iptables", "-t", "nat", "-C", "POSTROUTING", "-o", iface, "-m", "comment", "--comment", "sdwan-auto-nat", "-j", "MASQUERADE")
	if check.Run() != nil {
		_ = exec.Command("iptables", "-t", "nat", "-A", "POSTROUTING", "-o", iface, "-m", "comment", "--comment", "sdwan-auto-nat", "-j", "MASQUERADE").Run()
	}
}

func (e *KernelEngine) cleanupMangle() {
	_ = exec.Command("iptables", "-t", "mangle", "-D", "PREROUTING", "-j", "DEBTIK_SDWAN").Run()
	_ = exec.Command("iptables", "-t", "mangle", "-D", "OUTPUT", "-j", "DEBTIK_SDWAN").Run()
	_ = exec.Command("iptables", "-t", "mangle", "-F", "DEBTIK_SDWAN").Run()
	_ = exec.Command("iptables", "-t", "mangle", "-X", "DEBTIK_SDWAN").Run()
}

func (e *KernelEngine) cleanupIPRules() {
	out, err := exec.Command("ip", "rule", "show").CombinedOutput()
	if err != nil {
		return
	}
	lines := strings.Split(string(out), "\n")
	for _, l := range lines {
		if strings.Contains(l, "0x5d") {
			// Extract fwmark and table
			fields := strings.Fields(l)
			for i, f := range fields {
				if f == "fwmark" && i+1 < len(fields) {
					mark := fields[i+1]
					_ = exec.Command("ip", "rule", "del", "fwmark", mark).Run()
				}
			}
		}
	}
}

func (e *KernelEngine) ensureVirtualDevice(name string) {
	if name == "" {
		return
	}
	link, err := netlink.LinkByName(name)
	if err == nil {
		if link.Attrs().Flags&net.FlagUp == 0 {
			_ = netlink.LinkSetUp(link)
		}
		return
	}

	dummy := &netlink.Dummy{
		LinkAttrs: netlink.LinkAttrs{
			Name: name,
			MTU:  1500,
		},
	}
	if err := netlink.LinkAdd(dummy); err != nil {
		_ = exec.Command("ip", "link", "add", name, "type", "dummy").Run()
		_ = exec.Command("ip", "link", "set", name, "up", "mtu", "1500").Run()
		return
	}
	_ = netlink.LinkSetUp(dummy)
}

func (e *KernelEngine) removeVirtualDevice(name string) {
	if name == "" {
		return
	}
	link, err := netlink.LinkByName(name)
	if err == nil {
		_ = netlink.LinkDel(link)
	} else {
		_ = exec.Command("ip", "link", "del", name).Run()
	}
}