WIP

peer/holepunch.go

@@ -0,0 +1,155 @@
+package peer
+
+import (
+	"log"
+	"net/netip"
+	"sync"
+	"time"
+
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+)
+
+const (
+	probeWait    = 15 * time.Second
+	backoffStart = 5 * time.Minute
+	backoffMax   = time.Hour
+)
+
+type probeState struct {
+	pubKey    wgtypes.Key
+	endpoint  netip.AddrPort
+	probing   bool
+	direct    bool
+	backoff   time.Duration
+	backoffAt time.Time
+}
+
+type HolePunch struct {
+	Globals
+	mu    sync.Mutex
+	peers [256]*probeState
+}
+
+func NewHolePunch(g Globals) *HolePunch {
+	return &HolePunch{Globals: g}
+}
+
+// OnEndpointLearned is called when a peer's external WG endpoint becomes known,
+// either from the VPN control channel (MsgMyEndpoint) or from the hub poller.
+// fromHub=true resets any existing backoff so the hub-reported change is acted
+// on immediately.
+func (hp *HolePunch) OnEndpointLearned(peerIP byte, pubKey wgtypes.Key, endpoint netip.AddrPort, fromHub bool) {
+	hp.mu.Lock()
+	defer hp.mu.Unlock()
+
+	ps := hp.peers[peerIP]
+	if ps == nil {
+		ps = &probeState{pubKey: pubKey}
+		hp.peers[peerIP] = ps
+	}
+
+	if fromHub {
+		ps.backoff = 0
+		ps.backoffAt = time.Time{}
+		ps.direct = false
+		ps.pubKey = pubKey
+	}
+	ps.endpoint = endpoint
+
+	if ps.probing || ps.direct {
+		return
+	}
+	if ps.backoff > 0 && time.Now().Before(ps.backoffAt) {
+		return
+	}
+
+	ps.probing = true
+	go hp.runProbe(peerIP)
+}
+
+func (hp *HolePunch) runProbe(peerIP byte) {
+	hp.mu.Lock()
+	ps := hp.peers[peerIP]
+	if ps == nil {
+		hp.mu.Unlock()
+		return
+	}
+	pubKey := ps.pubKey
+	endpoint := ps.endpoint
+	hp.mu.Unlock()
+
+	vpnIP := netip.AddrFrom4([4]byte{
+		hp.Network[0], hp.Network[1], hp.Network[2], peerIP,
+	})
+
+	probeStart := time.Now()
+	if err := addProbeEntry(hp.WGClient, hp.WGDevName, pubKey, endpoint); err != nil {
+		log.Printf("[HolePunch] addProbeEntry peer %d: %v", peerIP, err)
+		hp.finishProbe(peerIP, false)
+		return
+	}
+
+	time.Sleep(probeWait)
+
+	_, handshakeTime, err := getPeerEndpoint(hp.WGClient, hp.WGDevName, pubKey)
+	if err == nil && handshakeTime.After(probeStart) {
+		if err := promoteToDirect(hp.WGClient, hp.WGDevName, pubKey, vpnIP); err != nil {
+			log.Printf("[HolePunch] promoteToDirect peer %d: %v", peerIP, err)
+		}
+		hp.finishProbe(peerIP, true)
+		return
+	}
+
+	// Probe failed — remove entry and schedule backoff retry.
+	if err := removePeerEntry(hp.WGClient, hp.WGDevName, pubKey); err != nil {
+		log.Printf("[HolePunch] removePeerEntry peer %d: %v", peerIP, err)
+	}
+
+	hp.mu.Lock()
+	ps = hp.peers[peerIP]
+	var delay time.Duration
+	if ps != nil {
+		if ps.backoff == 0 {
+			ps.backoff = backoffStart
+		} else {
+			ps.backoff = min(ps.backoff*2, backoffMax)
+		}
+		ps.backoffAt = time.Now().Add(ps.backoff)
+		ps.probing = false
+		delay = ps.backoff
+	}
+	hp.mu.Unlock()
+
+	if delay > 0 {
+		go func() {
+			time.Sleep(delay)
+			hp.retryProbe(peerIP)
+		}()
+	}
+}
+
+func (hp *HolePunch) finishProbe(peerIP byte, success bool) {
+	hp.mu.Lock()
+	defer hp.mu.Unlock()
+	ps := hp.peers[peerIP]
+	if ps == nil {
+		return
+	}
+	ps.probing = false
+	if success {
+		ps.direct = true
+		ps.backoff = 0
+	}
+}
+
+func (hp *HolePunch) retryProbe(peerIP byte) {
+	hp.mu.Lock()
+	ps := hp.peers[peerIP]
+	if ps == nil || ps.probing || ps.direct {
+		hp.mu.Unlock()
+		return
+	}
+	ps.probing = true
+	hp.mu.Unlock()
+	hp.runProbe(peerIP)
+}