vppn/peer/peer.go

package peer

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

type peerMain struct {
	Globals
	ifReader   *IFReader
	connReader *ConnReader
	hubPoller  *HubPoller
}

type mainArgs struct {
	NetName    string
	HubAddress string
	APIKey     string
}

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

	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)
	}

	startupCount, err := loadStartupCount(args.NetName)
	if err != nil {
		if !os.IsNotExist(err) {
			log.Fatalf("Failed to load startup count: %v", err)
		}
	}

	if startupCount.Count == math.MaxUint16 {
		log.Fatalf("Startup counter overflow.")
	}
	startupCount.Count += 1

	if err := storeStartupCount(args.NetName, startupCount); err != nil {
		log.Fatalf("Failed to write startup count: %v", err)
	}

	iface, err := openInterface(config.Network, config.LocalPeerIP, args.NetName)
	if err != nil {
		log.Fatalf("Failed to open interface: %v", err)
	}

	localPeer := state.Peers[config.LocalPeerIP]

	log.Printf("Local peer %d: %v", config.LocalPeerIP, localPeer)
	myAddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf(":%d", localPeer.Port))
	if err != nil {
		log.Fatalf("Failed to resolve UDP address: %v", err)
	}

	logf("Listening on %v...", myAddr)
	conn, err := net.ListenUDP("udp", myAddr)
	if err != nil {
		log.Fatalf("Failed to open UDP port: %v", err)
	}

	conn.SetReadBuffer(1024 * 1024 * 8)
	conn.SetWriteBuffer(1024 * 1024 * 8)

	var localAddr netip.AddrPort
	ip, localAddrValid := netip.AddrFromSlice(localPeer.PublicIP)
	if localAddrValid {
		localAddr = netip.AddrPortFrom(ip, localPeer.Port)
	}

	g := NewGlobals(config, startupCount, localAddr, conn, iface)

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

	return &peerMain{
		Globals:    g,
		ifReader:   NewIFReader(g),
		connReader: NewConnReader(g, conn),
		hubPoller:  hubPoller,
	}
}

func (p *peerMain) Run() {
	for i := range p.RemotePeers {
		remote := p.RemotePeers[i].Load()
		go newRemoteFSM(remote).Run()
	}

	go p.ifReader.Run()
	go p.connReader.Run()

	if !p.LocalAddrValid {
		go RunMCWriter(p.LocalPeerIP, p.PrivSignKey)
		go RunMCReader(p.Globals)
	}

	go p.hubPoller.Run()

	select {}
}

func initPeerWithHub(args mainArgs) {
	keys := generateKeys()

	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{
		EncPubKey:  keys.PubKey,
		PubSignKey: keys.PubSignKey,
	}

	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)
	}

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

	config := LocalConfig{}
	config.LocalPeerIP = initResp.PeerIP
	config.Network = initResp.Network
	config.PubKey = keys.PubKey
	config.PrivKey = keys.PrivKey
	config.PubSignKey = keys.PubSignKey
	config.PrivSignKey = keys.PrivSignKey

	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.")
}