From 315f8ff16c32dcfc5df2add44d1511da637ee1bc Mon Sep 17 00:00:00 2001
From: Alex Grintsvayg <git@grin.io>
Date: Tue, 22 May 2018 12:27:49 -0400
Subject: [PATCH] self-store

---
 dht/dht.go           | 12 +++++++++++-
 dht/node.go          |  6 +++++-
 dht/routing_table.go |  4 ++++
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/dht/dht.go b/dht/dht.go
index 76877ef..afd9491 100644
--- a/dht/dht.go
+++ b/dht/dht.go
@@ -219,7 +219,11 @@ func (dht *DHT) Announce(hash Bitmap) error {
 		return err
 	}
 
-	// TODO: if this node is closer than farthest peer, store locally and pop farthest peer
+	// if we found less than K contacts, or current node is closer than farthest contact
+	if len(res.Contacts) < bucketSize || dht.node.id.Xor(hash).Less(res.Contacts[bucketSize-1].ID.Xor(hash)) {
+		// pop last contact, and self-store instead
+		res.Contacts[bucketSize-1] = dht.contact
+	}
 
 	wg := &sync.WaitGroup{}
 	for _, c := range res.Contacts {
@@ -261,6 +265,12 @@ func (dht *DHT) storeOnNode(hash Bitmap, c Contact) {
 	dht.stopWG.Add(1)
 	defer dht.stopWG.Done()
 
+	// self-store
+	if dht.contact.Equals(c) {
+		dht.node.Store(hash, c)
+		return
+	}
+
 	resCh, cancel := dht.node.SendCancelable(c, Request{
 		Method: findValueMethod,
 		Arg:    &hash,
diff --git a/dht/node.go b/dht/node.go
index 26dbd8a..1ad8fd8 100644
--- a/dht/node.go
+++ b/dht/node.go
@@ -238,7 +238,7 @@ func (n *Node) handleRequest(addr *net.UDPAddr, request Request) {
 		// TODO: we should be sending the IP in the request, not just using the sender's IP
 		// TODO: should we be using StoreArgs.NodeID or StoreArgs.Value.LbryID ???
 		if n.tokens.Verify(request.StoreArgs.Value.Token, request.NodeID, addr) {
-			n.store.Upsert(request.StoreArgs.BlobHash, Contact{ID: request.StoreArgs.NodeID, IP: addr.IP, Port: request.StoreArgs.Value.Port})
+			n.Store(request.StoreArgs.BlobHash, Contact{ID: request.StoreArgs.NodeID, IP: addr.IP, Port: request.StoreArgs.Value.Port})
 			n.sendMessage(addr, Response{ID: request.ID, NodeID: n.id, Data: storeSuccessResponse})
 		} else {
 			n.sendMessage(addr, Error{ID: request.ID, NodeID: n.id, ExceptionType: "invalid-token"})
@@ -441,3 +441,7 @@ func (n *Node) startRoutingTableGrooming() {
 		}
 	}()
 }
+
+func (n *Node) Store(hash Bitmap, c Contact) {
+	n.store.Upsert(hash, c)
+}
diff --git a/dht/routing_table.go b/dht/routing_table.go
index 6b18fd6..b0f1739 100644
--- a/dht/routing_table.go
+++ b/dht/routing_table.go
@@ -27,6 +27,10 @@ type Contact struct {
 	Port int
 }
 
+func (c Contact) Equals(other Contact) bool {
+	return c.ID == other.ID
+}
+
 func (c Contact) Addr() *net.UDPAddr {
 	return &net.UDPAddr{IP: c.IP, Port: c.Port}
 }