prx/internal/nft/nft.go

// SPDX-License-Identifier: MIT

// Package nft renders and applies the prx-managed nftables table.
//
// The whole table is (re)defined in a single `nft -f` transaction that begins
// with `destroy table`, which makes application atomic and idempotent: running
// it twice never duplicates rules and there is no window without protection.
package nft

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

	"prx/internal/execx"
)

// TableName is the dedicated nftables table owned by prx.
const TableName = "prx"

var ifaceRE = regexp.MustCompile(`^[A-Za-z0-9_.:@-]{1,15}$`)

// Params are the inputs required to render the table.
type Params struct {
	WANInterface string
	Mark         uint32
	Local4       []string
	Local6       []string
}

// Table returns the fully qualified table reference.
func Table() string { return "inet " + TableName }

// Script renders the nftables ruleset text.
//
// Safety model:
//   - New inbound connections arriving on the WAN interface get the conntrack
//     mark.
//   - Every outgoing packet belonging to a connection whose original
//     destination is one of our WAN addresses is marked too. This covers
//     connections that existed before prx was enabled (e.g. an already open
//     SSH session), which a `ct state new` rule alone would miss.
//
// A separate ip rule then routes marked packets back out the WAN interface.
func Script(p Params) (string, error) {
	if !ifaceRE.MatchString(p.WANInterface) {
		return "", fmt.Errorf("invalid interface name %q", p.WANInterface)
	}
	if p.Mark == 0 {
		return "", errors.New("mark must be non-zero")
	}

	var b strings.Builder
	fmt.Fprintf(&b, "destroy table %s\n", Table())
	fmt.Fprintf(&b, "table %s {\n", Table())

	if len(p.Local4) > 0 {
		fmt.Fprintf(&b, "\tset wan_local4 {\n\t\ttype ipv4_addr;\n\t\telements = { %s }\n\t}\n", strings.Join(p.Local4, ", "))
	}
	if len(p.Local6) > 0 {
		fmt.Fprintf(&b, "\tset wan_local6 {\n\t\ttype ipv6_addr;\n\t\telements = { %s }\n\t}\n", strings.Join(p.Local6, ", "))
	}

	fmt.Fprintf(&b, "\tchain pre {\n")
	fmt.Fprintf(&b, "\t\ttype filter hook prerouting priority mangle; policy accept;\n")
	fmt.Fprintf(&b, "\t\tiifname %q ct state new ct mark set 0x%x\n", p.WANInterface, p.Mark)
	fmt.Fprintf(&b, "\t}\n")

	fmt.Fprintf(&b, "\tchain out {\n")
	fmt.Fprintf(&b, "\t\ttype filter hook output priority mangle; policy accept;\n")
	fmt.Fprintf(&b, "\t\tct mark 0x%x meta mark set 0x%x\n", p.Mark, p.Mark)
	if len(p.Local4) > 0 {
		fmt.Fprintf(&b, "\t\tct original ip daddr @wan_local4 meta mark set 0x%x\n", p.Mark)
	}
	if len(p.Local6) > 0 {
		fmt.Fprintf(&b, "\t\tct original ip6 daddr @wan_local6 meta mark set 0x%x\n", p.Mark)
	}
	fmt.Fprintf(&b, "\t}\n")
	fmt.Fprintf(&b, "}\n")
	return b.String(), nil
}

// Apply installs the ruleset atomically.
func Apply(ctx context.Context, r execx.Runner, script string) error {
	if _, err := r.Run(ctx, []byte(script), "nft", "-f", "-"); err != nil {
		return interpret("nft -f", err)
	}
	return nil
}

// Check validates the ruleset without committing it.
func Check(ctx context.Context, r execx.Runner, script string) error {
	if _, err := r.Run(ctx, []byte(script), "nft", "-c", "-f", "-"); err != nil {
		return interpret("nft -c -f", err)
	}
	return nil
}

// Remove destroys the prx table. It is safe to call when the table is absent.
func Remove(ctx context.Context, r execx.Runner) error {
	_, err := r.Run(ctx, nil, "nft", "destroy", "table", "inet", TableName)
	if err == nil || isAbsent(err) {
		return nil
	}
	return interpret("nft destroy", err)
}

// interpret turns an nft failure into an actionable message. A syntax-style
// failure usually means the host's nftables is older than 0.9.4 and does not
// understand `destroy`, which would otherwise surface as a cryptic error on a
// freshly installed system.
func interpret(op string, err error) error {
	var ee *execx.ExitError
	if !errors.As(err, &ee) {
		return fmt.Errorf("%s: %w", op, err)
	}
	msg := strings.TrimSpace(string(ee.Stderr))
	if msg == "" {
		return fmt.Errorf("%s: exit status %d", op, ee.Code)
	}
	lower := strings.ToLower(msg)
	for _, marker := range []string{"syntax error", "invalid", "unknown command", "unexpected"} {
		if strings.Contains(lower, marker) {
			return fmt.Errorf("%s: %s (nftables may be too old: prx needs >= 0.9.4 — upgrade nftables)", op, msg)
		}
	}
	return fmt.Errorf("%s: %s", op, msg)
}

// Exists reports whether the prx table currently exists.
func Exists(ctx context.Context, r execx.Runner) (bool, error) {
	_, err := r.Run(ctx, nil, "nft", "list", "table", "inet", TableName)
	if err != nil {
		if isAbsent(err) {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

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