package discovery import ( "context" "fmt" "time" "chorus/internal/logging" "github.com/libp2p/go-libp2p/core/host" "github.com/libp2p/go-libp2p/core/peer" "github.com/libp2p/go-libp2p/p2p/discovery/mdns" "github.com/rs/zerolog" ) // MDNSDiscovery handles mDNS peer discovery for local network type MDNSDiscovery struct { host host.Host service mdns.Service notifee *mdnsNotifee ctx context.Context cancel context.CancelFunc serviceTag string logger zerolog.Logger } // mdnsNotifee handles discovered peers type mdnsNotifee struct { h host.Host ctx context.Context peersChan chan peer.AddrInfo logger zerolog.Logger } // NewMDNSDiscovery creates a new mDNS discovery service func NewMDNSDiscovery(ctx context.Context, h host.Host, serviceTag string) (*MDNSDiscovery, error) { if serviceTag == "" { serviceTag = "CHORUS-peer-discovery" } discoveryCtx, cancel := context.WithCancel(ctx) logger := logging.ForComponent(logging.ComponentP2P) // Create notifee to handle discovered peers notifee := &mdnsNotifee{ h: h, ctx: discoveryCtx, peersChan: make(chan peer.AddrInfo, 10), logger: logger, } // Create mDNS service service := mdns.NewMdnsService(h, serviceTag, notifee) discovery := &MDNSDiscovery{ host: h, service: service, notifee: notifee, ctx: discoveryCtx, cancel: cancel, serviceTag: serviceTag, logger: logger, } // Start the service if err := service.Start(); err != nil { cancel() return nil, fmt.Errorf("failed to start mDNS service: %w", err) } // Start background peer connection handler go discovery.handleDiscoveredPeers() logger.Info().Str("service_tag", serviceTag).Msg("mDNS Discovery started") return discovery, nil } // PeersChan returns a channel that receives discovered peers func (d *MDNSDiscovery) PeersChan() <-chan peer.AddrInfo { return d.notifee.peersChan } // handleDiscoveredPeers processes discovered peers and attempts connections func (d *MDNSDiscovery) handleDiscoveredPeers() { for { select { case <-d.ctx.Done(): return case peerInfo := <-d.notifee.peersChan: // Skip self if peerInfo.ID == d.host.ID() { continue } // Check if already connected if d.host.Network().Connectedness(peerInfo.ID) == 1 { // Connected continue } // Attempt to connect d.logger.Info().Str("peer_id", peerInfo.ID.ShortString()).Msg("Discovered peer, attempting connection") connectCtx, cancel := context.WithTimeout(d.ctx, 10*time.Second) if err := d.host.Connect(connectCtx, peerInfo); err != nil { d.logger.Warn().Err(err).Str("peer_id", peerInfo.ID.ShortString()).Msg("Failed to connect to peer") } else { d.logger.Info().Str("peer_id", peerInfo.ID.ShortString()).Msg("Successfully connected to peer") } cancel() } } } // Close shuts down the mDNS discovery service func (d *MDNSDiscovery) Close() error { d.cancel() close(d.notifee.peersChan) return d.service.Close() } // HandlePeerFound is called when a peer is discovered via mDNS func (n *mdnsNotifee) HandlePeerFound(pi peer.AddrInfo) { select { case <-n.ctx.Done(): return case n.peersChan <- pi: // Peer info sent to channel default: // Channel is full, skip this peer n.logger.Warn().Str("peer_id", pi.ID.ShortString()).Msg("Discovery channel full, skipping peer") } }