vppn/node/main.go

package node

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/netip"
	"os"
	"runtime/debug"
	"sync/atomic"
	"vppn/m"
)

func panicHandler() {
	if r := recover(); r != nil {
		log.Fatalf("\n %v\n\nstacktrace from panic: %s\n", r, string(debug.Stack()))
	}
}

func Main() {
	defer panicHandler()

	var (
		initURL  string
		listenIP string
		port     int
	)

	flag.StringVar(&netName, "name", "", "[REQUIRED] The network name.")
	flag.StringVar(&initURL, "init-url", "", "Initializes peer from the hub URL.")
	flag.StringVar(&listenIP, "listen-ip", "", "IP address to listen on.")
	flag.IntVar(&port, "port", 0, "Port to listen on.")
	flag.Parse()

	if netName == "" {
		flag.Usage()
		os.Exit(1)
	}

	if initURL != "" {
		mainInit(initURL)
		return
	}

	main(listenIP, uint16(port))
}

func mainInit(initURL string) {
	if _, err := loadPeerConfig(netName); err == nil {
		log.Fatalf("Network is already initialized.")
	}

	resp, err := http.Get(initURL)
	if err != nil {
		log.Fatalf("Failed to fetch data from 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 := m.PeerConfig{}
	if err := json.Unmarshal(data, &peerConfig); err != nil {
		log.Fatalf("Failed to parse configuration: %v", err)
	}

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

	log.Print("Initialization successful.")
}

// ----------------------------------------------------------------------------

func main(listenIP string, port uint16) {
	config, err := loadPeerConfig(netName)
	if err != nil {
		log.Fatalf("Failed to load configuration: %v", err)
	}

	port = determinePort(config.Port, port)

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

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

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

	// Intialize globals.
	_iface = newIFWriter(iface)
	_conn = newConnWriter(conn)

	localIP = config.PeerIP
	discoveryPackets = make(chan controlPacket, 256)
	localAddr = &atomic.Pointer[netip.AddrPort]{}
	relayIP = &atomic.Pointer[byte]{}

	ip, ok := netip.AddrFromSlice(config.PublicIP)
	if ok {
		localPub = true
		addr := netip.AddrPortFrom(ip, config.Port)
		localAddr.Store(&addr)
	}

	privateKey = config.PrivKey

	// Start supervisors.
	for i := range 256 {
		go newPeerSupervisor(i).Run()
	}

	if localPub {
		go addrDiscoveryServer()
	} else {
		go addrDiscoveryClient()
		go relayManager()
	}
	go newHubPoller(config).Run()
	go readFromConn(conn)
	readFromIFace(iface)
}

// ----------------------------------------------------------------------------

func determinePort(confPort, portFromCommandLine uint16) uint16 {
	if portFromCommandLine != 0 {
		return portFromCommandLine
	}
	if confPort != 0 {
		return confPort
	}
	return 456
}

// ----------------------------------------------------------------------------

func readFromConn(conn *net.UDPConn) {

	defer panicHandler()

	var (
		remoteAddr netip.AddrPort
		n          int
		err        error
		buf        = make([]byte, bufferSize)
		decBuf     = make([]byte, bufferSize)
		data       []byte
		h          header
	)

	for {
		n, remoteAddr, err = conn.ReadFromUDPAddrPort(buf[:bufferSize])
		if err != nil {
			log.Fatalf("Failed to read from UDP port: %v", err)
		}

		remoteAddr = netip.AddrPortFrom(remoteAddr.Addr().Unmap(), remoteAddr.Port())

		data = buf[:n]

		if n < headerSize {
			continue // Packet it soo short.
		}

		h.Parse(data)
		switch h.StreamID {
		case controlStreamID:
			handleControlPacket(remoteAddr, h, data, decBuf)

		case dataStreamID:
			handleDataPacket(h, data, decBuf)

		default:
			log.Printf("Unknown stream ID: %d", h.StreamID)
		}
	}
}

func handleControlPacket(addr netip.AddrPort, h header, data, decBuf []byte) {
	route := routingTable[h.SourceIP].Load()
	if route.ControlCipher == nil {
		//log.Printf("Not connected (control).")
		return
	}

	if h.DestIP != localIP {
		log.Printf("Incorrect destination IP on control packet: %d != %d", h.DestIP, localIP)
		return
	}

	out, ok := route.ControlCipher.Decrypt(data, decBuf)
	if !ok {
		//log.Printf("Failed to decrypt control packet.")
		return
	}

	if len(out) == 0 {
		//log.Printf("Empty control packet from: %d", h.SourceIP)
		return
	}

	if dupChecks[h.SourceIP].IsDup(h.Counter) {
		//log.Printf("[%03d] Duplicate control packet: %d", h.SourceIP, h.Counter)
		return
	}

	pkt := controlPacket{
		SrcIP:   h.SourceIP,
		SrcAddr: addr,
	}

	if err := pkt.ParsePayload(out); err != nil {
		log.Printf("Failed to parse control packet: %v", err)
		return
	}

	switch pkt.Payload.(type) {

	case addrDiscoveryPacket:
		select {
		case discoveryPackets <- pkt:
		default:
			log.Printf("Dropping discovery packet.")
		}

	default:
		select {
		case controlPackets[h.SourceIP] <- pkt:
		default:
			log.Printf("Dropping control packet.")
		}
	}

}

func handleDataPacket(h header, data []byte, decBuf []byte) {
	route := routingTable[h.SourceIP].Load()
	if !route.Up {
		//log.Printf("Not connected (recv).")
		return
	}

	dec, ok := route.DataCipher.Decrypt(data, decBuf)
	if !ok {
		log.Printf("Failed to decrypt data packet.")
		return
	}

	if dupChecks[h.SourceIP].IsDup(h.Counter) {
		//log.Printf("[%03d] Duplicate data packet: %d", h.SourceIP, h.Counter)
		return
	}

	if h.DestIP == localIP {
		_iface.Write(dec)
		return
	}

	destRoute := routingTable[h.DestIP].Load()
	if !destRoute.Up {
		log.Printf("Not connected (relay): %v", destRoute)
		return
	}

	_conn.WriteTo(dec, destRoute.RemoteAddr)
}

// ----------------------------------------------------------------------------

func readFromIFace(iface io.ReadWriteCloser) {
	var (
		packet   = make([]byte, bufferSize)
		buf1     = make([]byte, bufferSize)
		buf2     = make([]byte, bufferSize)
		remoteIP byte
		err      error
	)

	for {
		packet, remoteIP, err = readNextPacket(iface, packet)
		if err != nil {
			log.Fatalf("Failed to read from interface: %v", err)
		}

		route := routingTable[remoteIP].Load()
		if !route.Up {
			log.Printf("Route not connected: %d", remoteIP)
			continue
		}

		_sendDataPacket(route, packet, buf1, buf2)
	}
}