package discovery import ( "context" "fmt" "net" "strconv" "strings" "sync" "time" "go.uber.org/zap" ) const defaultBroadcastInterval = time.Second type Nodes map[string]struct{} type NewNodes <-chan string //goland:noinspection GoNameStartsWithPackageName type DiscoverySet []*Discovery func (d DiscoverySet) NewNodes() NewNodes { res := make(chan string, 20) for _, discovery := range d { go func(v *Discovery) { for node := range v.Nodes() { if len(res) == cap(res) { panic("discovery set nodes not reading!") } res <- node } }(discovery) } return res } func NewDiscoverySet(log *zap.Logger, discoverPort uint16, opts ...Option) (DiscoverySet, error) { iFaces, err := net.Interfaces() if err != nil { return nil, fmt.Errorf("list interfaces: %w", err) } set := make(DiscoverySet, 0, len(iFaces)) for _, iFace := range iFaces { if iFace.Flags&net.FlagLoopback == net.FlagLoopback { continue } if iFace.Flags&net.FlagRunning != net.FlagRunning { continue } discover, err := NewDiscovery(iFace, log, discoverPort, opts...) if err != nil { return nil, fmt.Errorf("new discover for %s: %w", iFace.Name, err) } set = append(set, discover) } return set, nil } func NewDiscovery(iFace net.Interface, log *zap.Logger, discoverPort uint16, opts ...Option) (*Discovery, error) { addrs, err := iFace.Addrs() if err != nil { return nil, fmt.Errorf("get interface address: %w", err) } broadcast := net.IP(make([]byte, 4)) var ownAddr net.IP for _, addr := range addrs { if ipnet, ok := addr.(*net.IPNet); ok { ip4 := ipnet.IP.To4() if ip4 == nil { continue } for i := range ip4 { broadcast[i] = ip4[i] | ^ipnet.Mask[i] } ownAddr = ipnet.IP break } } if broadcast.To4() == nil { return nil, fmt.Errorf("no broadcast address") } portString := strconv.Itoa(int(discoverPort)) conn, err := net.ListenPacket("udp4", broadcast.String()+":"+portString) if err != nil { return nil, fmt.Errorf("listen packet: %w", err) } udpAddr, err := net.ResolveUDPAddr("udp4", broadcast.String()+":"+portString) if err != nil { return nil, fmt.Errorf("resolve udp address: %w", err) } d := Discovery{ log: log.With(zap.Stringer("broadcast_address", broadcast)), nodes: make(Nodes), newNode: make(chan string, 20), ownAddr: ownAddr.To4(), conn: conn, broadcastAddr: udpAddr, broadcastInterval: defaultBroadcastInterval, } for _, opt := range opts { opt(&d) } return &d, nil } type Discovery struct { log *zap.Logger debug bool mu sync.Mutex nodes Nodes newNode chan string ownAddr net.IP conn net.PacketConn broadcastAddr *net.UDPAddr broadcastInterval time.Duration } func (d *Discovery) Start(ctx context.Context, wg *sync.WaitGroup) { defer wg.Done() d.log.Info("start discovery", zap.Stringer("address", &d.ownAddr)) timer := time.NewTicker(d.broadcastInterval) defer timer.Stop() listenStop := make(chan struct{}) go func() { if err := d.listen(); err != nil { d.log.Error("listen failed", zap.Error(err)) } close(listenStop) }() for { select { case <-ctx.Done(): if err := d.conn.Close(); err != nil { d.log.Warn("close connection", zap.Error(err)) } close(d.newNode) return case <-listenStop: d.log.Error("listener stopped, stop discovery") return case <-timer.C: d.broadcast() } } } func (d *Discovery) listen() error { buf := make([]byte, 4) for { n, addr, err := d.conn.ReadFrom(buf) if err != nil { if strings.Contains(err.Error(), "use of closed network connection") { d.log.Warn("listen connection closed") return nil } if opErr, ok := err.(*net.OpError); ok && !opErr.Temporary() { continue } return fmt.Errorf("read from: %w", err) } nodeAddr := net.IP(buf[:n]) d.log.Debug("received node address", zap.Stringer("address", nodeAddr)) clientIP, _, _ := strings.Cut(addr.String(), ":") if nodeAddr.String() != clientIP { d.log.Warn("received addr mismatch", zap.Stringer("received", nodeAddr), zap.String("detected", clientIP)) } d.addNode(nodeAddr.String()) } } func (d *Discovery) broadcast() { d.log.Debug("broadcast") if _, err := d.conn.WriteTo(d.ownAddr, d.broadcastAddr); err != nil { d.log.Error("write broadcast message", zap.Error(err)) } } func (d *Discovery) addNode(addr string) { if !d.debug && addr == d.ownAddr.String() { return } d.mu.Lock() if _, ok := d.nodes[addr]; !ok { d.log.Info("new node address", zap.String("address", addr)) d.nodes[addr] = struct{}{} d.newNode <- addr } d.mu.Unlock() } func (d *Discovery) Nodes() NewNodes { return d.newNode }