tomswall/internal/nftables/engine.go

package nftables

import (
	"fmt"

	"github.com/google/nftables"

	"git.unkin.net/unkin/tomswall/internal/config"
)

type Engine struct {
	cfg  *config.Config
	conn *nftables.Conn
}

func NewEngine(cfg *config.Config) (*Engine, error) {
	conn, err := nftables.New()
	if err != nil {
		return nil, fmt.Errorf("connecting to nftables: %w", err)
	}
	return &Engine{cfg: cfg, conn: conn}, nil
}

func (e *Engine) ensureTable() *nftables.Table {
	return e.conn.AddTable(&nftables.Table{
		Family: nftables.TableFamilyINet,
		Name:   e.cfg.Settings.TableName,
	})
}

func (e *Engine) ensureChains(table *nftables.Table) map[string]*nftables.Chain {
	chains := map[string]*nftables.Chain{
		"input": {
			Name:     "input",
			Table:    table,
			Type:     nftables.ChainTypeFilter,
			Hooknum:  nftables.ChainHookInput,
			Priority: nftables.ChainPriorityFilter,
			Policy:   policyPtr(nftables.ChainPolicyDrop),
		},
		"forward": {
			Name:     "forward",
			Table:    table,
			Type:     nftables.ChainTypeFilter,
			Hooknum:  nftables.ChainHookForward,
			Priority: nftables.ChainPriorityFilter,
			Policy:   policyPtr(nftables.ChainPolicyDrop),
		},
		"output": {
			Name:     "output",
			Table:    table,
			Type:     nftables.ChainTypeFilter,
			Hooknum:  nftables.ChainHookOutput,
			Priority: nftables.ChainPriorityFilter,
			Policy:   policyPtr(nftables.ChainPolicyAccept),
		},
		"postrouting": {
			Name:     "postrouting",
			Table:    table,
			Type:     nftables.ChainTypeNAT,
			Hooknum:  nftables.ChainHookPostrouting,
			Priority: nftables.ChainPriorityNATSource,
		},
		"prerouting": {
			Name:     "prerouting",
			Table:    table,
			Type:     nftables.ChainTypeNAT,
			Hooknum:  nftables.ChainHookPrerouting,
			Priority: nftables.ChainPriorityNATDest,
		},
	}

	for name, chain := range chains {
		chains[name] = e.conn.AddChain(chain)
	}
	return chains
}

func (e *Engine) Plan() (*ChangeSet, error) {
	compiler := NewCompiler(e.cfg)

	desired, err := compiler.Compile()
	if err != nil {
		return nil, fmt.Errorf("compiling config: %w", err)
	}

	current, err := e.readCurrentState()
	if err != nil {
		return nil, fmt.Errorf("reading current state: %w", err)
	}

	return computeDiff(current, desired), nil
}

func (e *Engine) Apply(changes *ChangeSet) error {
	table := e.ensureTable()
	chains := e.ensureChains(table)

	for _, r := range changes.Remove {
		e.conn.DelRule(&nftables.Rule{
			Table:  table,
			Chain:  chains[r.Chain],
			Handle: r.Handle,
		})
	}

	for _, r := range changes.Add {
		chain, ok := chains[r.Chain]
		if !ok {
			return fmt.Errorf("unknown chain %q", r.Chain)
		}
		e.conn.AddRule(&nftables.Rule{
			Table:    table,
			Chain:    chain,
			Exprs:    r.Exprs,
			UserData: []byte(r.Tag),
		})
	}

	return e.conn.Flush()
}

func (e *Engine) Flush() error {
	tables, err := e.conn.ListTables()
	if err != nil {
		return fmt.Errorf("listing tables: %w", err)
	}

	for _, t := range tables {
		if t.Name == e.cfg.Settings.TableName {
			e.conn.DelTable(t)
			return e.conn.Flush()
		}
	}
	return nil
}

func (e *Engine) findTable() (*nftables.Table, error) {
	tables, err := e.conn.ListTables()
	if err != nil {
		return nil, fmt.Errorf("listing tables: %w", err)
	}
	for _, t := range tables {
		if t.Name == e.cfg.Settings.TableName && t.Family == nftables.TableFamilyINet {
			return t, nil
		}
	}
	return nil, nil
}

func (e *Engine) readCurrentState() (*FirewallState, error) {
	state := &FirewallState{
		Rules: make(map[string][]ManagedRule),
	}

	ourTable, err := e.findTable()
	if err != nil || ourTable == nil {
		return state, err
	}

	chains, err := e.conn.ListChainsOfTableFamily(nftables.TableFamilyINet)
	if err != nil {
		return nil, fmt.Errorf("listing chains: %w", err)
	}

	for _, chain := range chains {
		if chain.Table.Name != e.cfg.Settings.TableName {
			continue
		}
		rules, err := e.conn.GetRules(ourTable, chain)
		if err != nil {
			return nil, fmt.Errorf("listing rules of %s: %w", chain.Name, err)
		}
		for _, rule := range rules {
			state.Rules[chain.Name] = append(state.Rules[chain.Name], ManagedRule{
				Chain:  chain.Name,
				Handle: rule.Handle,
				Exprs:  rule.Exprs,
				Tag:    string(rule.UserData),
			})
		}
	}

	return state, nil
}

// Snapshot is the tomswall table as captured live; a nil state means it was absent.
type Snapshot struct {
	state *FirewallState
}

// Snapshot captures the live tomswall table so Restore can roll back to it.
func (e *Engine) Snapshot() (*Snapshot, error) {
	t, err := e.findTable()
	if err != nil || t == nil {
		return &Snapshot{}, err
	}
	state, err := e.readCurrentState()
	if err != nil {
		return nil, err
	}
	return &Snapshot{state: state}, nil
}

// Restore atomically returns the tomswall table to the snapshot, rule order included.
func (e *Engine) Restore(s *Snapshot) error {
	if s.state == nil {
		return e.Flush()
	}
	current, err := e.readCurrentState()
	if err != nil {
		return err
	}
	return e.Apply(restoreChangeSet(current, s.state))
}

func policyPtr(p nftables.ChainPolicy) *nftables.ChainPolicy {
	return &p
}