WIP

peer/control.go

@@ -0,0 +1,150 @@
+package peer
+
+import (
+	"encoding/binary"
+	"fmt"
+	"log"
+	"net"
+	"net/netip"
+
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+)
+
+const (
+	controlPort    = uint16(4561)
+	controlMsgLen  = 7 // 1 type + 4 IPv4 + 2 port
+	msgYourEndpt   = uint8(1)
+	msgMyEndpt     = uint8(2)
+)
+
+type ControlServer struct {
+	localPeerIP byte
+	network     []byte
+	conn        *net.UDPConn
+	hp          *HolePunch
+	netName     string
+}
+
+func NewControlServer(g Globals, hp *HolePunch, netName string) (*ControlServer, error) {
+	vpnIP := netip.AddrFrom4([4]byte{
+		g.Network[0], g.Network[1], g.Network[2], g.LocalPeerIP,
+	})
+	listenAddr := net.UDPAddrFromAddrPort(netip.AddrPortFrom(vpnIP, controlPort))
+	conn, err := net.ListenUDP("udp4", listenAddr)
+	if err != nil {
+		return nil, fmt.Errorf("listen on control port: %w", err)
+	}
+	return &ControlServer{
+		localPeerIP: g.LocalPeerIP,
+		network:     g.Network,
+		conn:        conn,
+		hp:          hp,
+		netName:     netName,
+	}, nil
+}
+
+func (cs *ControlServer) Run() {
+	buf := make([]byte, 64)
+	for {
+		n, src, err := cs.conn.ReadFromUDPAddrPort(buf)
+		if err != nil {
+			log.Printf("[Control] read: %v", err)
+			continue
+		}
+		if n < controlMsgLen {
+			continue
+		}
+		cs.handle(buf[:n], src)
+	}
+}
+
+func (cs *ControlServer) handle(msg []byte, src netip.AddrPort) {
+	msgType, ep, ok := decodeControlMsg(msg)
+	if !ok {
+		return
+	}
+	switch msgType {
+	case msgYourEndpt:
+		cs.onYourEndpoint(ep)
+	case msgMyEndpt:
+		cs.onMyEndpoint(src, ep)
+	}
+}
+
+// onYourEndpoint is called when the relay tells us our external WG endpoint.
+// We broadcast MsgMyEndpoint to all known peers so they can attempt direct
+// connections to us.
+func (cs *ControlServer) onYourEndpoint(ourEndpoint netip.AddrPort) {
+	log.Printf("[Control] external endpoint: %v", ourEndpoint)
+
+	state, err := loadNetworkState(cs.netName)
+	if err != nil {
+		log.Printf("[Control] load state: %v", err)
+		return
+	}
+
+	for _, p := range state.Peers {
+		if p == nil || p.PeerIP == cs.localPeerIP || len(p.WGPubKey) != wgtypes.KeyLen {
+			continue
+		}
+		peerVPNIP := netip.AddrFrom4([4]byte{
+			cs.network[0], cs.network[1], cs.network[2], p.PeerIP,
+		})
+		cs.sendMsg(msgMyEndpt, peerVPNIP, ourEndpoint)
+	}
+}
+
+// onMyEndpoint is called when a peer tells us their external WG endpoint.
+// We start a probe-before-commit attempt to reach them directly.
+func (cs *ControlServer) onMyEndpoint(src netip.AddrPort, theirEndpoint netip.AddrPort) {
+	peerIPByte := src.Addr().As4()[3]
+
+	state, err := loadNetworkState(cs.netName)
+	if err != nil {
+		log.Printf("[Control] load state: %v", err)
+		return
+	}
+
+	peer := state.Peers[peerIPByte]
+	if peer == nil || len(peer.WGPubKey) != wgtypes.KeyLen {
+		return
+	}
+
+	pubKey, err := wgtypes.NewKey(peer.WGPubKey)
+	if err != nil {
+		return
+	}
+
+	cs.hp.OnEndpointLearned(peerIPByte, pubKey, theirEndpoint, false)
+}
+
+// SendYourEndpoint sends MsgYourEndpoint to a peer, informing them of their
+// own external WG endpoint. Used by the relay's endpoint reporter (Phase 6).
+func (cs *ControlServer) SendYourEndpoint(peerVPNIP netip.Addr, theirEndpoint netip.AddrPort) {
+	cs.sendMsg(msgYourEndpt, peerVPNIP, theirEndpoint)
+}
+
+func (cs *ControlServer) sendMsg(msgType uint8, peerVPNIP netip.Addr, ep netip.AddrPort) {
+	dst := net.UDPAddrFromAddrPort(netip.AddrPortFrom(peerVPNIP, controlPort))
+	if _, err := cs.conn.WriteTo(encodeControlMsg(msgType, ep), dst); err != nil {
+		log.Printf("[Control] send to %v: %v", peerVPNIP, err)
+	}
+}
+
+func encodeControlMsg(msgType uint8, ep netip.AddrPort) []byte {
+	msg := make([]byte, controlMsgLen)
+	msg[0] = msgType
+	a4 := ep.Addr().Unmap().As4()
+	copy(msg[1:5], a4[:])
+	binary.BigEndian.PutUint16(msg[5:7], ep.Port())
+	return msg
+}
+
+func decodeControlMsg(msg []byte) (msgType uint8, ep netip.AddrPort, ok bool) {
+	if len(msg) < controlMsgLen {
+		return 0, netip.AddrPort{}, false
+	}
+	ip := netip.AddrFrom4([4]byte(msg[1:5]))
+	port := binary.BigEndian.Uint16(msg[5:7])
+	return msg[0], netip.AddrPortFrom(ip, port), true
+}