banger/internal/daemon/nat.go

package daemon

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"banger/internal/model"
	"banger/internal/system"
)

type natRule struct {
	table string
	chain string
	args  []string
}

func (d *Daemon) ensureNAT(ctx context.Context, vm model.VMRecord, enable bool) error {
	if err := system.RequireCommands(ctx, "iptables", "sysctl"); err != nil {
		return err
	}
	uplink, err := d.defaultUplink(ctx)
	if err != nil {
		return err
	}
	rules, err := natRulesForVM(vm, uplink)
	if err != nil {
		return err
	}
	if enable {
		if _, err := d.runner.RunSudo(ctx, "sysctl", "-w", "net.ipv4.ip_forward=1"); err != nil {
			return err
		}
		for _, rule := range rules {
			if err := d.addNATRule(ctx, rule); err != nil {
				return err
			}
		}
		return nil
	}
	for _, rule := range rules {
		if err := d.removeNATRule(ctx, rule); err != nil {
			return err
		}
	}
	return nil
}

func (d *Daemon) defaultUplink(ctx context.Context) (string, error) {
	out, err := d.runner.Run(ctx, "ip", "route", "show", "default")
	if err != nil {
		return "", err
	}
	return parseDefaultUplink(string(out))
}

func parseDefaultUplink(output string) (string, error) {
	for _, line := range strings.Split(output, "\n") {
		fields := strings.Fields(line)
		if len(fields) == 0 || fields[0] != "default" {
			continue
		}
		for i := 0; i < len(fields)-1; i++ {
			if fields[i] == "dev" && fields[i+1] != "" {
				return fields[i+1], nil
			}
		}
	}
	return "", errors.New("failed to detect uplink interface")
}

func natRulesForVM(vm model.VMRecord, uplink string) ([]natRule, error) {
	guestIP := strings.TrimSpace(vm.Runtime.GuestIP)
	if guestIP == "" {
		return nil, errors.New("nat requires a guest IP")
	}
	tap := strings.TrimSpace(vm.Runtime.TapDevice)
	if tap == "" {
		return nil, errors.New("nat requires a tap device")
	}
	uplink = strings.TrimSpace(uplink)
	if uplink == "" {
		return nil, errors.New("nat requires an uplink interface")
	}
	guestCIDR := guestIP + "/32"
	return []natRule{
		{
			table: "nat",
			chain: "POSTROUTING",
			args:  []string{"-s", guestCIDR, "-o", uplink, "-j", "MASQUERADE"},
		},
		{
			chain: "FORWARD",
			args:  []string{"-i", tap, "-o", uplink, "-j", "ACCEPT"},
		},
		{
			chain: "FORWARD",
			args:  []string{"-i", uplink, "-o", tap, "-m", "state", "--state", "RELATED,ESTABLISHED", "-j", "ACCEPT"},
		},
	}, nil
}

func natRuleArgs(action string, rule natRule) []string {
	args := make([]string, 0, len(rule.args)+4)
	if rule.table != "" {
		args = append(args, "-t", rule.table)
	}
	args = append(args, action, rule.chain)
	args = append(args, rule.args...)
	return args
}

func natAddPlan(rules []natRule) [][]string {
	plan := make([][]string, 0, len(rules)+1)
	plan = append(plan, []string{"sysctl", "-w", "net.ipv4.ip_forward=1"})
	for _, rule := range rules {
		plan = append(plan, natRuleArgs("-A", rule))
	}
	return plan
}

func natRemovePlan(rules []natRule) [][]string {
	plan := make([][]string, 0, len(rules))
	for _, rule := range rules {
		plan = append(plan, natRuleArgs("-D", rule))
	}
	return plan
}

func (d *Daemon) addNATRule(ctx context.Context, rule natRule) error {
	if _, err := d.runner.RunSudo(ctx, append([]string{"iptables"}, natRuleArgs("-C", rule)...)...); err == nil {
		return nil
	}
	_, err := d.runner.RunSudo(ctx, append([]string{"iptables"}, natRuleArgs("-A", rule)...)...)
	return err
}

func (d *Daemon) removeNATRule(ctx context.Context, rule natRule) error {
	if _, err := d.runner.RunSudo(ctx, append([]string{"iptables"}, natRuleArgs("-C", rule)...)...); err != nil {
		return nil
	}
	_, err := d.runner.RunSudo(ctx, append([]string{"iptables"}, natRuleArgs("-D", rule)...)...)
	return err
}

func natRuleKey(rule natRule) string {
	return fmt.Sprintf("%s:%s:%s", rule.table, rule.chain, strings.Join(rule.args, " "))
}