Persist try snapshot and arm a systemd revert timer

internal/nftables/engine.go

@@ -28,7 +28,8 @@ func (e *Engine) ensureTable() *nftables.Table {
 	})
 }
 
-func (e *Engine) ensureChains(table *nftables.Table) map[string]*nftables.Chain {
+// ensureChains declares the base chains; policies overrides their default policy.
+func (e *Engine) ensureChains(table *nftables.Table, policies map[string]nftables.ChainPolicy) map[string]*nftables.Chain {
 	chains := map[string]*nftables.Chain{
 		"input": {
 			Name:     "input",
@@ -71,6 +72,9 @@ func (e *Engine) ensureChains(table *nftables.Table) map[string]*nftables.Chain
 	}
 
 	for name, chain := range chains {
+		if p, ok := policies[name]; ok {
+			chain.Policy = policyPtr(p)
+		}
 		chains[name] = e.conn.AddChain(chain)
 	}
 	return chains
@@ -93,8 +97,12 @@ func (e *Engine) Plan() (*ChangeSet, error) {
 }
 
 func (e *Engine) Apply(changes *ChangeSet) error {
+	return e.apply(changes, nil)
+}
+
+func (e *Engine) apply(changes *ChangeSet, policies map[string]nftables.ChainPolicy) error {
 	table := e.ensureTable()
-	chains := e.ensureChains(table)
+	chains := e.ensureChains(table, policies)
 
 	for _, r := range changes.Remove {
 		e.conn.DelRule(&nftables.Rule{
@@ -184,34 +192,70 @@ func (e *Engine) readCurrentState() (*FirewallState, error) {
 	return state, nil
 }
 
-// Snapshot is the tomswall table as captured live; a nil state means it was absent.
+// Snapshot is the tomswall table as captured live, serialisable so a revert
+// survives the process that took it.
 type Snapshot struct {
-	state *FirewallState
+	Table    string                          `json:"table"`
+	Present  bool                            `json:"present"`
+	Policies map[string]nftables.ChainPolicy `json:"policies,omitempty"`
+	Rules    map[string][]SnapshotRule       `json:"rules,omitempty"`
+}
+
+// SnapshotRule is a managed rule with its expressions in netlink wire format.
+type SnapshotRule struct {
+	Tag   string   `json:"tag"`
+	Exprs [][]byte `json:"exprs"`
 }
 
 // Snapshot captures the live tomswall table so Restore can roll back to it.
 func (e *Engine) Snapshot() (*Snapshot, error) {
+	snap := &Snapshot{Table: e.cfg.Settings.TableName}
 	t, err := e.findTable()
 	if err != nil || t == nil {
-		return &Snapshot{}, err
+		return snap, err
 	}
+	snap.Present = true
+
+	chains, err := e.conn.ListChainsOfTableFamily(nftables.TableFamilyINet)
+	if err != nil {
+		return nil, fmt.Errorf("listing chains: %w", err)
+	}
+	snap.Policies = make(map[string]nftables.ChainPolicy)
+	for _, c := range chains {
+		if c.Table.Name == snap.Table && c.Policy != nil {
+			snap.Policies[c.Name] = *c.Policy
+		}
+	}
+
 	state, err := e.readCurrentState()
 	if err != nil {
 		return nil, err
 	}
-	return &Snapshot{state: state}, nil
+	snap.Rules, err = encodeState(state)
+	if err != nil {
+		return nil, err
+	}
+	return snap, nil
 }
 
-// Restore atomically returns the tomswall table to the snapshot, rule order included.
+// Restore atomically returns the tomswall table to the snapshot: rule order
+// and chain policies included, or removed if it was absent.
 func (e *Engine) Restore(s *Snapshot) error {
-	if s.state == nil {
+	if s.Table != e.cfg.Settings.TableName {
+		return fmt.Errorf("snapshot is of table %q, engine manages %q", s.Table, e.cfg.Settings.TableName)
+	}
+	if !s.Present {
 		return e.Flush()
 	}
+	want, err := decodeState(s.Rules)
+	if err != nil {
+		return err
+	}
 	current, err := e.readCurrentState()
 	if err != nil {
 		return err
 	}
-	return e.Apply(restoreChangeSet(current, s.state))
+	return e.apply(restoreChangeSet(current, want), s.Policies)
 }
 
 func policyPtr(p nftables.ChainPolicy) *nftables.ChainPolicy {

internal/nftables/snapshot.go

@@ -0,0 +1,119 @@
+package nftables
+
+import (
+	"encoding/binary"
+	"fmt"
+
+	"github.com/google/nftables"
+	"github.com/google/nftables/expr"
+	"github.com/mdlayher/netlink"
+	"golang.org/x/sys/unix"
+)
+
+const inet = byte(nftables.TableFamilyINet)
+
+// exprByName mirrors the expression types google/nftables can parse back from the kernel.
+var exprByName = map[string]func() expr.Any{
+	"ct":           func() expr.Any { return &expr.Ct{} },
+	"range":        func() expr.Any { return &expr.Range{} },
+	"meta":         func() expr.Any { return &expr.Meta{} },
+	"cmp":          func() expr.Any { return &expr.Cmp{} },
+	"counter":      func() expr.Any { return &expr.Counter{} },
+	"objref":       func() expr.Any { return &expr.Objref{} },
+	"payload":      func() expr.Any { return &expr.Payload{} },
+	"lookup":       func() expr.Any { return &expr.Lookup{} },
+	"immediate":    func() expr.Any { return &expr.Immediate{} },
+	"bitwise":      func() expr.Any { return &expr.Bitwise{} },
+	"redir":        func() expr.Any { return &expr.Redir{} },
+	"nat":          func() expr.Any { return &expr.NAT{} },
+	"limit":        func() expr.Any { return &expr.Limit{} },
+	"quota":        func() expr.Any { return &expr.Quota{} },
+	"dynset":       func() expr.Any { return &expr.Dynset{} },
+	"log":          func() expr.Any { return &expr.Log{} },
+	"exthdr":       func() expr.Any { return &expr.Exthdr{} },
+	"connlimit":    func() expr.Any { return &expr.Connlimit{} },
+	"queue":        func() expr.Any { return &expr.Queue{} },
+	"flow_offload": func() expr.Any { return &expr.FlowOffload{} },
+	"reject":       func() expr.Any { return &expr.Reject{} },
+	"masq":         func() expr.Any { return &expr.Masq{} },
+	"hash":         func() expr.Any { return &expr.Hash{} },
+	"notrack":      func() expr.Any { return &expr.Notrack{} },
+}
+
+func encodeState(state *FirewallState) (map[string][]SnapshotRule, error) {
+	out := make(map[string][]SnapshotRule, len(state.Rules))
+	for chain, rules := range state.Rules {
+		for _, r := range rules {
+			sr := SnapshotRule{Tag: r.Tag}
+			for _, e := range r.Exprs {
+				b, err := expr.Marshal(inet, e)
+				if err != nil {
+					return nil, fmt.Errorf("encoding %s rule %q: %w", chain, r.Tag, err)
+				}
+				sr.Exprs = append(sr.Exprs, b)
+			}
+			out[chain] = append(out[chain], sr)
+		}
+	}
+	return out, nil
+}
+
+func decodeState(rules map[string][]SnapshotRule) (*FirewallState, error) {
+	state := &FirewallState{Rules: make(map[string][]ManagedRule, len(rules))}
+	for chain, rs := range rules {
+		for _, sr := range rs {
+			r := ManagedRule{Chain: chain, Tag: sr.Tag}
+			for _, b := range sr.Exprs {
+				e, err := decodeExpr(b)
+				if err != nil {
+					return nil, fmt.Errorf("decoding %s rule %q: %w", chain, sr.Tag, err)
+				}
+				r.Exprs = append(r.Exprs, e)
+			}
+			state.Rules[chain] = append(state.Rules[chain], r)
+		}
+	}
+	return state, nil
+}
+
+// decodeExpr reverses expr.Marshal, as google/nftables does when reading rules.
+func decodeExpr(b []byte) (expr.Any, error) {
+	ad, err := netlink.NewAttributeDecoder(b)
+	if err != nil {
+		return nil, err
+	}
+	ad.ByteOrder = binary.BigEndian
+	var name string
+	var data []byte
+	for ad.Next() {
+		switch ad.Type() {
+		case unix.NFTA_EXPR_NAME:
+			name = ad.String()
+		case unix.NFTA_EXPR_DATA:
+			data = ad.Bytes()
+		}
+	}
+	if err := ad.Err(); err != nil {
+		return nil, err
+	}
+	newExpr, ok := exprByName[name]
+	if !ok {
+		return nil, fmt.Errorf("unsupported expression %q", name)
+	}
+	e := newExpr()
+	if name == "notrack" {
+		return e, nil
+	}
+	if err := expr.Unmarshal(inet, data, e); err != nil {
+		return nil, err
+	}
+	// A verdict is an immediate into the verdict register with no data.
+	if imm, ok := e.(*expr.Immediate); ok && imm.Register == unix.NFT_REG_VERDICT && len(imm.Data) == 0 {
+		v := &expr.Verdict{}
+		if err := expr.Unmarshal(inet, data, v); err != nil {
+			return nil, err
+		}
+		return v, nil
+	}
+	return e, nil
+}

internal/nftables/snapshot_test.go

@@ -0,0 +1,131 @@
+package nftables
+
+import (
+	"encoding/json"
+	"reflect"
+	"testing"
+
+	"github.com/google/nftables"
+	"github.com/google/nftables/expr"
+	"github.com/mdlayher/netlink"
+	"golang.org/x/sys/unix"
+
+	"git.unkin.net/unkin/tomswall/internal/config"
+)
+
+func TestSnapshotRulesRoundTrip(t *testing.T) {
+	exprs := []expr.Any{
+		&expr.Meta{Key: expr.MetaKeyL4PROTO, Register: 1},
+		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: []byte{unix.IPPROTO_TCP}},
+		&expr.Payload{DestRegister: 1, Base: expr.PayloadBaseTransportHeader, Offset: 2, Len: 2},
+		&expr.Cmp{Op: expr.CmpOpEq, Register: 1, Data: []byte{0, 22}},
+		&expr.Ct{Register: 1, Key: expr.CtKeySTATE},
+		&expr.Notrack{},
+		&expr.Verdict{Kind: expr.VerdictAccept},
+	}
+	state := &FirewallState{Rules: map[string][]ManagedRule{
+		"input": {{Chain: "input", Tag: "ssh", Exprs: exprs}, {Chain: "input", Tag: "drop", Exprs: []expr.Any{&expr.Verdict{Kind: expr.VerdictDrop}}}},
+	}}
+
+	rules, err := encodeState(state)
+	if err != nil {
+		t.Fatal(err)
+	}
+	b, err := json.Marshal(&Snapshot{Table: "tomswall", Present: true, Rules: rules,
+		Policies: map[string]nftables.ChainPolicy{"input": nftables.ChainPolicyAccept}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	var snap Snapshot
+	if err := json.Unmarshal(b, &snap); err != nil {
+		t.Fatal(err)
+	}
+	if snap.Policies["input"] != nftables.ChainPolicyAccept {
+		t.Errorf("policy lost: %v", snap.Policies)
+	}
+	got, err := decodeState(snap.Rules)
+	if err != nil {
+		t.Fatal(err)
+	}
+	in := got.Rules["input"]
+	if len(in) != 2 || in[0].Tag != "ssh" || in[1].Tag != "drop" {
+		t.Fatalf("rules/order lost: %+v", in)
+	}
+	if !reflect.DeepEqual(in[0].Exprs, exprs) {
+		t.Errorf("exprs changed:\n got %#v\nwant %#v", in[0].Exprs, exprs)
+	}
+	if _, ok := in[1].Exprs[0].(*expr.Verdict); !ok {
+		t.Errorf("verdict decoded as %T", in[1].Exprs[0])
+	}
+}
+
+func TestEnsureChainsPolicyOverride(t *testing.T) {
+	e := testEngine(t, nil)
+	chains := e.ensureChains(e.ensureTable(), map[string]nftables.ChainPolicy{"input": nftables.ChainPolicyAccept})
+	if *chains["input"].Policy != nftables.ChainPolicyAccept {
+		t.Error("input policy not overridden")
+	}
+	if *chains["forward"].Policy != nftables.ChainPolicyDrop {
+		t.Error("forward policy should keep its default")
+	}
+}
+
+func TestSnapshotAndRestoreAbsentTable(t *testing.T) {
+	tablePresent := false
+	var sent []netlink.HeaderType
+	e := testEngine(t, func(req []netlink.Message) ([]netlink.Message, error) {
+		for _, m := range req {
+			sent = append(sent, m.Header.Type)
+			if m.Header.Type == nftType(unix.NFT_MSG_GETTABLE) && tablePresent {
+				data := []byte{inet, 0, 0, 0}
+				attrs, _ := netlink.MarshalAttributes([]netlink.Attribute{{Type: unix.NFTA_TABLE_NAME, Data: []byte("tomswall\x00")}})
+				return []netlink.Message{{Header: netlink.Header{Type: nftType(unix.NFT_MSG_NEWTABLE), Sequence: m.Header.Sequence}, Data: append(data, attrs...)}}, nil
+			}
+		}
+		return nil, nil
+	})
+
+	snap, err := e.Snapshot()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if snap.Present || snap.Table != "tomswall" {
+		t.Fatalf("want absent tomswall snapshot, got %+v", snap)
+	}
+
+	// The try created the table; restoring the absent snapshot deletes it.
+	tablePresent = true
+	sent = nil
+	if err := e.Restore(snap); err != nil {
+		t.Fatal(err)
+	}
+	deleted := false
+	for _, ht := range sent {
+		deleted = deleted || ht == nftType(unix.NFT_MSG_DELTABLE)
+	}
+	if !deleted {
+		t.Errorf("table not deleted; sent %v", sent)
+	}
+}
+
+func TestRestoreRejectsOtherTable(t *testing.T) {
+	e := testEngine(t, nil)
+	if err := e.Restore(&Snapshot{Table: "other"}); err == nil {
+		t.Error("expected table mismatch error")
+	}
+}
+
+func nftType(msg int) netlink.HeaderType {
+	return netlink.HeaderType(unix.NFNL_SUBSYS_NFTABLES<<8 | msg)
+}
+
+func testEngine(t *testing.T, dial func([]netlink.Message) ([]netlink.Message, error)) *Engine {
+	if dial == nil {
+		dial = func([]netlink.Message) ([]netlink.Message, error) { return nil, nil }
+	}
+	conn, err := nftables.New(nftables.WithTestDial(dial))
+	if err != nil {
+		t.Fatal(err)
+	}
+	return &Engine{cfg: &config.Config{Settings: config.Settings{TableName: "tomswall"}}, conn: conn}
+}