thread safe map

server/main.go

@@ -4,7 +4,6 @@
 // TODO: handle multiple peers behind one NAT
 // TODO: key rotation
 // TODO: sha256 -> HKDF
-// TODO: recovery
 package main
 
 import (
@@ -13,7 +12,6 @@ import (
 	"log"
 	"net"
 	"os"
-	"sync"
 
 	"baalvpn/shared"
 )
@@ -25,9 +23,8 @@ type Peer struct {
 }
 
 var (
-	peers      = map[string]*Peer{}
-	peersMutex sync.Mutex
-	conn       *net.UDPConn
+	peers = shared.NewTMap[string, *Peer]()
+	conn  *net.UDPConn
 )
 
 func main() {
@@ -67,6 +64,12 @@ func receivePackets() {
 }
 
 func handleReq(req []byte, addr *net.UDPAddr) {
+	defer func() {
+		if r := recover(); r != nil {
+			log.Printf("recovered from panic while handling request from %s: %v\n", addr.IP.String(), r)
+		}
+	}()
+
 	if len(req) < 4 {
 		log.Println("packet too short")
 		return
@@ -81,29 +84,24 @@ func handleReq(req []byte, addr *net.UDPAddr) {
 
 	switch pktType {
 	case shared.REQ_REGISTER:
-		peersMutex.Lock()
-		defer peersMutex.Unlock()
-
 		pubKey := req[4:]
 
 		internalIP := randomIP().String()
 
-		if destPeer, ok := peers[addr.IP.String()]; ok {
+		if destPeer, ok := peers.GetOK(addr.IP.String()); ok {
 			internalIP = destPeer.InternalIP
 		} else {
-			peers[addr.IP.String()] = &Peer{
+			peers.Set(addr.IP.String(), &Peer{
 				RealAddr:   addr,
 				InternalIP: internalIP,
 				PublicKey:  pubKey,
-			}
+			})
 		}
 
 		resp := shared.BuildPkt(shared.RESP_REGISTER, net.ParseIP(internalIP))
 		sendTo(addr, resp)
 	case shared.ENC_PKT:
-		peersMutex.Lock()
-		peer := peers[addr.IP.String()]
-		peersMutex.Unlock()
+		peer := peers.Get(addr.IP.String())
 		if peer == nil {
 			log.Println("data from unregistered peer:", addr.String())
 			return
@@ -123,9 +121,7 @@ func handleReq(req []byte, addr *net.UDPAddr) {
 			log.Println(srcIP + " -/> " + destIP + " (unrecognized IP)")
 		}
 	case shared.BROADCAST_PKT:
-		peersMutex.Lock()
-		peer := peers[addr.IP.String()]
-		peersMutex.Unlock()
+		peer := peers.Get(addr.IP.String())
 		if peer == nil {
 			log.Println("data from unregistered peer:", addr.String())
 			return
@@ -145,9 +141,11 @@ func handleReq(req []byte, addr *net.UDPAddr) {
 		resp := shared.BuildPkt(shared.RESP_GET_PUBKEY, append(net.ParseIP(peer.InternalIP), peer.PublicKey...))
 		sendTo(addr, resp)
 	case shared.REQ_ESTABLISH, shared.RESP_ESTABLISH:
-		peersMutex.Lock()
-		peer := peers[addr.IP.String()]
-		peersMutex.Unlock()
+		peer := peers.Get(addr.IP.String())
+		if peer == nil {
+			log.Println("unregistered peer", addr.String(), "tried to ESTABLISH")
+			return
+		}
 
 		srcIP := net.IP(req[20:36]).String()
 		if srcIP != peer.InternalIP {
@@ -164,14 +162,12 @@ func handleReq(req []byte, addr *net.UDPAddr) {
 }
 
 func broadcast(pkt []byte, except string) {
-	peersMutex.Lock()
-	defer peersMutex.Unlock()
-	for _, p := range peers {
-		if p.InternalIP == except {
-			continue
+	peers.ForEach(func(k string, v *Peer) {
+		if v.InternalIP == except {
+			return
 		}
-		sendTo(p.RealAddr, pkt)
-	}
+		go sendTo(v.RealAddr, pkt)
+	})
 }
 
 func sendTo(addr *net.UDPAddr, pkt []byte) {
@@ -181,14 +177,13 @@ func sendTo(addr *net.UDPAddr, pkt []byte) {
 }
 
 func getPeerByInternalIP(internalIP string) *Peer {
-	peersMutex.Lock()
-	defer peersMutex.Unlock()
-	for _, p := range peers {
-		if p.InternalIP == internalIP {
-			return p
+	var out *Peer = nil
+	peers.ForEach(func(k string, v *Peer) {
+		if v.InternalIP == internalIP {
+			out = v
 		}
-	}
-	return nil
+	})
+	return out
 }
 
 // fd00:baa1::/32