Implement DHT server and update client

adnl/dht/client.go

@@ -170,13 +170,12 @@ func (c *Client) addNode(node *Node) (_ *dhtNode, err error) {
 	addr := node.AddrList.Addresses[0].IP.String() + ":" + fmt.Sprint(node.AddrList.Addresses[0].Port)
 
 	if hf := bucket.findNode(kid); hf != nil {
-		if hf.addr == addr {
-			return nil, fmt.Errorf("node already exists")
-		}
-		// updated address otherwise
+		hf.updateEndpoint(addr, pub.Key)
+		hf.setNode(node)
+		return hf, nil
 	}
 
-	kNode := c.initNode(kid, addr, pub.Key)
+	kNode := c.initNode(kid, addr, pub.Key, node)
 	bucket.addNode(kNode)
 
 	return kNode, nil
@@ -551,3 +550,53 @@ func (c *Client) buildPriorityList(id []byte) *priorityList {
 
 	return plistGood
 }
+
+type nodeAffinity struct {
+	node     *Node
+	affinity int
+}
+
+func (c *Client) nearestNodeAffinities(id []byte, limit int) []nodeAffinity {
+	if limit <= 0 {
+		return nil
+	}
+
+	plist := c.buildPriorityList(id)
+	result := make([]nodeAffinity, 0, limit)
+	seen := map[string]struct{}{}
+
+	for len(result) < limit {
+		node, pr := plist.Get()
+		if node == nil {
+			break
+		}
+
+		key := node.id()
+		if _, ok := seen[key]; ok {
+			continue
+		}
+		seen[key] = struct{}{}
+
+		nn := node.getNode()
+		if nn == nil {
+			continue
+		}
+
+		result = append(result, nodeAffinity{node: nn, affinity: pr})
+	}
+
+	return result
+}
+
+func (c *Client) nearestNodes(id []byte, limit int) ([]*Node, []int) {
+	details := c.nearestNodeAffinities(id, limit)
+	nodes := make([]*Node, 0, len(details))
+	affinities := make([]int, 0, len(details))
+
+	for _, d := range details {
+		nodes = append(nodes, d.node)
+		affinities = append(affinities, d.affinity)
+	}
+
+	return nodes, affinities
+}

adnl/dht/client_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/xssnick/tonutils-go/liteclient"
 	"github.com/xssnick/tonutils-go/tl"
 	"net"
+	"os"
 	"reflect"
 	"strconv"
 	"testing"
@@ -406,6 +407,10 @@ func TestClient_FindAddressesUnit(t *testing.T) {
 }
 
 func TestClient_FindAddressesIntegration(t *testing.T) {
+	if _, ok := os.LookupEnv("TON_DHT_INTEGRATION_TEST"); !ok {
+		t.Skip("TON_DHT_INTEGRATION_TEST is not set; skipping integration test that requires public DHT network access")
+	}
+
 	_, priv, err := ed25519.GenerateKey(nil)
 	if err != nil {
 		t.Fatal(err)
@@ -472,6 +477,10 @@ func TestClient_Close(t *testing.T) {
 }
 
 func TestClient_StoreAddressIntegration(t *testing.T) {
+	if _, ok := os.LookupEnv("TON_DHT_INTEGRATION_TEST"); !ok {
+		t.Skip("TON_DHT_INTEGRATION_TEST is not set; skipping integration test that requires public DHT network access")
+	}
+
 	_, priv, err := ed25519.GenerateKey(nil)
 	if err != nil {
 		t.Fatal(err)

adnl/dht/node.go

@@ -6,8 +6,10 @@ import (
 	"crypto/ed25519"
 	"encoding/hex"
 	"fmt"
+	"github.com/xssnick/tonutils-go/adnl/address"
 	"github.com/xssnick/tonutils-go/adnl/keys"
 	"math/bits"
+	"net"
 	"reflect"
 	"sync"
 	"sync/atomic"
@@ -27,6 +29,8 @@ type dhtNode struct {
 	addr      string
 	serverKey ed25519.PublicKey
 
+	node *Node
+
 	currentState int
 	badScore     int32
 
@@ -58,16 +62,50 @@ func (l dhtNodeList) Less(i, j int) bool {
 	return atomic.LoadInt64(&l[i].ping) < atomic.LoadInt64(&l[j].ping)
 }
 
-func (c *Client) initNode(id []byte, addr string, serverKey ed25519.PublicKey) *dhtNode {
+func (c *Client) initNode(id []byte, addr string, serverKey ed25519.PublicKey, node *Node) *dhtNode {
 	n := &dhtNode{
 		adnlId:    id,
 		addr:      addr,
 		serverKey: serverKey,
 		client:    c,
 	}
+	if node != nil {
+		n.setNode(node)
+	}
 	return n
 }
 
+func (n *dhtNode) setNode(node *Node) {
+	if node == nil {
+		return
+	}
+
+	clone := cloneNode(node)
+	n.mx.Lock()
+	n.node = clone
+	n.mx.Unlock()
+}
+
+func (n *dhtNode) updateEndpoint(addr string, key ed25519.PublicKey) {
+	n.mx.Lock()
+	if addr != "" {
+		n.addr = addr
+	}
+	if key != nil {
+		n.serverKey = key
+	}
+	n.mx.Unlock()
+}
+
+func (n *dhtNode) getNode() *Node {
+	n.mx.Lock()
+	defer n.mx.Unlock()
+	if n.node == nil {
+		return nil
+	}
+	return cloneNode(n.node)
+}
+
 func (n *dhtNode) findNodes(ctx context.Context, id []byte, K int32) (result []*Node, err error) {
 	val, err := tl.Serialize(FindNode{
 		Key: id,
@@ -316,3 +354,48 @@ func affinity(x, y []byte) uint {
 	}
 	return result
 }
+
+func cloneNode(node *Node) *Node {
+	if node == nil {
+		return nil
+	}
+
+	clone := *node
+
+	switch id := node.ID.(type) {
+	case keys.PublicKeyED25519:
+		cp := make([]byte, len(id.Key))
+		copy(cp, id.Key)
+		clone.ID = keys.PublicKeyED25519{Key: cp}
+	case keys.PublicKeyAES:
+		cp := make([]byte, len(id.Key))
+		copy(cp, id.Key)
+		clone.ID = keys.PublicKeyAES{Key: cp}
+	}
+
+	if node.AddrList != nil {
+		addrClone := *node.AddrList
+		if len(node.AddrList.Addresses) > 0 {
+			addrClone.Addresses = make([]*address.UDP, len(node.AddrList.Addresses))
+			for i, addr := range node.AddrList.Addresses {
+				if addr == nil {
+					continue
+				}
+				udpClone := *addr
+				if addr.IP != nil {
+					ip := make(net.IP, len(addr.IP))
+					copy(ip, addr.IP)
+					udpClone.IP = ip
+				}
+				addrClone.Addresses[i] = &udpClone
+			}
+		}
+		clone.AddrList = &addrClone
+	}
+
+	if len(node.Signature) > 0 {
+		clone.Signature = append([]byte(nil), node.Signature...)
+	}
+
+	return &clone
+}