diff --git a/p2p/enr/enr.go b/p2p/enr/enr.go index 2d155450e6..594c385d0d 100644 --- a/p2p/enr/enr.go +++ b/p2p/enr/enr.go @@ -14,7 +14,7 @@ // You should have received a copy of the GNU Lesser General Public License // along with the go-ethereum library. If not, see . -// Package discover implements the Ethereum Node Record as per https://github.com/ethereum/EIPs/pull/778 +// Package enr implements the Ethereum Node Record as per https://github.com/ethereum/EIPs/pull/778 package enr import ( @@ -104,20 +104,23 @@ func (r *Record) Set(k Key) { if err != nil { panic(fmt.Errorf("enr: can't encode %s: %v", k.ENRKey(), err)) } - for i, p := range r.pairs { - if p.k == k.ENRKey() { - // replace value of pair - r.pairs[i].v = blob - return - } 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 - return - } + + 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() { + // element is present at r.pairs[i] + r.pairs[i].v = blob + return + } else if i < len(r.pairs) { + // 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 + return } + + // element should be placed at the end of r.pairs r.pairs = append(r.pairs, pair{k.ENRKey(), blob}) } diff --git a/p2p/enr/enr_test.go b/p2p/enr/enr_test.go index e5de4ccda8..d6bafd4b65 100644 --- a/p2p/enr/enr_test.go +++ b/p2p/enr/enr_test.go @@ -76,6 +76,48 @@ func TestGetSetIP4(t *testing.T) { } } +// TestSortedGetAndSet tests that Set produced a sorted pairs slice. +func TestSortedGetAndSet(t *testing.T) { + type pair struct { + k string + v uint32 + } + + for _, tt := range []struct { + input []pair + want []pair + }{ + { + input: []pair{{"a", 1}, {"c", 2}, {"b", 3}}, + want: []pair{{"a", 1}, {"b", 3}, {"c", 2}}, + }, + { + input: []pair{{"a", 1}, {"c", 2}, {"b", 3}, {"d", 4}, {"a", 5}, {"bb", 6}}, + want: []pair{{"a", 5}, {"b", 3}, {"bb", 6}, {"c", 2}, {"d", 4}}, + }, + { + input: []pair{{"c", 2}, {"b", 3}, {"d", 4}, {"a", 5}, {"bb", 6}}, + want: []pair{{"a", 5}, {"b", 3}, {"bb", 6}, {"c", 2}, {"d", 4}}, + }, + } { + var r Record + for _, i := range tt.input { + r.Set(WithKey(i.k, &i.v)) + } + for i, w := range tt.want { + // set got's key from r.pair[i], so that we preserve order of pairs + got := pair{k: r.pairs[i].k} + if ok, err := r.Load(WithKey(w.k, &got.v)); !ok || err != nil { + t.Fatal(err) + } + + if got != w { + t.Fatalf("expected %#v, got %#v", w, got) + } + } + } +} + // TestGetSetIP6 tests encoding/decoding and setting/getting of the enr.IP6 type func TestGetSetIP6(t *testing.T) { ip := IP6(net.IP{0x20, 0x01, 0x48, 0x60, 0, 0, 0x20, 0x01, 0, 0, 0, 0, 0, 0, 0x00, 0x68})