prx/internal/iprule/iprule.go

// SPDX-License-Identifier: MIT

// Package iprule manages the policy routing rules and routes that prx relies
// on to keep inbound reply traffic on the physical interface.
package iprule

import (
	"context"
	"errors"
	"fmt"
	"regexp"
	"strconv"
	"strings"

	"prx/internal/execx"
)

// Spec describes the mark-based policy rule that steers replies back out the WAN.
type Spec struct {
	Family   int // 4 or 6
	Mark     uint32
	Table    int
	Priority int
	Gateway  string
	Dev      string
}

// BypassSpec describes a destination-based rule that keeps traffic direct.
type BypassSpec struct {
	Family   int
	CIDR     string
	Table    int
	Priority int
	Gateway  string
	Dev      string
}

func ipArgs(family int, rest ...string) []string {
	if family == 6 {
		return append([]string{"-6"}, rest...)
	}
	return rest
}

func markString(m uint32) string { return fmt.Sprintf("0x%x", m) }

type ruleEntry struct {
	Priority int
	HasMark  bool
	Mark     uint32
	HasTo    bool
	To       string
	Table    int
}

var ruleLineRE = regexp.MustCompile(`^\s*(\d+):\s+(.*)$`)

func parseRules(out string) []ruleEntry {
	var entries []ruleEntry
	for _, line := range strings.Split(out, "\n") {
		m := ruleLineRE.FindStringSubmatch(line)
		if m == nil {
			continue
		}
		prio, err := strconv.Atoi(m[1])
		if err != nil {
			continue
		}
		e := ruleEntry{Priority: prio}
		fields := strings.Fields(m[2])
		for i := 0; i < len(fields); i++ {
			switch fields[i] {
			case "fwmark":
				if i+1 < len(fields) {
					e.HasMark = true
					e.Mark = parseMark(fields[i+1])
					i++
				}
			case "lookup":
				if i+1 < len(fields) {
					e.Table, _ = strconv.Atoi(fields[i+1])
					i++
				}
			case "to":
				if i+1 < len(fields) {
					e.HasTo = true
					e.To = fields[i+1]
					i++
				}
			}
		}
		entries = append(entries, e)
	}
	return entries
}

func parseMark(s string) uint32 {
	if idx := strings.IndexByte(s, '/'); idx >= 0 {
		s = s[:idx]
	}
	s = strings.TrimPrefix(strings.TrimPrefix(s, "0x"), "0X")
	v, _ := strconv.ParseUint(s, 16, 32)
	return uint32(v)
}

func listRules(ctx context.Context, r execx.Runner, family int) ([]ruleEntry, error) {
	res, err := r.Run(ctx, nil, "ip", ipArgs(family, "rule", "show")...)
	if err != nil {
		return nil, fmt.Errorf("ip rule show: %w", err)
	}
	return parseRules(string(res.Stdout)), nil
}

// RuleExists reports whether a mark-based rule with the given attributes exists.
func RuleExists(ctx context.Context, r execx.Runner, s Spec) (bool, error) {
	entries, err := listRules(ctx, r, s.Family)
	if err != nil {
		return false, err
	}
	for _, e := range entries {
		if e.Priority == s.Priority && e.HasMark && e.Mark == s.Mark && e.Table == s.Table {
			return true, nil
		}
	}
	return false, nil
}

// EnsureRule adds the mark-based rule if it is missing.
func EnsureRule(ctx context.Context, r execx.Runner, s Spec) (bool, error) {
	exists, err := RuleExists(ctx, r, s)
	if err != nil {
		return false, err
	}
	if exists {
		return false, nil
	}
	_, err = r.Run(ctx, nil, "ip", ipArgs(s.Family,
		"rule", "add", "fwmark", markString(s.Mark),
		"lookup", strconv.Itoa(s.Table),
		"priority", strconv.Itoa(s.Priority))...)
	if err != nil {
		return false, fmt.Errorf("add fwmark rule: %w", err)
	}
	return true, nil
}

// BypassRuleExists reports whether a destination-based rule exists.
func BypassRuleExists(ctx context.Context, r execx.Runner, s BypassSpec) (bool, error) {
	entries, err := listRules(ctx, r, s.Family)
	if err != nil {
		return false, err
	}
	for _, e := range entries {
		if e.Priority == s.Priority && e.HasTo && e.To == s.CIDR && e.Table == s.Table {
			return true, nil
		}
	}
	return false, nil
}

// EnsureBypassRule adds a destination-based bypass rule if missing.
func EnsureBypassRule(ctx context.Context, r execx.Runner, s BypassSpec) (bool, error) {
	exists, err := BypassRuleExists(ctx, r, s)
	if err != nil {
		return false, err
	}
	if exists {
		return false, nil
	}
	_, err = r.Run(ctx, nil, "ip", ipArgs(s.Family,
		"rule", "add", "to", s.CIDR,
		"lookup", strconv.Itoa(s.Table),
		"priority", strconv.Itoa(s.Priority))...)
	if err != nil {
		return false, fmt.Errorf("add bypass rule: %w", err)
	}
	return true, nil
}

// RouteExists reports whether the table already holds a default route via the
// given gateway and device.
func RouteExists(ctx context.Context, r execx.Runner, family, table int, gateway, dev string) (bool, error) {
	res, err := r.Run(ctx, nil, "ip", ipArgs(family, "route", "show", "table", strconv.Itoa(table))...)
	if err != nil {
		if isMissing(err) {
			return false, nil
		}
		return false, err
	}
	for _, line := range strings.Split(string(res.Stdout), "\n") {
		fields := strings.Fields(line)
		if len(fields) == 0 || fields[0] != "default" {
			continue
		}
		if via, d := extractViaDev(fields[1:]); via == gateway && (dev == "" || d == dev) {
			return true, nil
		}
	}
	return false, nil
}

// EnsureRoute installs a default route in the policy table (idempotent).
func EnsureRoute(ctx context.Context, r execx.Runner, family, table int, gateway, dev string) (bool, error) {
	exists, err := RouteExists(ctx, r, family, table, gateway, dev)
	if err != nil {
		return false, err
	}
	if exists {
		return false, nil
	}
	_, err = r.Run(ctx, nil, "ip", ipArgs(family,
		"route", "replace", "default", "via", gateway, "dev", dev, "table", strconv.Itoa(table))...)
	if err != nil {
		return false, fmt.Errorf("add default route to table %d: %w", table, err)
	}
	return true, nil
}

// EnsureCIDRRoute installs a route for a specific CIDR in a table (idempotent).
func EnsureCIDRRoute(ctx context.Context, r execx.Runner, family, table int, cidr, gateway, dev string) (bool, error) {
	res, err := r.Run(ctx, nil, "ip", ipArgs(family, "route", "show", "table", strconv.Itoa(table), cidr)...)
	if err != nil {
		return false, fmt.Errorf("ip route show %s: %w", cidr, err)
	}
	if strings.TrimSpace(string(res.Stdout)) != "" {
		return false, nil
	}
	_, err = r.Run(ctx, nil, "ip", ipArgs(family,
		"route", "replace", cidr, "via", gateway, "dev", dev, "table", strconv.Itoa(table))...)
	if err != nil {
		return false, fmt.Errorf("add route for %s: %w", cidr, err)
	}
	return true, nil
}

// FlushTable removes every route from the given table. Missing table is not an
// error.
func FlushTable(ctx context.Context, r execx.Runner, family, table int) error {
	_, err := r.Run(ctx, nil, "ip", ipArgs(family, "route", "flush", "table", strconv.Itoa(table))...)
	if err != nil && !isMissing(err) {
		return fmt.Errorf("flush table %d: %w", table, err)
	}
	return nil
}

// TableHasRoutes reports whether the FIB table contains any route.
func TableHasRoutes(ctx context.Context, r execx.Runner, family, table int) (bool, error) {
	res, err := r.Run(ctx, nil, "ip", ipArgs(family, "route", "show", "table", strconv.Itoa(table))...)
	if err != nil {
		if isMissing(err) {
			return false, nil
		}
		return false, err
	}
	return strings.TrimSpace(string(res.Stdout)) != "", nil
}

// RemoveRulesByTable deletes every rule that looks up the given table,
// reconstructing an exact `ip rule del` for each. It is robust against
// priority drift between runs and returns the number of rules removed.
func RemoveRulesByTable(ctx context.Context, r execx.Runner, family, table int) (int, error) {
	entries, err := listRules(ctx, r, family)
	if err != nil {
		return 0, err
	}
	removed := 0
	for _, e := range entries {
		if e.Table != table {
			continue
		}
		args := []string{"rule", "del"}
		if e.HasMark {
			args = append(args, "fwmark", markString(e.Mark))
		}
		if e.HasTo {
			args = append(args, "to", e.To)
		}
		args = append(args, "lookup", strconv.Itoa(table), "priority", strconv.Itoa(e.Priority))
		if _, err := r.Run(ctx, nil, "ip", ipArgs(family, args...)...); err != nil {
			return removed, fmt.Errorf("delete rule: %w", err)
		}
		removed++
	}
	return removed, nil
}

func extractViaDev(fields []string) (string, string) {
	var via, dev string
	for i := 0; i < len(fields); i++ {
		switch fields[i] {
		case "via":
			if i+1 < len(fields) {
				via = fields[i+1]
			}
		case "dev":
			if i+1 < len(fields) {
				dev = fields[i+1]
			}
		}
	}
	return via, dev
}

func isMissing(err error) bool {
	var ee *execx.ExitError
	if errors.As(err, &ee) {
		msg := string(ee.Stderr)
		return strings.Contains(msg, "No such") || strings.Contains(msg, "does not exist")
	}
	return false
}