zeromesh/cmd/testclient/main.go

package main

import (
	"flag"
	"fmt"
	"log"
	"log/slog"
	"net"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"zeromesh/sdk"
)

var (
	serverURL    = flag.String("server", "http://server:10001", "Controller HTTP URL")
	name         = flag.String("name", "client", "Node name")
	udpPort      = flag.Int("udp", 7001, "Agent UDP listen port")
	tapName      = flag.String("tap", "ztmesh0", "TAP device name")
	identityPath = flag.String("identity", "", "Identity file path")
	adminUser    = flag.String("admin", "admin", "Admin username")
	adminPass    = flag.String("pass", "admin123456", "Admin password")
)

func main() {
	flag.Parse()
	slog.SetLogLoggerLevel(slog.LevelInfo)

	cfg := sdk.DefaultConfig()
	cfg.ControllerURL = *serverURL
	cfg.ListenPort = *udpPort
	cfg.TapName = *tapName
	cfg.PSK = "default-psk-change-me"
	if *identityPath != "" {
		cfg.IdentityPath = *identityPath
	}
	cli := sdk.New(cfg)

	var token string
	for i := 0; i < 30; i++ {
		if err := cli.Login(*adminUser, *adminPass); err == nil {
			token = cli.Token()
			break
		}
		if _, err := cli.InitAdmin(*adminUser, *adminPass); err == nil {
			token = cli.Token()
			log.Print("initialized admin")
			break
		}
		log.Printf("waiting for server (attempt %d)...", i+1)
		time.Sleep(2 * time.Second)
	}
	if token == "" {
		log.Fatal("could not login or init admin after 30 attempts")
	}
	log.Printf("logged in as %s", *adminUser)

	idPath := *identityPath
	if idPath == "" {
		idPath = fmt.Sprintf("/tmp/%s.id", *name)
	}
	id, err := sdk.LoadOrGenerateIdentity(idPath)
	if err != nil {
		log.Fatalf("identity: %v", err)
	}
	log.Printf("identity: %s pub=%s", id.Address, id.PublicKeyHex())

	myIP := sdk.GetPreferredIP()
	node, err := cli.RegisterNode(id.Address.String(), id.PublicKeyHex(), *name, myIP, *udpPort, "1.0.0")
	if err != nil {
		log.Fatalf("register node: %v", err)
	}
	log.Printf("registered node: %s", node.NodeID)

	findNetwork := func() *sdk.Network {
		networks, err := cli.ListNetworks()
		if err != nil {
			return nil
		}
		for i := range networks {
			if !strings.HasPrefix(networks[i].IPRange, "10.7.") && !strings.HasPrefix(networks[i].IPRange, "172.") && !strings.HasPrefix(networks[i].IPRange, "192.") {
				continue
			}
			return &networks[i]
		}
		if len(networks) > 0 {
			return &networks[0]
		}
		return nil
	}

	myNet := findNetwork()
	if myNet == nil {
		created, err := cli.CreateNetwork(fmt.Sprintf("testnet-%s", *name), "")
		if err != nil {
			log.Fatalf("create network: %v", err)
		}
		myNet = created
		log.Printf("created network: %s (%s) id=%d", myNet.Name, myNet.IPRange, myNet.NetworkID)
	} else {
		log.Printf("using network: %s (%s) id=%d", myNet.Name, myNet.IPRange, myNet.NetworkID)
	}

	var myIPAddr string
	for i := 0; i < 10; i++ {
		members, err := cli.NetworkMembers(myNet.NetworkID)
		if err == nil {
			for _, m := range members {
				if m.NodeID == id.Address.String() {
					myIPAddr = m.IPAddress
					break
				}
			}
		}
		if myIPAddr != "" {
			break
		}
		log.Printf("authorizing self on network %d...", myNet.NetworkID)
		cli.AuthorizeMember(myNet.NetworkID, id.Address.String())
		time.Sleep(1 * time.Second)
	}
	if myIPAddr == "" {
		log.Fatalf("could not get virtual IP after authorization")
	}
	log.Printf("my virtual IP: %s", myIPAddr)

	agent := sdk.NewAgent(cli)
	agent.SetIdentity(id)

	if err := agent.Start(); err != nil {
		log.Fatalf("agent start: %v", err)
	}
	defer agent.Stop()
	log.Printf("agent started: tap=%s udp=%d", agent.TAPName(), agent.LocalPort())

	agent.ControllerHello()
	time.Sleep(2 * time.Second)

	for i := 0; i < 6; i++ {
		agent.SyncPeers()
		time.Sleep(3 * time.Second)
		c := agent.PeerCount()
		if c > 0 {
			break
		}
	}

	cidrFull := myIPAddr + "/24"
	if err := agent.SetTAPIP(cidrFull); err != nil {
		log.Printf("set tap ip warning: %v", err)
	} else {
		log.Printf("tap ip set: %s", cidrFull)
	}

	listener, err := net.Listen("tcp", fmt.Sprintf("%s:7777", myIPAddr))
	if err != nil {
		log.Printf("echo server listen error (expected if TAP not ready): %v", err)
	} else {
		defer listener.Close()
		log.Printf("echo server listening on %s:7777", myIPAddr)
		go func() {
			for {
				conn, err := listener.Accept()
				if err != nil {
					return
				}
				go func(c net.Conn) {
					defer c.Close()
					buf := make([]byte, 4096)
					for {
						n, err := c.Read(buf)
						if err != nil {
							return
						}
						c.Write(buf[:n])
					}
				}(conn)
			}
		}()
	}

	peerCount := agent.PeerCount()
	log.Printf("%s ready: ip=%s tap=%s peers=%d", *name, myIPAddr, agent.TAPName(), peerCount)

	sc := make(chan os.Signal, 1)
	signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM)
	<-sc
	log.Printf("shutting down %s", *name)
}