vppn/node/main.go

package node

import (
	"bytes"
	"crypto/rand"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/netip"
	"net/url"
	"os"
	"runtime/debug"
	"time"
	"vppn/m"

	"golang.org/x/crypto/nacl/box"
	"golang.org/x/crypto/nacl/sign"
)

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 hubAddress string

	flag.StringVar(&netName, "name", "", "[REQUIRED] The network name.")
	flag.StringVar(&hubAddress, "hub-address", "", "[REQUIRED] The hub address.")
	flag.StringVar(&apiKey, "api-key", "", "[REQUIRED] The node's API key.")
	flag.Parse()

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

	var err error

	hubURL, err = url.Parse(hubAddress)
	if err != nil {
		log.Fatalf("Failed to parse hub address: %v", err)
	}

	main()
}

func initPeerWithHub() {
	encPubKey, encPrivKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		log.Fatalf("Failed to generate encryption keys: %v", err)
	}

	signPubKey, signPrivKey, err := sign.GenerateKey(rand.Reader)
	if err != nil {
		log.Fatalf("Failed to generate signing keys: %v", err)
	}

	initURL := *hubURL
	initURL.Path = "/peer/init/"

	args := m.PeerInitArgs{
		EncPubKey:  encPubKey[:],
		PubSignKey: signPubKey[:],
	}

	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("", 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 = encPubKey[:]
	peerConfig.PrivKey = encPrivKey[:]
	peerConfig.PubSignKey = signPubKey[:]
	peerConfig.PrivSignKey = signPrivKey[:]

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

	log.Print("Initialization successful.")
}

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

func main() {
	config, err := loadPeerConfig(netName)
	if err != nil {
		log.Printf("Failed to load configuration: %v", err)
		log.Printf("Initializing...")
		initPeerWithHub()

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

	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(":%d", config.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)
	}

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

	localIP = config.PeerIP

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

	privKey = config.PrivKey
	privSignKey = config.PrivSignKey

	if !localPub {
		go relayManager()
		go localDiscovery()
	}

	go func() {
		for range time.Tick(pingInterval) {
			messages <- pingTimerMsg{}
		}
	}()

	sender := newPacketSender(conn)

	go startPeerSuper(routingTable, messages, sender)

	go newHubPoller().Run()
	go readFromConn(conn, iface, sender)

	readFromIFace(iface, sender)
}

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

func readFromConn(conn *net.UDPConn, iface io.ReadWriteCloser, sender dataPacketSender) {

	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, iface, sender)

		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: %#v", h)
		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
	}

	msg, err := parseControlMsg(h.SourceIP, addr, out)
	if err != nil {
		log.Printf("Failed to parse control packet: %v", err)
		return
	}

	select {
	case messages <- msg:
	default:
		log.Printf("Dropping control packet.")
	}

}

func handleDataPacket(h header, data []byte, decBuf []byte, iface ifWriter, sender dataPacketSender) {
	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 {
		if _, err := iface.Write(dec); err != nil {
			log.Fatalf("Failed to write to interface: %v", err)
		}
		return
	}

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

	sender.SendEncryptedDataPacket(dec, destRoute.RemoteAddr)
}

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

func readFromIFace(iface io.ReadWriteCloser, sender dataPacketSender) {
	var (
		packet   = 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
		}

		sender.SendDataPacket(packet, *route)
	}
}