vppn/peer/peer.go

package peer

import (
	"bytes"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"net/netip"
	"net/url"
	"os"
	"vppn/m"

	"git.crumpington.com/lib/go/flock"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
)

type peerMain struct {
	Globals
	netName          string
	holePunch        *HolePunch
	controlServer    *ControlServer
	endpointReporter *EndpointReporter // non-nil on relay peers only
	hubPoller        *HubPoller
	lockFile         *os.File
}

func newPeerMain(args mainArgs) *peerMain {
	logf := func(s string, args ...any) {
		log.Printf("[Main] "+s, args...)
	}

	if err := os.MkdirAll(configDir(args.NetName), 0700); err != nil {
		log.Fatalf("Failed to create config directory: %v", err)
	}

	lockFile, err := flock.TryLock(lockFilePath(args.NetName))
	if err != nil {
		log.Fatalf("Failed to open lock file: %v", err)
	}
	if lockFile == nil {
		log.Fatalf("Failed to obtain file lock.")
	}

	config, err := loadPeerConfig(args.NetName)
	if err != nil {
		logf("Failed to load configuration: %v", err)
		logf("Initializing...")
		initPeerWithHub(args)

		config, err = loadPeerConfig(args.NetName)
		if err != nil {
			log.Fatalf("Failed to load configuration: %v", err)
		}
	}

	state, err := loadNetworkState(args.NetName)
	if err != nil {
		log.Fatalf("Failed to load network state: %v", err)
	}

	wgPrivKey, err := wgtypes.ParseKey(config.WGPrivKey)
	if err != nil {
		log.Fatalf("Failed to parse WireGuard private key: %v", err)
	}

	localPeer := state.Peers[config.LocalPeerIP]
	var listenPort int
	if localPeer != nil {
		listenPort = int(localPeer.Port1)
	}

	vpnIP := netip.AddrFrom4([4]byte{
		config.Network[0],
		config.Network[1],
		config.Network[2],
		config.LocalPeerIP,
	})

	wgClient, err := createWGDevice(args.NetName, wgPrivKey, listenPort, vpnIP, config.Network)
	if err != nil {
		log.Fatalf("Failed to create WireGuard device: %v", err)
	}

	for _, p := range state.Peers {
		if p == nil || !p.Relay || p.PeerIP == config.LocalPeerIP {
			continue
		}
		if len(p.WGPubKey) != wgtypes.KeyLen || len(p.PublicIP1) == 0 || p.Port1 == 0 {
			continue
		}
		relayPubKey, err := wgtypes.NewKey(p.WGPubKey)
		if err != nil {
			logf("Invalid relay WG key: %v", err)
			continue
		}
		relayIP, ok := netip.AddrFromSlice(p.PublicIP1)
		if !ok {
			continue
		}
		relayEndpoint := netip.AddrPortFrom(relayIP.Unmap(), p.Port1)
		if err := applyBaseConfig(wgClient, args.NetName, relayPubKey, relayEndpoint, config.Network); err != nil {
			logf("Failed to apply relay base config: %v", err)
		}
		break
	}

	g := NewGlobals(config, netip.AddrPort{})
	g.WGPrivKey = wgPrivKey
	g.WGClient = wgClient
	g.WGDevName = args.NetName

	holePunch := NewHolePunch(g)

	controlServer, err := NewControlServer(g, holePunch, args.NetName)
	if err != nil {
		log.Fatalf("Failed to create control server: %v", err)
	}

	var endpointReporter *EndpointReporter
	if localPeer != nil && localPeer.Relay {
		if err := enableForwarding(args.NetName); err != nil {
			log.Fatalf("Failed to enable IP forwarding: %v", err)
		}
		endpointReporter = NewEndpointReporter(g, controlServer, args.NetName)
	}

	hubPoller, err := NewHubPoller(g, holePunch, args.NetName, args.HubAddress, args.APIKey)
	if err != nil {
		log.Fatalf("Failed to create hub poller: %v", err)
	}

	go runStatusServer(g, statusSocketPath(args.NetName))

	return &peerMain{
		Globals:          g,
		netName:          args.NetName,
		holePunch:        holePunch,
		controlServer:    controlServer,
		endpointReporter: endpointReporter,
		hubPoller:        hubPoller,
		lockFile:         lockFile,
	}
}

func (p *peerMain) Run() {
	go p.controlServer.Run()
	if p.endpointReporter != nil {
		go p.endpointReporter.Run()
	}
	go RunMCWriter(p.Globals)
	go RunMCReader(p.Globals, p.holePunch, p.netName)
	go p.hubPoller.Run()
	select {}
}

func initPeerWithHub(args mainArgs) {
	privKey := generateWGKey()
	pubKey := privKey.PublicKey()

	initURL, err := url.Parse(args.HubAddress)
	if err != nil {
		log.Fatalf("Failed to parse hub URL: %v", err)
	}
	initURL.Path = "/peer/init/"

	initArgs := m.PeerInitArgs{
		WGPubKey: pubKey[:],
	}

	buf := &bytes.Buffer{}
	if err := json.NewEncoder(buf).Encode(initArgs); err != nil {
		log.Fatalf("Failed to encode init args: %v", err)
	}

	req, err := http.NewRequest(http.MethodPost, initURL.String(), buf)
	if err != nil {
		log.Fatalf("Failed to construct request: %v", err)
	}
	req.SetBasicAuth("", args.APIKey)

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		log.Fatalf("Failed to init with hub: %v", err)
	}
	defer resp.Body.Close()

	data, err := io.ReadAll(resp.Body)
	if err != nil {
		log.Fatalf("Failed to read response body: %v", err)
	}

	if resp.StatusCode == http.StatusConflict {
		log.Fatalf("WireGuard key already registered (HTTP 409). Delete and re-create the peer to re-register.")
	}
	if resp.StatusCode != http.StatusOK {
		log.Fatalf("Hub returned unexpected status %d: %s", resp.StatusCode, data)
	}

	initResp := m.PeerInitResp{}
	if err := json.Unmarshal(data, &initResp); err != nil {
		log.Fatalf("Failed to parse configuration: %v\n%s", err, data)
	}

	config := LocalConfig{
		LocalPeerIP: initResp.PeerIP,
		Network:     initResp.Network,
		WGPrivKey:   privKey.String(),
	}

	if err := storeNetworkState(args.NetName, initResp.NetworkState); err != nil {
		log.Fatalf("Failed to store network state: %v", err)
	}

	if err := storePeerConfig(args.NetName, config); err != nil {
		log.Fatalf("Failed to store configuration: %v", err)
	}

	log.Print("Initialization successful.")
}