PlugAI_Launcher/linux-server/broadcaster.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"net"
	"os/exec"
	"runtime"
	"strings"
	"time"
)

type BroadcastMessage struct {
	Type string `json:"type"`
	IP   string `json:"ip"`
	Port int    `json:"port"`
	Name string `json:"name"`
}

// 获取默认网关 IP
func getGatewayIP() (string, error) {
	var cmd *exec.Cmd
	if runtime.GOOS == "linux" {
		cmd = exec.Command("sh", "-c", "ip route | grep default | awk '{print $3}'")
	} else if runtime.GOOS == "darwin" {
		cmd = exec.Command("sh", "-c", "netstat -rn | grep default | awk '{print $2}'")
	} else {
		return "", fmt.Errorf("unsupported OS")
	}
	out, err := cmd.Output()
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(out)), nil
}

// 判断两个 IP 是否在同一子网内
func ipInSameSubnet(ip1 net.IP, ip2 net.IP, mask net.IPMask) bool {
	if mask == nil {
		return false
	}
	return ip1.Mask(mask).Equal(ip2.Mask(mask))
}

// 获取本地与网关同网段的 IP 地址
func getLocalIP() string {
	gatewayStr, err := getGatewayIP()
	if err != nil {
		return "127.0.0.1"
	}
	gateway := net.ParseIP(gatewayStr)
	if gateway == nil {
		return "127.0.0.1"
	}

	interfaces, err := net.Interfaces()
	if err != nil {
		return "127.0.0.1"
	}

	for _, iface := range interfaces {
		if iface.Flags&net.FlagUp == 0 {
			continue
		}
		addrs, err := iface.Addrs()
		if err != nil {
			continue
		}
		for _, addr := range addrs {
			var ip net.IP
			var mask net.IPMask
			switch v := addr.(type) {
			case *net.IPNet:
				ip = v.IP
				mask = v.Mask
			case *net.IPAddr:
				ip = v.IP
			}
			if ip == nil || ip.IsLoopback() || ip.To4() == nil {
				continue
			}
			if ipInSameSubnet(ip, gateway, mask) {
				return ip.String()
			}
		}
	}
	return "127.0.0.1"
}

// 判断 IP 是否有效（用于判断网络是否连通）
func isValidIP(ip string) bool {
	if ip == "" || ip == "127.0.0.1" {
		return false
	}
	parsed := net.ParseIP(ip)
	if parsed == nil || parsed.IsLoopback() || parsed.IsUnspecified() || parsed.IsLinkLocalUnicast() || parsed.IsLinkLocalMulticast() {
		return false
	}
	return true
}

func main() {
	// 命令行参数
	var (
		port     = flag.Int("port", 9876, "Device service port")
		interval = flag.Int("interval", 2, "Broadcast interval in seconds")
		name     = flag.String("name", "PlugAI ZeroStack AI Server", "T1")
	)
	flag.Parse()

	var (
		prevIP string
		conn   *net.UDPConn
	)

	for {
		currentIP := getLocalIP()

		if !isValidIP(currentIP) {
			fmt.Println("网络未就绪，等待中...")
			time.Sleep(time.Duration(*interval) * time.Second)
			continue
		}

		// 如果 IP 改变或连接尚未建立，则重建 UDP 连接
		if currentIP != prevIP || conn == nil {
			if conn != nil {
				conn.Close()
			}
			addr := net.UDPAddr{
				IP:   net.IPv4bcast,
				Port: 9876,
			}
			newConn, err := net.DialUDP("udp4", nil, &addr)
			if err != nil {
				fmt.Println("UDP 连接失败:", err)
				time.Sleep(time.Duration(*interval) * time.Second)
				continue
			}
			conn = newConn
			prevIP = currentIP
			fmt.Println("绑定新 IP 广播:", currentIP)
		}

		// 构建并发送广播消息
		msg := BroadcastMessage{
			Type: "ai_server_announce",
			IP:   currentIP,
			Port: *port,
			Name: *name,
		}
		data, _ := json.Marshal(msg)
		_, err := conn.Write(data)
		if err != nil {
			fmt.Println("广播失败:", err)
		} else {
			fmt.Printf("广播中: %s:%d (%s)\n", currentIP, *port, *name)
		}

		time.Sleep(time.Duration(*interval) * time.Second)
	}
}