NetworkTool/server/internal/discovery/discovery.go

package discovery

import (
	"context"
	"encoding/json"
	"fmt"
	"net"
	"strings"

	"nettool/internal/config"
	"nettool/internal/deviceinfo"
	"nettool/internal/logger"
	"nettool/internal/model"
)

type Server struct {
	cfg       config.Config
	log       *logger.Logger
	deviceSvc *deviceinfo.Service
}

func New(cfg config.Config, log *logger.Logger, deviceSvc *deviceinfo.Service) *Server {
	return &Server{cfg: cfg, log: log, deviceSvc: deviceSvc}
}

func (s *Server) Run(ctx context.Context) error {
	addr, err := net.ResolveUDPAddr("udp4", fmt.Sprintf("%s:%d", s.cfg.UDPHost, s.cfg.UDPPort))
	if err != nil {
		return err
	}
	conn, err := net.ListenUDP("udp4", addr)
	if err != nil {
		return err
	}
	defer conn.Close()

	go func() {
		<-ctx.Done()
		_ = conn.Close()
	}()

	s.log.Info("udp discovery listening", "addr", conn.LocalAddr().String())
	if err := enableLocalAddrControl(conn); err != nil {
		s.log.Warn("udp discovery local interface detection is unavailable", "error", err.Error())
	}

	buf := make([]byte, 2048)
	for {
		n, remote, packetInfo, err := readFromUDP(conn, buf)
		if err != nil {
			if ctx.Err() != nil {
				return nil
			}
			return err
		}

		var req model.DiscoverRequest
		if err := json.Unmarshal(buf[:n], &req); err != nil || req.MessageType != "discover" {
			continue
		}

		lan2IP, mac := s.maintenanceEndpoint(packetInfo)
		if lan2IP == "" {
			s.log.Warn("skip discovery response because no 169.254 maintenance address was found")
			continue
		}

		device := s.deviceSvc.Get()
		resp := model.DiscoverResponse{
			ProtocolVersion: 1,
			MessageType:     "discover_response",
			RequestID:       req.RequestID,
			DeviceID:        device.DeviceID,
			Hostname:        device.Hostname,
			ServerVersion:   device.ServerVersion,
			MAC:             mac,
			LAN2IP:          lan2IP,
			HTTPPort:        s.cfg.HTTPPort,
			AuthRequired:    true,
		}
		payload, _ := json.Marshal(resp)
		responseSourceIP := findResponseSourceIP(packetInfo.ifIndex, remote.IP)
		responseTarget := remote
		broadcastResponse := false
		if responseSourceIP == nil && packetInfo.ifIndex > 0 {
			responseSourceIP = net.ParseIP(lan2IP)
			responseTarget = broadcastResponseTarget(remote)
			broadcastResponse = true
		}
		if _, err := writeToUDP(conn, payload, responseTarget, packetInfo, responseSourceIP); err != nil {
			s.log.Warn("failed to send udp discovery response", "remote", remote.String(), "target", responseTarget.String(), "source_ip", responseSourceIP.String(), "maintenance_ip", lan2IP, "interface_index", packetInfo.ifIndex, "broadcast", broadcastResponse, "error", err.Error())
		}
	}
}

func broadcastResponseTarget(remote *net.UDPAddr) *net.UDPAddr {
	return &net.UDPAddr{IP: net.IPv4bcast, Port: remote.Port}
}

type udpPacketInfo struct {
	localIP net.IP
	ifIndex int
}

func (s *Server) maintenanceEndpoint(packetInfo udpPacketInfo) (string, string) {
	if packetInfo.ifIndex > 0 {
		lan2IP, mac := findLinkLocalEndpointByInterfaceIndex(packetInfo.ifIndex)
		if lan2IP != "" {
			return lan2IP, mac
		}
		return "", ""
	}
	if packetInfo.localIP != nil {
		lan2IP, mac := findLinkLocalEndpointByInterfaceIP(packetInfo.localIP.String())
		if lan2IP != "" {
			return lan2IP, mac
		}
	}
	if s.cfg.MaintenanceIP != "" {
		mac := findMACByIP(s.cfg.MaintenanceIP)
		if mac != "" {
			return s.cfg.MaintenanceIP, mac
		}
	}
	return findFirstLinkLocalEndpoint()
}

func findResponseSourceIP(interfaceIndex int, remoteIP net.IP) net.IP {
	if interfaceIndex <= 0 {
		return nil
	}
	iface, err := net.InterfaceByIndex(interfaceIndex)
	if err != nil {
		return nil
	}
	addresses, err := iface.Addrs()
	if err != nil {
		return nil
	}
	return bestSourceIPForRemote(addresses, remoteIP)
}

func bestSourceIPForRemote(addresses []net.Addr, remoteIP net.IP) net.IP {
	remoteIPv4 := remoteIP.To4()
	if remoteIPv4 == nil {
		return nil
	}
	var best net.IP
	bestPrefix := -1
	for _, address := range addresses {
		ipNet, ok := address.(*net.IPNet)
		if !ok || ipNet.IP.To4() == nil || !ipNet.Contains(remoteIPv4) {
			continue
		}
		prefix, bits := ipNet.Mask.Size()
		if bits != 32 || prefix <= bestPrefix {
			continue
		}
		best = ipNet.IP.To4()
		bestPrefix = prefix
	}
	return best
}

func findLinkLocalEndpointByInterfaceIndex(index int) (string, string) {
	iface, err := net.InterfaceByIndex(index)
	if err != nil {
		return "", ""
	}
	return findLinkLocalEndpointOnInterface(*iface)
}

func findLinkLocalEndpointByInterfaceIP(ip string) (string, string) {
	ifaces, err := net.Interfaces()
	if err != nil {
		return "", ""
	}
	for _, iface := range ifaces {
		lan2IP, mac := findLinkLocalEndpointOnInterface(iface)
		if lan2IP == ip {
			return lan2IP, mac
		}
	}
	return "", ""
}

func findFirstLinkLocalEndpoint() (string, string) {
	ifaces, err := net.Interfaces()
	if err != nil {
		return "", ""
	}
	for _, iface := range ifaces {
		lan2IP, mac := findLinkLocalEndpointOnInterface(iface)
		if lan2IP != "" {
			return lan2IP, mac
		}
	}
	return "", ""
}

func findLinkLocalEndpointOnInterface(iface net.Interface) (string, string) {
	if iface.Flags&net.FlagLoopback != 0 || iface.Flags&net.FlagUp == 0 || len(iface.HardwareAddr) == 0 {
		return "", ""
	}
	addrs, err := iface.Addrs()
	if err != nil {
		return "", ""
	}
	for _, addr := range addrs {
		current := ipv4FromAddr(addr)
		if current == nil || !strings.HasPrefix(current.String(), "169.254.") {
			continue
		}
		return current.String(), iface.HardwareAddr.String()
	}
	return "", ""
}

func findMACByIP(ip string) string {
	ifaces, err := net.Interfaces()
	if err != nil {
		return ""
	}
	var fallback string
	for _, iface := range ifaces {
		if iface.Flags&net.FlagLoopback != 0 || len(iface.HardwareAddr) == 0 {
			continue
		}
		addrs, err := iface.Addrs()
		if err != nil {
			continue
		}
		for _, addr := range addrs {
			current := ipv4FromAddr(addr)
			if current == nil {
				continue
			}
			if current.String() == ip {
				return iface.HardwareAddr.String()
			}
			if fallback == "" && strings.HasPrefix(current.String(), "169.254.") {
				fallback = iface.HardwareAddr.String()
			}
		}
	}
	return fallback
}

func ipv4FromAddr(addr net.Addr) net.IP {
	var current net.IP
	switch value := addr.(type) {
	case *net.IPNet:
		current = value.IP
	case *net.IPAddr:
		current = value.IP
	}
	if current == nil {
		return nil
	}
	return current.To4()
}