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 {