From 79755d3ce53baaa19c68e24cae149b6bb5141a40 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Tue, 11 Oct 2022 20:32:24 -0700
Subject: [PATCH] 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 <bradfitz@tailscale.com>
---
 tstest/natlab/firewall.go | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/tstest/natlab/firewall.go b/tstest/natlab/firewall.go
index 1e8cb4e67..c6d6c5a27 100644
--- a/tstest/natlab/firewall.go
+++ b/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.