153 lines
3.5 KiB
Go
153 lines
3.5 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[:])
|
|
}
|
|
|
|
// Bucket represents a Kademlia k-bucket.
|
|
type Bucket struct {
|
|
nodes []Node
|
|
}
|
|
|
|
// RoutingTable manages known nodes using Kademlia k-buckets.
|
|
type RoutingTable struct {
|
|
mu sync.RWMutex
|
|
ownID NodeID
|
|
buckets [160]*Bucket
|
|
}
|
|
|
|
// NewRoutingTable creates a new routing table.
|
|
func NewRoutingTable(ownID NodeID) *RoutingTable {
|
|
rt := &RoutingTable{
|
|
ownID: ownID,
|
|
}
|
|
for i := 0; i < 160; i++ {
|
|
rt.buckets[i] = &Bucket{nodes: make([]Node, 0, MaxNodes)}
|
|
}
|
|
return rt
|
|
}
|
|
|
|
// bucketIndex calculates the appropriate bucket index for a given node ID.
|
|
// Returns an index from 0 to 159, or -1 if the ID is our own.
|
|
func (rt *RoutingTable) bucketIndex(target NodeID) int {
|
|
for i := 0; i < 20; i++ {
|
|
xor := rt.ownID[i] ^ target[i]
|
|
if xor != 0 {
|
|
// Find the most significant bit set in the byte
|
|
for j := 7; j >= 0; j-- {
|
|
if (xor & (1 << j)) != 0 {
|
|
return 159 - (i*8 + (7 - j))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return -1
|
|
}
|
|
|
|
// AddNode adds a node to the routing table or updates it.
|
|
func (rt *RoutingTable) AddNode(n Node) {
|
|
idx := rt.bucketIndex(n.ID)
|
|
if idx == -1 {
|
|
return // Do not add ourselves
|
|
}
|
|
|
|
rt.mu.Lock()
|
|
defer rt.mu.Unlock()
|
|
|
|
bucket := rt.buckets[idx]
|
|
|
|
// Check if already exists and update
|
|
for i, existing := range bucket.nodes {
|
|
if existing.ID == n.ID {
|
|
bucket.nodes[i].Addr = n.Addr
|
|
// Move to end (most recently seen)
|
|
node := bucket.nodes[i]
|
|
bucket.nodes = append(bucket.nodes[:i], bucket.nodes[i+1:]...)
|
|
bucket.nodes = append(bucket.nodes, node)
|
|
return
|
|
}
|
|
}
|
|
|
|
// Add new node if bucket is not full
|
|
if len(bucket.nodes) < MaxNodes {
|
|
bucket.nodes = append(bucket.nodes, n)
|
|
} else {
|
|
// In a full Kademlia implementation, we would ping the oldest node
|
|
// and replace it if it doesn't respond. For now, just drop the new one.
|
|
}
|
|
}
|
|
|
|
// 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()
|
|
|
|
var allNodes []Node
|
|
|
|
// Fast path: find the target's bucket
|
|
idx := rt.bucketIndex(target)
|
|
if idx != -1 {
|
|
allNodes = append(allNodes, rt.buckets[idx].nodes...)
|
|
}
|
|
|
|
// Collect nodes from neighboring buckets if we don't have enough
|
|
if len(allNodes) < count {
|
|
// We just collect all nodes and sort them for simplicity.
|
|
// Optimizing this is possible but not strictly necessary for < 1280 nodes total.
|
|
allNodes = nil
|
|
for i := 0; i < 160; i++ {
|
|
allNodes = append(allNodes, rt.buckets[i].nodes...)
|
|
}
|
|
}
|
|
|
|
sort.Slice(allNodes, func(i, j int) bool {
|
|
distI := Distance(allNodes[i].ID, target)
|
|
distJ := Distance(allNodes[j].ID, target)
|
|
return distI.Cmp(distJ) < 0
|
|
})
|
|
|
|
if len(allNodes) > count {
|
|
return allNodes[:count]
|
|
}
|
|
return allNodes
|
|
}
|
|
|
|
// Len returns the total number of nodes in the routing table.
|
|
func (rt *RoutingTable) Len() int {
|
|
rt.mu.RLock()
|
|
defer rt.mu.RUnlock()
|
|
count := 0
|
|
for i := 0; i < 160; i++ {
|
|
count += len(rt.buckets[i].nodes)
|
|
}
|
|
return count
|
|
}
|