DHT-kademlia-P2Psec/pkg/routing_table.go

package kademlia

// TODO: Tests
const (
	threshold  = 10
	bucketSize = 20
)

// RoutingTableNode interface implemented by both InnerNode and Leaf
type RoutingTableNode interface {
	isNode()
	findClosest(id [32]byte, numberOfNodes int) []P2PNode
	insert(node P2PNode) RoutingTableNode
}

type Leaf struct {
	bucket         []P2PNode
	hostClientNode P2PNode
	depth          uint64
	maxDepth       uint64
	prefixLen      uint64
}

// IsNode ensures Leaf and InnerNode implement RoutingTableNode
func (i *InnerNode) isNode() {}
func (l *Leaf) isNode()      {}

type RoutingTable struct {
	root                RoutingTableNode
	hostClientNode      P2PNode
	matchedPrefixLength uint64
}

func NewRoutingTable(clientNode P2PNode, matchedPrefixLength uint64) *RoutingTable {
	return &RoutingTable{
		root:                &Leaf{bucket: make([]P2PNode, 0), depth: 0, maxDepth: 256, prefixLen: matchedPrefixLength, hostClientNode: clientNode},
		hostClientNode:      clientNode,
		matchedPrefixLength: matchedPrefixLength,
	}

}

type InnerNode struct {
	branches       []RoutingTableNode
	hostClientNode P2PNode
	depth          uint64
	prefixLength   uint64
}

// FindClosest Return a list of nodes in the routing table closest to the provided ID
func (r *RoutingTable) FindClosest(targetID [32]byte, numberOfNodes int) *[]P2PNode {
	ret := r.root.findClosest(targetID, numberOfNodes)
	return &ret
}

func (i *InnerNode) findClosest(id [32]byte, numberOfNodes int) []P2PNode {
	distance := uint64(0)
	result := make([]P2PNode, 0)

	// go through the buckets in increasing distance from the key until sufficiently many nodes were found
	for distance < 1<<i.prefixLength && len(result) < numberOfNodes {
		index := getBitsAt(id, i.depth, i.prefixLength)
		index = index ^ distance
		result = append(result, i.branches[index].findClosest(id, numberOfNodes)...)
		distance++
	}
	return result
}

func (l *Leaf) findClosest(_ [32]byte, _ int) []P2PNode {
	return l.bucket
}

func (r *RoutingTable) Insert(node P2PNode) {
	r.root = r.root.insert(node)
}

func (i *InnerNode) insert(node P2PNode) RoutingTableNode {
	// There is one bucket for every bit sequence. The index for a sequence is its numerical value
	indexToInsert := getBitsAt(node.NodeID, i.depth, i.prefixLength)
	i.branches[indexToInsert] = i.branches[indexToInsert].insert(node)
	return i
}

func (l *Leaf) insert(node P2PNode) RoutingTableNode {
	// If the bit at the current depth matches that of our client, i.e. we are on the 'right' path,
	// and there are enough nodes in the bucket, the node is eligible to be split
	if len(l.bucket) > threshold && getBitsAt(node.NodeID, 0, l.depth) == getBitsAt(l.hostClientNode.NodeID, 0, l.depth) {
		branches := make([]RoutingTableNode, 1<<l.prefixLen)
		for i := range len(branches) {
			branches[i] = &Leaf{bucket: make([]P2PNode, 0), depth: l.depth + l.prefixLen, maxDepth: 256, prefixLen: l.prefixLen, hostClientNode: l.hostClientNode}
		}
		for _, n := range l.bucket {
			index := getBitsAt(n.NodeID, l.depth, l.prefixLen)
			branches[index] = branches[index].insert(n)
		}

		resultingInnerNode := &InnerNode{branches: branches, prefixLength: l.prefixLen, depth: l.depth, hostClientNode: l.hostClientNode}
		return resultingInnerNode.insert(node)
	} else {
		l.addToBucket(node)
		return l
	}
}

func (l *Leaf) addToBucket(node P2PNode) {
	// If the node is already contained, skip
	if ContainsID(l.bucket, node.NodeID) {
		return
	}

	// If not and there is room, simply add the node
	if len(l.bucket) < bucketSize {
		l.bucket = append(l.bucket, node)
	} else {
		// If not, only add it if the currently first node has gone offline, since if it is still alive, the likelihood
		// of it staying online is very high (see Kademlia paper)
		if Ping(l.bucket[0], l.hostClientNode) {
			l.bucket = append(l.bucket[1:], l.bucket[0])
		} else {
			l.bucket = append(l.bucket[1:], node)
		}
	}
}