From 36d4b275a0252862ba4593ecc436b626a93bcbb4 Mon Sep 17 00:00:00 2001
From: sbruens <sbruens@users.noreply.github.com>
Date: Tue, 10 Dec 2024 16:51:46 -0500
Subject: [PATCH] Simplify the natmap a little.

---
 service/udp.go | 32 +++++++++-----------------------
 1 file changed, 9 insertions(+), 23 deletions(-)

diff --git a/service/udp.go b/service/udp.go
index 4b822ce0..57818a78 100644
--- a/service/udp.go
+++ b/service/udp.go
@@ -179,7 +179,7 @@ func PacketServe(clientConn net.PacketConn, handle AssocationHandleFunc, metrics
 			}
 			metrics.AddNATEntry()
 			deleteEntry := nm.Add(addr, conn)
-			go func(conn *natconn) {
+			go func(conn net.Conn) {
 				defer func() {
 					conn.Close()
 					deleteEntry()
@@ -239,6 +239,7 @@ func (c *natconn) Close() error {
 	close(c.bytesReadCh)
 	return c.PacketConn.Close()
 }
+
 func (c *natconn) RemoteAddr() net.Addr {
 	return c.raddr
 }
@@ -437,33 +438,18 @@ func (m *natmap) Get(key string) *natconn {
 	return m.keyConn[key]
 }
 
-func (m *natmap) set(key string, pc *natconn) {
-	m.Lock()
-	defer m.Unlock()
-
-	m.keyConn[key] = pc
-	return
-}
-
-func (m *natmap) del(key string) *natconn {
-	m.Lock()
-	defer m.Unlock()
-
-	entry, ok := m.keyConn[key]
-	if ok {
-		delete(m.keyConn, key)
-		return entry
-	}
-	return nil
-}
-
 // Add adds a new UDP NAT entry to the natmap and returns a closure to delete
 // the entry.
 func (m *natmap) Add(addr net.Addr, pc *natconn) func() {
+	m.Lock()
+	defer m.Unlock()
+
 	key := addr.String()
-	m.set(key, pc)
+	m.keyConn[key] = pc
 	return func() {
-		m.del(key)
+		m.Lock()
+		defer m.Unlock()
+		delete(m.keyConn, key)
 	}
 }