158 lines
3.2 KiB
Go
158 lines
3.2 KiB
Go
//go:build ignore
|
|
|
|
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)
|
|
}
|