thread safe map

client/main.go

@@ -8,6 +8,7 @@ import (
 	"net"
 	"os"
 	"os/exec"
+	"sync"
 	"time"
 
 	"baalvpn/shared"
@@ -28,12 +29,10 @@ var (
 	conn       *net.UDPConn
 	privKey    *xwing.DecapsulationKey
 
-	// TODO: this should probably have a mutex
+	establishMutex    sync.Mutex
-	peerPubKeys = map[string][]byte{}
+	peerPubKeys       = shared.NewTMap[string, []byte]()
-	// TODO: this should probably have a mutex
+	peerSessionKeys   = shared.NewTMap[string, []byte]()
-	peerSessionKeys = map[string][]byte{}
+	peerEstablishAcks = shared.NewTMap[string, bool]()
-	// TODO: this should probably have a mutex
-	peerEstablishAcks = map[string]bool{}
 )
 
 func main() {
@@ -77,45 +76,53 @@ func register() {
 }
 
 func getOrEstablishSessionKey(ip string) []byte {
-	if key, ok := peerSessionKeys[ip]; ok {
+	if key, ok := peerSessionKeys.GetOK(ip); ok {
 		return key
-	} else {
+	}
-		log.Println("requesting pubkey of " + ip)
-		req := shared.BuildPkt(shared.REQ_GET_PUBKEY, net.ParseIP(ip))
-		if _, err := conn.WriteToUDP(req, serverAddr); err != nil {
-			panic(err)
-		}
 
-		for {
+	establishMutex.Lock()
-			// TODO: eww
+	defer establishMutex.Unlock()
-			time.Sleep(50 * time.Millisecond)
-			if peerPubKey, ok := peerPubKeys[ip]; ok {
-				log.Println("got pubkey of " + ip)
 
-				ciphertext, sharedSecret, err := xwing.Encapsulate(peerPubKey)
+	// check again after the other goroutine finished
-				if err != nil {
+	if key, ok := peerSessionKeys.GetOK(ip); ok {
-					panic(err)
+		return key
-				}
+	}
 
-				reqData := append(net.ParseIP(ip), append(internalIP, ciphertext...)...)
+	log.Println("requesting pubkey of " + ip)
-				req := shared.BuildPkt(shared.REQ_ESTABLISH, reqData)
+	req := shared.BuildPkt(shared.REQ_GET_PUBKEY, net.ParseIP(ip))
-				if _, err := conn.WriteToUDP(req, serverAddr); err != nil {
+	if _, err := conn.WriteToUDP(req, serverAddr); err != nil {
-					panic(err)
+		panic(err)
-				}
+	}
 
-				for {
+	for {
-					// TODO: eww
+		// TODO: eww
-					time.Sleep(50 * time.Millisecond)
+		time.Sleep(50 * time.Millisecond)
-					if peerEstablishAcks[ip] {
+		if peerPubKey, ok := peerPubKeys.GetOK(ip); ok {
-						delete(peerEstablishAcks, ip)
+			log.Println("got pubkey of " + ip)
-						break
-					}
-				}
 
-				sessionKey := sha256.Sum256(sharedSecret)
+			ciphertext, sharedSecret, err := xwing.Encapsulate(peerPubKey)
-				peerSessionKeys[ip] = sessionKey[:]
+			if err != nil {
-				return sessionKey[:]
+				panic(err)
 			}
+
+			reqData := append(net.ParseIP(ip), append(internalIP, ciphertext...)...)
+			req := shared.BuildPkt(shared.REQ_ESTABLISH, reqData)
+			if _, err := conn.WriteToUDP(req, serverAddr); err != nil {
+				panic(err)
+			}
+
+			for {
+				// TODO: eww
+				time.Sleep(50 * time.Millisecond)
+				if peerEstablishAcks.Get(ip) {
+					peerEstablishAcks.Delete(ip)
+					break
+				}
+			}
+
+			sessionKey := sha256.Sum256(sharedSecret)
+			peerSessionKeys.Set(ip, sessionKey[:])
+			return sessionKey[:]
 		}
 	}
 }
@@ -136,11 +143,17 @@ func receivePackets() {
 			continue
 		}
 
-		go receivePacket(buffer[:n])
+		go handlePkt(buffer[:n])
 	}
 }
 
-func receivePacket(pkt []byte) {
+func handlePkt(pkt []byte) {
+	defer func() {
+		if r := recover(); r != nil {
+			log.Printf("recovered from panic while handling a packet: %v\n", r)
+		}
+	}()
+
 	if len(pkt) < 4 {
 		log.Println("packet too short")
 		return
@@ -181,7 +194,7 @@ func receivePacket(pkt []byte) {
 	case shared.RESP_GET_PUBKEY:
 		respIP := net.IP(pkt[4:20]).String()
 		pubKey := pkt[20:]
-		peerPubKeys[respIP] = pubKey
+		peerPubKeys.Set(respIP, pubKey)
 	case shared.REQ_ESTABLISH:
 		srcIP := net.IP(pkt[20:36]).String()
 		ciphertext := pkt[36:]
@@ -192,7 +205,7 @@ func receivePacket(pkt []byte) {
 		}
 
 		sessionKey := sha256.Sum256(sharedSecret)
-		peerSessionKeys[srcIP] = sessionKey[:]
+		peerSessionKeys.Set(srcIP, sessionKey[:])
 
 		log.Println("received session key from", srcIP)
 		reqData := append(net.ParseIP(srcIP), internalIP...)
@@ -203,7 +216,7 @@ func receivePacket(pkt []byte) {
 	case shared.RESP_ESTABLISH:
 		srcIP := net.IP(pkt[20:36]).String()
 		log.Println("established session key with", srcIP)
-		peerEstablishAcks[srcIP] = true
+		peerEstablishAcks.Set(srcIP, true)
 	default:
 		log.Println("unknown packet type")
 	}

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{}
+	peers = shared.NewTMap[string, *Peer]()
-	peersMutex sync.Mutex
+	conn  *net.UDPConn
-	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.Get(addr.IP.String())
-		peer := peers[addr.IP.String()]
-		peersMutex.Unlock()
 		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.Get(addr.IP.String())
-		peer := peers[addr.IP.String()]
-		peersMutex.Unlock()
 		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.Get(addr.IP.String())
-		peer := peers[addr.IP.String()]
+		if peer == nil {
-		peersMutex.Unlock()
+			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()
+	peers.ForEach(func(k string, v *Peer) {
-	defer peersMutex.Unlock()
+		if v.InternalIP == except {
-	for _, p := range peers {
+			return
-		if p.InternalIP == except {
-			continue
 		}
-		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()
+	var out *Peer = nil
-	defer peersMutex.Unlock()
+	peers.ForEach(func(k string, v *Peer) {
-	for _, p := range peers {
+		if v.InternalIP == internalIP {
-		if p.InternalIP == internalIP {
+			out = v
-			return p
 		}
-	}
+	})
-	return nil
+	return out
 }
 
 // fd00:baa1::/32

shared/tmap.go

@@ -0,0 +1,47 @@
+package shared
+
+import "sync"
+
+type TMap[K comparable, V any] struct {
+	data map[K]V
+	mu   sync.RWMutex
+}
+
+func NewTMap[K comparable, V any]() TMap[K, V] {
+	return TMap[K, V]{
+		data: map[K]V{},
+	}
+}
+
+func (m *TMap[K, V]) Get(k K) V {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return m.data[k]
+}
+
+func (m *TMap[K, V]) GetOK(k K) (V, bool) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	v, ok := m.data[k]
+	return v, ok
+}
+
+func (m *TMap[K, V]) Set(k K, v V) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.data[k] = v
+}
+
+func (m *TMap[K, V]) Delete(k K) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	delete(m.data, k)
+}
+
+func (m *TMap[K, V]) ForEach(f func(k K, v V)) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	for k, v := range m.data {
+		f(k, v)
+	}
+}