CHORUS/discovery/mdns.go

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")
	}
}