peersByInternal

README.md

@@ -1,6 +1,6 @@
 # baalvpn
 
-A post-quantum VPN for IPv6.
+A post-quantum VPN for IPv6
 
 ## Features
 * **Post-quantum security:** Uses [XWingMLKEM768X25519](https://datatracker.ietf.org/doc/html/draft-connolly-cfrg-xwing-kem/) + [XChaCha20-Poly1305](https://en.wikipedia.org/wiki/ChaCha20-Poly1305#XChaCha20-Poly1305_%E2%80%93_extended_nonce_variant) + [HKDF-SHA-256](https://en.wikipedia.org/wiki/HKDF) for end-to-end encryption

client/main.go

@@ -114,10 +114,9 @@ func handleIncomingPkt(pkt []byte) {
 	switch pktType {
 	case shared.RESP_REGISTER:
 		internalIP = net.IP(pkt[4:20])
-		encryptedMulticastKey := net.IP(pkt[20:])
 
 		var err error
-		multicastKey, err = hpke.Open(privKey, hpke.HKDFSHA256(), hpke.ChaCha20Poly1305(), []byte("baalvpn"), encryptedMulticastKey)
+		multicastKey, err = hpke.Open(privKey, hpke.HKDFSHA256(), hpke.ChaCha20Poly1305(), []byte("baalvpn"), pkt[20:])
 		if err != nil {
 			panic(err)
 		}
@@ -255,7 +254,7 @@ func relayPackets() {
 	log.Println("Listening for packets...")
 
 	for {
-		pkt := make([]byte, 1500)
+		pkt := make([]byte, 2000)
 		n, err := iface.Read([]byte(pkt))
 		if err != nil {
 			panic(err)

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)