vppn/peer/hub_poller.go

package peer

import (
	"encoding/json"
	"io"
	"log"
	"net/http"
	"net/netip"
	"net/url"
	"time"

	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

	"vppn/m"
)

const hubPollInterval = 64 * time.Second

type HubPoller struct {
	selfVPNIP netip.Addr
	vpnNet    netip.Prefix
	hubURL    string
	apiKey    string
	statePath string // where the network state cache is persisted
	addCh     chan<- HubPeer
	removeCh  chan<- wgtypes.Key
	known     map[wgtypes.Key]int64 // pubKey → last seen version
}

func NewHubPoller(
	selfVPNIP netip.Addr,
	vpnNet netip.Prefix,
	hubURL, apiKey string,
	statePath string,
	addCh chan<- HubPeer,
	removeCh chan<- wgtypes.Key,
) (*HubPoller, error) {
	u, err := url.Parse(hubURL)
	if err != nil {
		return nil, err
	}
	u.Path = "/peer/fetch-state/"

	return &HubPoller{
		selfVPNIP: selfVPNIP,
		vpnNet:    vpnNet,
		hubURL:    u.String(),
		apiKey:    apiKey,
		statePath: statePath,
		addCh:     addCh,
		removeCh:  removeCh,
		known:     make(map[wgtypes.Key]int64),
	}, nil
}

func (hp *HubPoller) Run() {
	// Prime from the on-disk cache before reaching the hub, so the peer
	// configures WireGuard from its last known state even if the hub is down.
	// known starts empty, so this emits every cached peer as an add and seeds
	// the version map; the first real poll then emits only deltas.
	if state, err := loadNetworkState(hp.statePath); err == nil {
		hp.apply(state)
	}

	hp.poll()
	for range time.Tick(hubPollInterval) {
		hp.poll()
	}
}

func (hp *HubPoller) poll() {
	req, err := http.NewRequest(http.MethodGet, hp.hubURL, nil)
	if err != nil {
		log.Printf("[HubPoller] build request: %v", err)
		return
	}
	req.SetBasicAuth("", hp.apiKey)

	client := &http.Client{Timeout: 32 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		log.Printf("[HubPoller] fetch: %v", err)
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		log.Printf("[HubPoller] unexpected status %d", resp.StatusCode)
		return
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		log.Printf("[HubPoller] read body: %v", err)
		return
	}

	var state m.NetworkState
	if err := json.Unmarshal(body, &state); err != nil {
		log.Printf("[HubPoller] unmarshal: %v", err)
		return
	}

	// Persist only when the state actually changed, to avoid needless writes
	// on every poll.
	if hp.apply(state) {
		if err := saveNetworkState(hp.statePath, state); err != nil {
			log.Printf("[HubPoller] save state: %v", err)
		}
	}
}

// apply diffs state against the known versions, emitting add events for new or
// changed peers and remove events for peers that disappeared. It returns true
// if anything changed.
func (hp *HubPoller) apply(state m.NetworkState) (changed bool) {
	seen := make(map[wgtypes.Key]struct{}, len(hp.known))

	netAddr := hp.vpnNet.Addr().As4()

	for _, p := range state.Peers {
		if p == nil || len(p.WGPubKey) != wgtypes.KeyLen || len(p.SignPubKey) != 32 {
			continue
		}

		pubKey, err := wgtypes.NewKey(p.WGPubKey)
		if err != nil {
			continue
		}

		octets := netAddr
		octets[3] = p.PeerIP
		vpnIP := netip.AddrFrom4(octets)
		if vpnIP == hp.selfVPNIP {
			continue
		}

		seen[pubKey] = struct{}{}

		if v, ok := hp.known[pubKey]; ok && v == p.Version {
			continue
		}
		hp.known[pubKey] = p.Version
		hp.addCh <- hubPeerFrom(pubKey, vpnIP, p)
		changed = true
	}

	for key := range hp.known {
		if _, ok := seen[key]; !ok {
			delete(hp.known, key)
			hp.removeCh <- key
			changed = true
		}
	}

	return changed
}

func hubPeerFrom(pubKey wgtypes.Key, vpnIP netip.Addr, p *m.Peer) HubPeer {
	var ep4, ep6 netip.AddrPort
	if len(p.Addr4) > 0 {
		if addr, ok := netip.AddrFromSlice(p.Addr4); ok {
			ep4 = netip.AddrPortFrom(addr.Unmap(), p.Port)
		}
	}
	if len(p.Addr6) > 0 {
		if addr, ok := netip.AddrFromSlice(p.Addr6); ok {
			ep6 = netip.AddrPortFrom(addr, p.Port)
		}
	}
	var signPubKey [32]byte
	copy(signPubKey[:], p.SignPubKey)
	return HubPeer{
		PubKey:     pubKey,
		VPNIP:      vpnIP,
		Name:       p.Name,
		IsRelay:    p.Relay,
		IsPublic:   ep4.IsValid() || ep6.IsValid(),
		EndpointV4: ep4,
		EndpointV6: ep6,
		SignPubKey: signPubKey,
	}
}