oc-discovery/daemons/node/common/dht_discovery.go

package common

import (
	"context"
	"math/rand"
	"strings"
	"time"

	oclib "cloud.o-forge.io/core/oc-lib"
	"github.com/ipfs/go-cid"
	dht "github.com/libp2p/go-libp2p-kad-dht"
	"github.com/libp2p/go-libp2p/core/host"
	pp "github.com/libp2p/go-libp2p/core/peer"
	ma "github.com/multiformats/go-multiaddr"
	mh "github.com/multiformats/go-multihash"
)

// FilterLoopbackAddrs strips loopback (127.x, ::1) and unspecified addresses
// from an AddrInfo so we never hand peers an address they cannot dial externally.
func FilterLoopbackAddrs(ai pp.AddrInfo) pp.AddrInfo {
	filtered := make([]ma.Multiaddr, 0, len(ai.Addrs))
	for _, addr := range ai.Addrs {
		ip, err := ExtractIP(addr.String())
		if err != nil || ip.IsLoopback() || ip.IsUnspecified() {
			continue
		}
		filtered = append(filtered, addr)
	}
	return pp.AddrInfo{ID: ai.ID, Addrs: filtered}
}

// RecommendedHeartbeatInterval is the target period between heartbeat ticks.
// Indexers use this as the DHT Provide refresh interval.
const RecommendedHeartbeatInterval = 60 * time.Second

// discoveryDHT is the DHT instance used for indexer discovery.
// Set by SetDiscoveryDHT once the indexer service initialises its DHT.
var discoveryDHT *dht.IpfsDHT

// SetDiscoveryDHT stores the DHT instance used by replenishIndexersFromDHT.
// Called by NewIndexerService once the DHT is ready.
func SetDiscoveryDHT(d *dht.IpfsDHT) {
	discoveryDHT = d
}

// initNodeDHT creates a lightweight DHT client for pure nodes (no IndexerService).
// Uses the seed indexers as bootstrap peers. Called lazily by ConnectToIndexers
// when discoveryDHT is still nil after the initial warm-up delay.
func initNodeDHT(h host.Host, seeds []Entry) {
	logger := oclib.GetLogger()
	bootstrapPeers := []pp.AddrInfo{}
	for _, s := range seeds {
		bootstrapPeers = append(bootstrapPeers, *s.Info)
	}
	d, err := dht.New(context.Background(), h,
		dht.Mode(dht.ModeClient),
		dht.ProtocolPrefix("oc"),
		dht.BootstrapPeers(bootstrapPeers...),
	)
	if err != nil {
		logger.Warn().Err(err).Msg("[dht] node DHT client init failed")
		return
	}
	SetDiscoveryDHT(d)
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	if err := d.Bootstrap(ctx); err != nil {
		logger.Warn().Err(err).Msg("[dht] node DHT client bootstrap failed")
	}
	logger.Info().Msg("[dht] node DHT client ready")
}

// IndexerCID returns the well-known CID under which all indexers advertise.
func IndexerCID() cid.Cid {
	h, _ := mh.Sum([]byte("/opencloud/indexers"), mh.SHA2_256, -1)
	return cid.NewCidV1(cid.Raw, h)
}

// DiscoverIndexersFromDHT uses the DHT to find up to count indexers advertising
// under the well-known key. Excludes self. Resolves addresses when the provider
// record carries none.
func DiscoverIndexersFromDHT(h host.Host, d *dht.IpfsDHT, count int) []pp.AddrInfo {
	logger := oclib.GetLogger()
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	c := IndexerCID()
	ch := d.FindProvidersAsync(ctx, c, count*2)
	seen := map[pp.ID]struct{}{}
	var results []pp.AddrInfo
	for ai := range ch {
		if ai.ID == h.ID() {
			continue
		}
		if _, dup := seen[ai.ID]; dup {
			continue
		}
		seen[ai.ID] = struct{}{}
		if len(ai.Addrs) == 0 {
			resolved, err := d.FindPeer(ctx, ai.ID)
			if err != nil {
				logger.Warn().Str("peer", ai.ID.String()).Msg("[dht] no addrs and FindPeer failed, skipping")
				continue
			}
			ai = resolved
		}
		ai = FilterLoopbackAddrs(ai)
		if len(ai.Addrs) == 0 {
			continue
		}
		results = append(results, ai)
		if len(results) >= count {
			break
		}
	}
	logger.Info().Int("found", len(results)).Msg("[dht] indexer discovery complete")
	return results
}

// SelectByFillRate picks up to want providers using fill-rate weighted random
// selection w(F) = F*(1-F) — peaks at F=0.5, prefers less-loaded indexers.
// Providers with unknown fill rate receive F=0.5 (neutral prior).
// Enforces subnet /24 diversity: at most one indexer per /24.
func SelectByFillRate(providers []pp.AddrInfo, fillRates map[pp.ID]float64, want int) []pp.AddrInfo {
	if len(providers) == 0 || want <= 0 {
		return nil
	}
	type weighted struct {
		ai     pp.AddrInfo
		weight float64
	}
	ws := make([]weighted, 0, len(providers))
	for _, ai := range providers {
		f, ok := fillRates[ai.ID]
		if !ok {
			f = 0.5
		}
		ws = append(ws, weighted{ai: ai, weight: f * (1 - f)})
	}
	// Shuffle first for fairness among equal-weight peers.
	rand.Shuffle(len(ws), func(i, j int) { ws[i], ws[j] = ws[j], ws[i] })
	// Sort descending by weight (simple insertion sort — small N).
	for i := 1; i < len(ws); i++ {
		for j := i; j > 0 && ws[j].weight > ws[j-1].weight; j-- {
			ws[j], ws[j-1] = ws[j-1], ws[j]
		}
	}

	subnets := map[string]struct{}{}
	var selected []pp.AddrInfo
	for _, w := range ws {
		if len(selected) >= want {
			break
		}
		subnet := subnetOf(w.ai)
		if subnet != "" {
			if _, dup := subnets[subnet]; dup {
				continue
			}
			subnets[subnet] = struct{}{}
		}
		selected = append(selected, w.ai)
	}
	return selected
}

// subnetOf returns the /24 subnet string for the first non-loopback address of ai.
func subnetOf(ai pp.AddrInfo) string {
	for _, ma := range ai.Addrs {
		ip, err := ExtractIP(ma.String())
		if err != nil || ip.IsLoopback() {
			continue
		}
		parts := strings.Split(ip.String(), ".")
		if len(parts) >= 3 {
			return parts[0] + "." + parts[1] + "." + parts[2]
		}
	}
	return ""
}

// replenishIndexersFromDHT is called when an indexer heartbeat fails and more
// indexers are needed. Queries the DHT and adds fresh entries to StaticIndexers.
func replenishIndexersFromDHT(h host.Host, need int) {
	if need <= 0 || discoveryDHT == nil {
		return
	}
	logger := oclib.GetLogger()
	logger.Info().Int("need", need).Msg("[dht] replenishing indexer pool from DHT")

	providers := DiscoverIndexersFromDHT(h, discoveryDHT, need*3)
	selected := SelectByFillRate(providers, nil, need)
	if len(selected) == 0 {
		logger.Warn().Msg("[dht] no indexers found in DHT for replenishment")
		return
	}

	added := 0
	for _, ai := range selected {
		addr := addrKey(ai)
		if !Indexers.ExistsAddr(addr) {
			adCopy := ai
			Indexers.SetAddr(addr, &adCopy)
			added++
		}
	}

	if added > 0 {
		logger.Info().Int("added", added).Msg("[dht] indexers added from DHT")
		Indexers.NudgeIt()
	}
}

// addrKey returns the canonical map key for an AddrInfo.
// The PeerID is used as key so the same peer is never stored twice regardless
// of which of its addresses was seen first.
func addrKey(ai pp.AddrInfo) string {
	return ai.ID.String()
}