108 lines
2.2 KiB
Go
108 lines
2.2 KiB
Go
package dht
|
|
|
|
import (
|
|
"crypto/rand"
|
|
"math/big"
|
|
"net"
|
|
"sort"
|
|
"sync"
|
|
)
|
|
|
|
// NodeID is a 160-bit Kademlia ID.
|
|
type NodeID [20]byte
|
|
|
|
// RandomNodeID generates a new random NodeID.
|
|
func RandomNodeID() NodeID {
|
|
var id NodeID
|
|
_, _ = rand.Read(id[:])
|
|
return id
|
|
}
|
|
|
|
// Node represents a contact in the DHT network.
|
|
type Node struct {
|
|
ID NodeID
|
|
Addr *net.UDPAddr
|
|
}
|
|
|
|
// Distance calculates the XOR distance between two NodeIDs.
|
|
func Distance(a, b NodeID) *big.Int {
|
|
var xor [20]byte
|
|
for i := 0; i < 20; i++ {
|
|
xor[i] = a[i] ^ b[i]
|
|
}
|
|
return new(big.Int).SetBytes(xor[:])
|
|
}
|
|
|
|
// RoutingTable manages known nodes.
|
|
type RoutingTable struct {
|
|
mu sync.RWMutex
|
|
ownID NodeID
|
|
nodes []Node
|
|
}
|
|
|
|
// NewRoutingTable creates a new routing table.
|
|
func NewRoutingTable(ownID NodeID) *RoutingTable {
|
|
return &RoutingTable{
|
|
ownID: ownID,
|
|
nodes: make([]Node, 0),
|
|
}
|
|
}
|
|
|
|
// AddNode adds a node to the routing table or updates it.
|
|
func (rt *RoutingTable) AddNode(n Node) {
|
|
rt.mu.Lock()
|
|
defer rt.mu.Unlock()
|
|
|
|
if n.ID == rt.ownID {
|
|
return
|
|
}
|
|
|
|
for i, existing := range rt.nodes {
|
|
if existing.ID == n.ID {
|
|
rt.nodes[i].Addr = n.Addr
|
|
return
|
|
}
|
|
}
|
|
|
|
rt.nodes = append(rt.nodes, n)
|
|
|
|
// Sort by distance to our own ID and keep top 1000 nodes for simplicity.
|
|
sort.Slice(rt.nodes, func(i, j int) bool {
|
|
distI := Distance(rt.nodes[i].ID, rt.ownID)
|
|
distJ := Distance(rt.nodes[j].ID, rt.ownID)
|
|
return distI.Cmp(distJ) < 0
|
|
})
|
|
|
|
if len(rt.nodes) > 1000 {
|
|
rt.nodes = rt.nodes[:1000]
|
|
}
|
|
}
|
|
|
|
// ClosestNodes returns the closest nodes to a given target ID.
|
|
func (rt *RoutingTable) ClosestNodes(target NodeID, count int) []Node {
|
|
rt.mu.RLock()
|
|
defer rt.mu.RUnlock()
|
|
|
|
// Since we don't have true k-buckets, we sort the entire list.
|
|
// This is O(n log n) but fine for a small simplified table of 1000 nodes.
|
|
sortedNodes := make([]Node, len(rt.nodes))
|
|
copy(sortedNodes, rt.nodes)
|
|
|
|
sort.Slice(sortedNodes, func(i, j int) bool {
|
|
distI := Distance(sortedNodes[i].ID, target)
|
|
distJ := Distance(sortedNodes[j].ID, target)
|
|
return distI.Cmp(distJ) < 0
|
|
})
|
|
|
|
if len(sortedNodes) > count {
|
|
return sortedNodes[:count]
|
|
}
|
|
return sortedNodes
|
|
}
|
|
|
|
// Len returns the number of nodes in the routing table.
|
|
func (rt *RoutingTable) Len() int {
|
|
rt.mu.RLock()
|
|
defer rt.mu.RUnlock()
|
|
return len(rt.nodes)
|
|
}
|