From bd9696c33ddc87843031a20e769b366e4e8eee6a Mon Sep 17 00:00:00 2001
From: Anton Evangelatov <anton.evangelatov@gmail.com>
Date: Mon, 4 Dec 2017 14:17:32 +0100
Subject: [PATCH] p2p/enr: sort pairs prior to RLP encoding; binary search on
 Load

---
 p2p/enr/enr.go      | 37 +++++++++++++++++++++++++++----------
 p2p/enr/enr_test.go | 21 +++++++++++++++++++++
 2 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/p2p/enr/enr.go b/p2p/enr/enr.go
index 1127a1847d..40d5cfec8d 100644
--- a/p2p/enr/enr.go
+++ b/p2p/enr/enr.go
@@ -69,11 +69,10 @@ func (r *Record) SetSeq(s uint32) {
 }
 
 func (r *Record) Load(k Key) (bool, error) {
-	for _, p := range r.pairs {
-		if p.k == k.ENRKey() {
-			err := rlp.DecodeBytes(p.v, k)
-			return true, err
-		}
+	i := sort.Search(len(r.pairs), func(i int) bool { return r.pairs[i].k >= k.ENRKey() })
+
+	if i < len(r.pairs) && r.pairs[i].k == k.ENRKey() {
+		return true, rlp.DecodeBytes(r.pairs[i].v, k)
 	}
 
 	return false, errors.New("record does not exist")
@@ -85,7 +84,29 @@ func (r *Record) Set(k Key) error {
 	if err != nil {
 		return err
 	}
-	r.pairs = append(r.pairs, pair{k.ENRKey(), blob})
+	var inserted bool
+	for i, p := range r.pairs {
+		if p.k == k.ENRKey() {
+			// replace value of pair
+			r.pairs[i].v = blob
+			inserted = true
+
+			break
+		} else if p.k > k.ENRKey() {
+			// insert pair before i-th elem
+			el := pair{k.ENRKey(), blob}
+
+			r.pairs = append(r.pairs, pair{})
+			copy(r.pairs[i+1:], r.pairs[i:])
+			r.pairs[i] = el
+
+			inserted = true
+			break
+		}
+	}
+	if !inserted {
+		r.pairs = append(r.pairs, pair{k.ENRKey(), blob})
+	}
 	return nil
 }
 
@@ -213,10 +234,6 @@ func (r *Record) Sign(privkey *ecdsa.PrivateKey) error {
 }
 
 func (r *Record) serialisedContent() ([]byte, error) {
-	sort.Slice(r.pairs, func(i, j int) bool {
-		return r.pairs[i].k < r.pairs[j].k
-	})
-
 	list := []interface{}{r.seq}
 
 	for _, p := range r.pairs {
diff --git a/p2p/enr/enr_test.go b/p2p/enr/enr_test.go
index d452af3662..4867c9e6c4 100644
--- a/p2p/enr/enr_test.go
+++ b/p2p/enr/enr_test.go
@@ -151,6 +151,27 @@ func TestDirty(t *testing.T) {
 	}
 }
 
+func TestGetSetOverwrite(t *testing.T) {
+	var r Record
+
+	ip := IP4(net.IP{192, 168, 0, 3})
+	r.Set(ip)
+
+	ip2 := IP4(net.IP{192, 168, 0, 4})
+	r.Set(ip2)
+
+	var ip3 IP4
+
+	_, err := r.Load(&ip3)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if bytes.Compare(ip2, ip3) != 0 {
+		t.Fatalf("got %#v, expected %#v", ip2, ip3)
+	}
+}
+
 func TestSignEncodeAndDecode(t *testing.T) {
 	privkey, err := crypto.HexToECDSA(privkeyHex)
 	if err != nil {