toni/baalvpn
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

thread safe map

Browse Source
This commit is contained in:
Antoni Piasecki
2026-07-22 14:28:11 +02:00
parent c9c4cc0af1
commit a2d647f3af
3 changed files with 132 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

99
client/main.go
View File

@@ -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
peerPubKeys = map[string][]byte{}
// TODO: this should probably have a mutex
peerSessionKeys = map[string][]byte{}
// TODO: this should probably have a mutex
peerEstablishAcks = map[string]bool{}
establishMutex sync.Mutex
peerPubKeys = shared.NewTMap[string, []byte]()
peerSessionKeys = shared.NewTMap[string, []byte]()
peerEstablishAcks = shared.NewTMap[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 {
// TODO: eww
time.Sleep(50 * time.Millisecond)
if peerPubKey, ok := peerPubKeys[ip]; ok {
log.Println("got pubkey of " + ip)
establishMutex.Lock()
defer establishMutex.Unlock()
ciphertext, sharedSecret, err := xwing.Encapsulate(peerPubKey)
if err != nil {
panic(err)
}
// check again after the other goroutine finished
if key, ok := peerSessionKeys.GetOK(ip); ok {
return key
}
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)
}
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 {
// TODO: eww
time.Sleep(50 * time.Millisecond)
if peerEstablishAcks[ip] {
delete(peerEstablishAcks, ip)
break
}
}
for {
// TODO: eww
time.Sleep(50 * time.Millisecond)
if peerPubKey, ok := peerPubKeys.GetOK(ip); ok {
log.Println("got pubkey of " + ip)
sessionKey := sha256.Sum256(sharedSecret)
peerSessionKeys[ip] = sessionKey[:]
return sessionKey[:]
ciphertext, sharedSecret, err := xwing.Encapsulate(peerPubKey)
if err != nil {
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")
}