vppn/peer/peer.go

package peer

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

type peerMain struct {
	conf       localConfig
	rt         *atomic.Pointer[routingTable]
	ifReader   *ifReader
	connReader *connReader
	iface      io.Writer
	hubPoller  *hubPoller
	super      *supervisor
}

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

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

	localPeer := state.Peers[config.PeerIP]

	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)

	// Wrap write function - this is necessary to avoid starvation.
	writeLock := sync.Mutex{}
	writeToUDPAddrPort := func(b []byte, addr netip.AddrPort) (n int, err error) {
		writeLock.Lock()
		n, err = conn.WriteToUDPAddrPort(b, addr)
		if err != nil {
			logf("Failed to write packet: %v", err)
		}
		writeLock.Unlock()
		return n, err
	}

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

	rt := newRoutingTable(localPeer.PeerIP, localAddr)
	rtPtr := &atomic.Pointer[routingTable]{}
	rtPtr.Store(&rt)

	ifReader := newIFReader(iface, writeToUDPAddrPort, rtPtr)
	super := newSupervisor(writeToUDPAddrPort, rtPtr, config.PrivKey)
	connReader := newConnReader(conn.ReadFromUDPAddrPort, writeToUDPAddrPort, iface, super.HandleControlMsg, rtPtr)
	hubPoller, err := newHubPoller(config.PeerIP, args.NetName, args.HubAddress, args.APIKey, super.HandleControlMsg)
	if err != nil {
		log.Fatalf("Failed to create hub poller: %v", err)
	}

	return &peerMain{
		conf:       config,
		rt:         rtPtr,
		iface:      iface,
		ifReader:   ifReader,
		connReader: connReader,
		hubPoller:  hubPoller,
		super:      super,
	}
}

func (p *peerMain) Run() {

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

	if !p.rt.Load().LocalAddr.IsValid() {
		go runMCWriter(p.conf.PeerIP, p.conf.PrivSignKey)
		go runMCReader(p.rt, p.super.HandleControlMsg)
	}

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