TAP -> TUN, handle multiple peers behind one NAT

server/main.go

@@ -1,14 +1,15 @@
 // TODO: slice bounds checking
 // TODO: authentication
 // TODO: somehow persist IPs
-// TODO: handle multiple peers behind one NAT
 // TODO: key rotation
+// TODO: dont trust claimed IP at all
 package main
 
 import (
 	"crypto/hpke"
 	"crypto/rand"
 	"encoding/binary"
+	"encoding/hex"
 	"log"
 	"net"
 	"os"
@@ -17,14 +18,15 @@ import (
 )
 
 type Peer struct {
-	RealAddr   *net.UDPAddr
-	InternalIP string
-	PublicKey  hpke.PublicKey
+	RealAddr     *net.UDPAddr
+	InternalIP   string
+	HexPublicKey string
 }
 
 var (
 	peers           = shared.NewTMap[string, *Peer]()
 	peersByInternal = shared.NewTMap[string, *Peer]()
+	peersByPubKey   = shared.NewTMap[string, *Peer]()
 	conn            *net.UDPConn
 	multicastKey    = make([]byte, 32)
 )
@@ -54,8 +56,8 @@ func receivePackets() {
 	}
 	defer conn.Close()
 
+	buffer := make([]byte, 50000)
 	for {
-		buffer := make([]byte, 50000)
 		n, addr, err := conn.ReadFromUDP(buffer)
 		if err != nil {
 			log.Println(err)
@@ -65,14 +67,16 @@ func receivePackets() {
 			continue
 		}
 
-		go handleReq(buffer[:n], addr)
+		req := make([]byte, n)
+		copy(req, buffer[:n])
+		go handleReq(req, addr)
 	}
 }
 
 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)
+			log.Printf("recovered from panic while handling request from %s: %v\n", addr.String(), r)
 		}
 	}()
 
@@ -96,19 +100,25 @@ func handleReq(req []byte, addr *net.UDPAddr) {
 			return
 		}
 
-		internalIP := randomIP().String()
+		hexPubKey := hex.EncodeToString(pubKey.Bytes())
 
-		if destPeer, ok := peers.GetOK(addr.IP.String()); ok {
-			internalIP = destPeer.InternalIP
-			destPeer.PublicKey = pubKey
+		var internalIP string
+		// TODO: authenticating only by PUBLIC key is obviously a bad idea
+		if peer, ok := peersByPubKey.GetOK(hexPubKey); ok {
+			internalIP = peer.InternalIP
+			peers.Delete(peer.RealAddr.String())
+			peer.RealAddr = addr
+			peers.Set(addr.String(), peer)
 		} else {
+			internalIP = randomIP().String()
 			peer := &Peer{
-				RealAddr:   addr,
-				InternalIP: internalIP,
-				PublicKey:  pubKey,
+				RealAddr:     addr,
+				InternalIP:   internalIP,
+				HexPublicKey: hexPubKey,
 			}
-			peers.Set(addr.IP.String(), peer)
+			peers.Set(addr.String(), peer)
 			peersByInternal.Set(internalIP, peer)
+			peersByPubKey.Set(hexPubKey, peer)
 		}
 
 		ciphertext, err := hpke.Seal(pubKey, hpke.HKDFSHA256(), hpke.ChaCha20Poly1305(), []byte("baalvpn"), multicastKey)
@@ -119,7 +129,7 @@ func handleReq(req []byte, addr *net.UDPAddr) {
 		resp := shared.BuildPkt(shared.RESP_REGISTER, append(net.ParseIP(internalIP), ciphertext...))
 		sendTo(addr, resp)
 	case shared.ENC_PKT:
-		peer := peers.Get(addr.IP.String())
+		peer := peers.Get(addr.String())
 		if peer == nil {
 			log.Println("data from unregistered peer:", addr.String())
 			return
@@ -139,7 +149,7 @@ func handleReq(req []byte, addr *net.UDPAddr) {
 			log.Println(srcIP + " -/> " + destIP + " (unrecognized IP)")
 		}
 	case shared.BROADCAST_PKT:
-		peer := peers.Get(addr.IP.String())
+		peer := peers.Get(addr.String())
 		if peer == nil {
 			log.Println("data from unregistered peer:", addr.String())
 			return
@@ -150,15 +160,16 @@ func handleReq(req []byte, addr *net.UDPAddr) {
 	case shared.REQ_GET_PUBKEY:
 		peerIP := net.IP(req[4:20]).String()
 		peer := peersByInternal.Get(peerIP)
-		if peer == nil {
-			log.Println("tried to get pubkey of unknown peer")
-			return
+		if peer != nil {
+			pubKey, err := hex.DecodeString(peer.HexPublicKey)
+			if err != nil {
+				panic(err)
+			}
+			resp := shared.BuildPkt(shared.RESP_GET_PUBKEY, append(net.ParseIP(peerIP), pubKey...))
+			sendTo(addr, resp)
 		}
-
-		resp := shared.BuildPkt(shared.RESP_GET_PUBKEY, append(net.ParseIP(peer.InternalIP), peer.PublicKey.Bytes()...))
-		sendTo(addr, resp)
 	case shared.REQ_ESTABLISH, shared.RESP_ESTABLISH:
-		peer := peers.Get(addr.IP.String())
+		peer := peers.Get(addr.String())
 		if peer == nil {
 			log.Println("unregistered peer tried to ESTABLISH")
 			return