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 Peer struct {
	ifReader   *IFReader
	connReader *ConnReader
	iface      io.Writer
	hubPoller  *hubPoller
	super      *Super
}

type Config struct {
	NetName    string
	HubAddress string
	APIKey     string
}

func New(conf Config) *Peer {
	config, err := loadPeerConfig(conf.NetName)
	if err != nil {
		log.Printf("Failed to load configuration: %v", err)
		log.Printf("Initializing...")
		initPeerWithHub(conf)

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

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

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

	log.Printf("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 {
			log.Printf("Failed to write packet: %v", err)
		}
		writeLock.Unlock()
		return n, err
	}

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

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

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

	return &Peer{
		iface:      iface,
		ifReader:   ifReader,
		connReader: connReader,
		hubPoller:  hubPoller,
		super:      super,
	}
}

func (p *Peer) Run() {
	go p.ifReader.Run()
	go p.connReader.Run()
	p.super.Start()
	p.hubPoller.Run()
}

func initPeerWithHub(conf Config) {
	keys := generateKeys()

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

	args := m.PeerInitArgs{
		EncPubKey:  keys.PubKey,
		PubSignKey: keys.PubSignKey,
	}

	buf := &bytes.Buffer{}
	if err := json.NewEncoder(buf).Encode(args); 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("", conf.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)
	}

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

	peerConfig.PubKey = keys.PubKey
	peerConfig.PrivKey = keys.PrivKey
	peerConfig.PubSignKey = keys.PubSignKey
	peerConfig.PrivSignKey = keys.PrivSignKey

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

	log.Print("Initialization successful.")
}