peersByInternal

server/main.go

@@ -23,9 +23,10 @@ type Peer struct {
 }
 
 var (
-	peers        = shared.NewTMap[string, *Peer]()
-	conn         *net.UDPConn
-	multicastKey = make([]byte, 32)
+	peers           = shared.NewTMap[string, *Peer]()
+	peersByInternal = shared.NewTMap[string, *Peer]()
+	conn            *net.UDPConn
+	multicastKey    = make([]byte, 32)
 )
 
 func main() {
@@ -99,12 +100,15 @@ func handleReq(req []byte, addr *net.UDPAddr) {
 
 		if destPeer, ok := peers.GetOK(addr.IP.String()); ok {
 			internalIP = destPeer.InternalIP
+			destPeer.PublicKey = pubKey
 		} else {
-			peers.Set(addr.IP.String(), &Peer{
+			peer := &Peer{
 				RealAddr:   addr,
 				InternalIP: internalIP,
 				PublicKey:  pubKey,
-			})
+			}
+			peers.Set(addr.IP.String(), peer)
+			peersByInternal.Set(internalIP, peer)
 		}
 
 		ciphertext, err := hpke.Seal(pubKey, hpke.HKDFSHA256(), hpke.ChaCha20Poly1305(), []byte("baalvpn"), multicastKey)
@@ -128,7 +132,7 @@ func handleReq(req []byte, addr *net.UDPAddr) {
 			return
 		}
 
-		if destPeer := getPeerByInternalIP(destIP); destPeer != nil {
+		if destPeer := peersByInternal.Get(destIP); destPeer != nil {
 			log.Println(srcIP + " -> " + destIP)
 			sendTo(destPeer.RealAddr, req)
 		} else {
@@ -145,9 +149,9 @@ func handleReq(req []byte, addr *net.UDPAddr) {
 		broadcast(req, peer.InternalIP)
 	case shared.REQ_GET_PUBKEY:
 		peerIP := net.IP(req[4:20]).String()
-		peer := getPeerByInternalIP(peerIP)
+		peer := peersByInternal.Get(peerIP)
 		if peer == nil {
-			log.Println(addr.String(), "tried to get pubkey of unknown peer", peerIP)
+			log.Println("tried to get pubkey of unknown peer")
 			return
 		}
 
@@ -156,7 +160,7 @@ func handleReq(req []byte, addr *net.UDPAddr) {
 	case shared.REQ_ESTABLISH, shared.RESP_ESTABLISH:
 		peer := peers.Get(addr.IP.String())
 		if peer == nil {
-			log.Println("unregistered peer", addr.String(), "tried to ESTABLISH")
+			log.Println("unregistered peer tried to ESTABLISH")
 			return
 		}
 
@@ -167,7 +171,11 @@ func handleReq(req []byte, addr *net.UDPAddr) {
 		}
 
 		dstIP := net.IP(req[4:20]).String()
-		dstPeer := getPeerByInternalIP(dstIP)
+		dstPeer := peersByInternal.Get(dstIP)
+		if dstPeer == nil {
+			log.Println("tried to ESTABLISH with an unknown peer")
+			return
+		}
 		sendTo(dstPeer.RealAddr, req)
 	default:
 		log.Println("unknown packet type")
@@ -189,16 +197,6 @@ func sendTo(addr *net.UDPAddr, pkt []byte) {
 	}
 }
 
-func getPeerByInternalIP(internalIP string) *Peer {
-	var out *Peer = nil
-	peers.ForEach(func(k string, v *Peer) {
-		if v.InternalIP == internalIP {
-			out = v
-		}
-	})
-	return out
-}
-
 // fd00:baa1::/32
 func randomIP() net.IP {
 	ip := make(net.IP, 16)