tstest/natlab: add Firewall.Reset method to drop firewall state

For future use in magicsock tests.

Updates #540

Change-Id: I2f07d1a2924f20b36e357c4533ff0a1a974d5061
Signed-off-by: Brad Fitzpatrick <[email protected]>

tstest/natlab/firewall.go

@@ -9,6 +9,8 @@ import (
 	"net/netip"
 	"sync"
 	"time"
+
+	"tailscale.com/util/mak"
 )
 
 // FirewallType is the type of filtering a stateful firewall
@@ -100,19 +102,19 @@ func (f *Firewall) timeNow() time.Time {
 	return time.Now()
 }
 
-func (f *Firewall) init() {
-	if f.seen == nil {
-		f.seen = map[fwKey]time.Time{}
-	}
+// Reset drops all firewall state, forgetting all flows.
+func (f *Firewall) Reset() {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.seen = nil
 }
 
 func (f *Firewall) HandleOut(p *Packet, oif *Interface) *Packet {
 	f.mu.Lock()
 	defer f.mu.Unlock()
-	f.init()
 
 	k := f.Type.key(p.Src, p.Dst)
-	f.seen[k] = f.timeNow().Add(f.sessionTimeoutLocked())
+	mak.Set(&f.seen, k, f.timeNow().Add(f.sessionTimeoutLocked()))
 	p.Trace("firewall out ok")
 	return p
 }
@@ -120,7 +122,6 @@ func (f *Firewall) HandleOut(p *Packet, oif *Interface) *Packet {
 func (f *Firewall) HandleIn(p *Packet, iif *Interface) *Packet {
 	f.mu.Lock()
 	defer f.mu.Unlock()
-	f.init()
 
 	// reverse src and dst because the session table is from the POV
 	// of outbound packets.