From d7bc4ec02988e5da059813c81a486439c3ee8922 Mon Sep 17 00:00:00 2001
From: David Anderson <danderson@tailscale.com>
Date: Fri, 3 Jul 2020 02:26:54 +0000
Subject: [PATCH] tstest/natlab: use common helper for conn registration.

Signed-off-by: David Anderson <danderson@tailscale.com>
---
 tstest/natlab/natlab.go | 29 +++++++++++++----------------
 1 file changed, 13 insertions(+), 16 deletions(-)

diff --git a/tstest/natlab/natlab.go b/tstest/natlab/natlab.go
index 7f75ba165..8137799b2 100644
--- a/tstest/natlab/natlab.go
+++ b/tstest/natlab/natlab.go
@@ -371,14 +371,7 @@ func (m *Machine) registerConn4(c *conn) error {
 	if c.ipp.IP.Is6() && c.ipp.IP != v6unspec {
 		return fmt.Errorf("registerConn4 got IPv6 %s", c.ipp)
 	}
-	if _, ok := m.conns4[c.ipp]; ok {
-		return fmt.Errorf("duplicate conn listening on %v", c.ipp)
-	}
-	if m.conns4 == nil {
-		m.conns4 = map[netaddr.IPPort]*conn{}
-	}
-	m.conns4[c.ipp] = c
-	return nil
+	return registerConn(&m.conns4, c)
 }
 
 func (m *Machine) unregisterConn4(c *conn) {
@@ -393,14 +386,7 @@ func (m *Machine) registerConn6(c *conn) error {
 	if c.ipp.IP.Is4() {
 		return fmt.Errorf("registerConn6 got IPv4 %s", c.ipp)
 	}
-	if _, ok := m.conns6[c.ipp]; ok {
-		return fmt.Errorf("duplicate conn listening on %v", c.ipp)
-	}
-	if m.conns6 == nil {
-		m.conns6 = map[netaddr.IPPort]*conn{}
-	}
-	m.conns6[c.ipp] = c
-	return nil
+	return registerConn(&m.conns6, c)
 }
 
 func (m *Machine) unregisterConn6(c *conn) {
@@ -409,6 +395,17 @@ func (m *Machine) unregisterConn6(c *conn) {
 	delete(m.conns6, c.ipp)
 }
 
+func registerConn(conns *map[netaddr.IPPort]*conn, c *conn) error {
+	if _, ok := (*conns)[c.ipp]; ok {
+		return fmt.Errorf("duplicate conn listening on %v", c.ipp)
+	}
+	if *conns == nil {
+		*conns = map[netaddr.IPPort]*conn{}
+	}
+	(*conns)[c.ipp] = c
+	return nil
+}
+
 func (m *Machine) AddNetwork(n *Network) {}
 
 func (m *Machine) ListenPacket(network, address string) (net.PacketConn, error) {