From 94055287b066e24086b2b6d25c1384d3afe2fdb9 Mon Sep 17 00:00:00 2001
From: Yaroslav Kolomiiets <yarikos@gmail.com>
Date: Mon, 9 Nov 2020 15:05:47 +0000
Subject: [PATCH] simplify pexMsgFactory
---
pex.go | 100 +++++++++++++++------------------------
pex_test.go | 132 ++++++++++++++++++++++++----------------------------
2 files changed, 97 insertions(+), 135 deletions(-)
diff --git a/pex.go b/pex.go
index 85d32667..a6295f36 100644
--- a/pex.go
+++ b/pex.go
@@ -28,16 +28,10 @@ type pexEvent struct {
f pp.PexPeerFlags
}
-// Combines the node addr, as required for pp.PexMsg.
-type pexMsgAdded struct {
- krpc.NodeAddr
- pp.PexPeerFlags
-}
-
-// Makes generating a PexMsg more efficient.
+// facilitates efficient de-duplication while generating PEX messages
type pexMsgFactory struct {
- added map[string]pexMsgAdded
- dropped map[string]krpc.NodeAddr
+ added map[addrKey]pexEvent
+ dropped map[addrKey]pexEvent
}
func (me *pexMsgFactory) DeltaLen() int {
@@ -46,90 +40,78 @@ func (me *pexMsgFactory) DeltaLen() int {
int64(len(me.dropped))))
}
-// Returns the key to use to identify a given addr in the factory. Panics if we can't support the
-// addr later in generating a PexMsg (since adding an unusable addr will cause DeltaLen to be out.)
-func (me *pexMsgFactory) addrKey(addr krpc.NodeAddr) string {
- if addr.IP.To4() != nil {
- addr.IP = addr.IP.To4()
- }
- keyBytes, err := addr.MarshalBinary()
- if err != nil {
- panic(err)
- }
- switch len(keyBytes) {
- case compactIpv4NodeAddrElemSize:
- case compactIpv6NodeAddrElemSize:
- default:
- panic(len(keyBytes))
- }
- return string(keyBytes)
+type addrKey string
+
+// Returns the key to use to identify a given addr in the factory.
+func (me *pexMsgFactory) addrKey(addr net.Addr) addrKey {
+ return addrKey(addr.String())
}
// Returns whether the entry was added (we can check if we're cancelling out another entry and so
// won't hit the limit consuming this event).
-func (me *pexMsgFactory) Add(addr krpc.NodeAddr, flags pp.PexPeerFlags) {
- key := me.addrKey(addr)
+func (me *pexMsgFactory) add(e pexEvent) {
+ key := me.addrKey(e.addr)
if _, ok := me.dropped[key]; ok {
delete(me.dropped, key)
return
}
if me.added == nil {
- me.added = make(map[string]pexMsgAdded, pexMaxDelta)
+ me.added = make(map[addrKey]pexEvent, pexMaxDelta)
}
- me.added[key] = pexMsgAdded{addr, flags}
-
+ me.added[key] = e
}
// Returns whether the entry was added (we can check if we're cancelling out another entry and so
// won't hit the limit consuming this event).
-func (me *pexMsgFactory) Drop(addr krpc.NodeAddr) {
- key := me.addrKey(addr)
+func (me *pexMsgFactory) drop(e pexEvent) {
+ key := me.addrKey(e.addr)
if _, ok := me.added[key]; ok {
delete(me.added, key)
return
}
if me.dropped == nil {
- me.dropped = make(map[string]krpc.NodeAddr, pexMaxDelta)
+ me.dropped = make(map[addrKey]pexEvent, pexMaxDelta)
}
- me.dropped[key] = addr
+ me.dropped[key] = e
}
func (me *pexMsgFactory) addEvent(event pexEvent) {
- addr, ok := nodeAddr(event.addr)
- if !ok {
- return
- }
switch event.t {
case pexAdd:
- me.Add(addr, event.f)
+ me.add(event)
case pexDrop:
- me.Drop(addr)
+ me.drop(event)
default:
panic(event.t)
}
}
-var compactIpv4NodeAddrElemSize = krpc.CompactIPv4NodeAddrs{}.ElemSize()
-var compactIpv6NodeAddrElemSize = krpc.CompactIPv6NodeAddrs{}.ElemSize()
-
func (me *pexMsgFactory) PexMsg() (ret pp.PexMsg) {
for key, added := range me.added {
- switch len(key) {
- case compactIpv4NodeAddrElemSize:
- ret.Added = append(ret.Added, added.NodeAddr)
- ret.AddedFlags = append(ret.AddedFlags, added.PexPeerFlags)
- case compactIpv6NodeAddrElemSize:
- ret.Added6 = append(ret.Added6, added.NodeAddr)
- ret.Added6Flags = append(ret.Added6Flags, added.PexPeerFlags)
+ addr, ok := nodeAddr(added.addr)
+ if !ok {
+ continue
+ }
+ switch len(addr.IP) {
+ case net.IPv4len:
+ ret.Added = append(ret.Added, addr)
+ ret.AddedFlags = append(ret.AddedFlags, added.f)
+ case net.IPv6len:
+ ret.Added6 = append(ret.Added6, addr)
+ ret.Added6Flags = append(ret.Added6Flags, added.f)
default:
panic(key)
}
}
- for key, addr := range me.dropped {
- switch len(key) {
- case compactIpv4NodeAddrElemSize:
+ for key, dropped := range me.dropped {
+ addr, ok := nodeAddr(dropped.addr)
+ if !ok {
+ continue
+ }
+ switch len(addr.IP) {
+ case net.IPv4len:
ret.Dropped = append(ret.Dropped, addr)
- case compactIpv6NodeAddrElemSize:
+ case net.IPv6len:
ret.Dropped6 = append(ret.Dropped6, addr)
default:
panic(key)
@@ -138,14 +120,6 @@ func (me *pexMsgFactory) PexMsg() (ret pp.PexMsg) {
return
}
-func mustNodeAddr(addr net.Addr) krpc.NodeAddr {
- ret, ok := nodeAddr(addr)
- if !ok {
- panic(addr)
- }
- return ret
-}
-
// Convert an arbitrary torrent peer Addr into one that can be represented by the compact addr
// format.
func nodeAddr(addr net.Addr) (_ krpc.NodeAddr, ok bool) {
diff --git a/pex_test.go b/pex_test.go
index a59c8ee4..006644c2 100644
--- a/pex_test.go
+++ b/pex_test.go
@@ -12,11 +12,23 @@ import (
)
var (
- addrs = []net.Addr{
+ addrs6 = []net.Addr{
&net.TCPAddr{IP: net.IPv6loopback, Port: 4747},
&net.TCPAddr{IP: net.IPv6loopback, Port: 4748},
+ &net.TCPAddr{IP: net.IPv6loopback, Port: 4749},
+ &net.TCPAddr{IP: net.IPv6loopback, Port: 4750},
+ }
+ addrs4 = []net.Addr{
&net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747},
&net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4748},
+ &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4749},
+ &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4750},
+ }
+ addrs = []net.Addr{
+ addrs6[0],
+ addrs6[1],
+ addrs4[0],
+ addrs4[1],
}
f = pp.PexOutgoingConn
)
@@ -111,6 +123,14 @@ func TestPexReset(t *testing.T) {
require.EqualValues(t, targ, s)
}
+func mustNodeAddr(addr net.Addr) krpc.NodeAddr {
+ ret, ok := nodeAddr(addr)
+ if !ok {
+ panic(addr)
+ }
+ return ret
+}
+
var testcases = []struct {
name string
in *pexState
@@ -239,24 +259,17 @@ var testcases = []struct {
// deterministic. Because the flags are in a different array, we can't just use testify's
// ElementsMatch because the ordering *does* still matter between an added addr and its flags.
type comparablePexMsg struct {
- added, added6 []pexMsgAdded
- dropped, dropped6 []krpc.NodeAddr
-}
-
-func (me *comparablePexMsg) makeAdded(addrs []krpc.NodeAddr, flags []pp.PexPeerFlags) (ret []pexMsgAdded) {
- for i, addr := range addrs {
- ret = append(ret, pexMsgAdded{
- NodeAddr: addr,
- PexPeerFlags: flags[i],
- })
- }
- return
+ added, added6 []krpc.NodeAddr
+ addedFlags, added6Flags []pp.PexPeerFlags
+ dropped, dropped6 []krpc.NodeAddr
}
// Such Rust-inspired.
func (me *comparablePexMsg) From(f pp.PexMsg) {
- me.added = me.makeAdded(f.Added, f.AddedFlags)
- me.added6 = me.makeAdded(f.Added6, f.Added6Flags)
+ me.added = f.Added
+ me.addedFlags = f.AddedFlags
+ me.added6 = f.Added6
+ me.added6Flags = f.Added6Flags
me.dropped = f.Dropped
me.dropped6 = f.Dropped6
}
@@ -265,7 +278,9 @@ func (me *comparablePexMsg) From(f pp.PexMsg) {
// in pexMsgFactory that preserve insert ordering.
func (actual comparablePexMsg) AssertEqual(t *testing.T, expected comparablePexMsg) {
assert.ElementsMatch(t, expected.added, actual.added)
+ assert.ElementsMatch(t, expected.addedFlags, actual.addedFlags)
assert.ElementsMatch(t, expected.added6, actual.added6)
+ assert.ElementsMatch(t, expected.added6Flags, actual.added6Flags)
assert.ElementsMatch(t, expected.dropped, actual.dropped)
assert.ElementsMatch(t, expected.dropped6, actual.dropped6)
}
@@ -321,60 +336,47 @@ func TestPexInitialNoCutoff(t *testing.T) {
}
func TestPexAdd(t *testing.T) {
- addrs4 := []krpc.NodeAddr{
- krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747}, // 0
- krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4748}, // 1
- krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 2), Port: 4747}, // 2
- krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 2), Port: 4748}, // 3
- }
- addrs6 := []krpc.NodeAddr{
- krpc.NodeAddr{IP: net.IPv6loopback, Port: 4747}, // 0
- krpc.NodeAddr{IP: net.IPv6loopback, Port: 4748}, // 1
- krpc.NodeAddr{IP: net.IPv6loopback, Port: 4749}, // 2
- krpc.NodeAddr{IP: net.IPv6loopback, Port: 4750}, // 3
- }
- f := pp.PexPrefersEncryption | pp.PexOutgoingConn
-
t.Run("ipv4", func(t *testing.T) {
addrs := addrs4
var m pexMsgFactory
- m.Drop(addrs[0])
- m.Add(addrs[1], f)
+ m.addEvent(pexEvent{pexDrop, addrs[0], 0})
+ m.addEvent(pexEvent{pexAdd, addrs[1], f})
for _, addr := range addrs {
- m.Add(addr, f)
+ m.addEvent(pexEvent{pexAdd, addr, f})
}
targ := pp.PexMsg{
Added: krpc.CompactIPv4NodeAddrs{
- addrs[1],
- addrs[2],
- addrs[3],
+ mustNodeAddr(addrs[1]),
+ mustNodeAddr(addrs[2]),
+ mustNodeAddr(addrs[3]),
},
AddedFlags: []pp.PexPeerFlags{f, f, f},
}
- assertPexMsgsEqual(t, targ, m.PexMsg())
+ out := m.PexMsg()
+ assertPexMsgsEqual(t, targ, out)
})
t.Run("ipv6", func(t *testing.T) {
addrs := addrs6
var m pexMsgFactory
- m.Drop(addrs[0])
- m.Add(addrs[1], f)
+ m.addEvent(pexEvent{pexDrop, addrs[0], 0})
+ m.addEvent(pexEvent{pexAdd, addrs[1], f})
for _, addr := range addrs {
- m.Add(addr, f)
+ m.addEvent(pexEvent{pexAdd, addr, f})
}
targ := pp.PexMsg{
Added6: krpc.CompactIPv6NodeAddrs{
- addrs[1],
- addrs[2],
- addrs[3],
+ mustNodeAddr(addrs[1]),
+ mustNodeAddr(addrs[2]),
+ mustNodeAddr(addrs[3]),
},
Added6Flags: []pp.PexPeerFlags{f, f, f},
}
assertPexMsgsEqual(t, targ, m.PexMsg())
})
t.Run("empty", func(t *testing.T) {
- addr := krpc.NodeAddr{}
+ nullAddr := &net.TCPAddr{}
var xm pexMsgFactory
- assert.Panics(t, func() { xm.Add(addr, f) })
+ xm.addEvent(pexEvent{pexAdd, nullAddr, f})
m := xm.PexMsg()
require.EqualValues(t, 0, len(m.Added))
require.EqualValues(t, 0, len(m.AddedFlags))
@@ -384,33 +386,19 @@ func TestPexAdd(t *testing.T) {
}
func TestPexDrop(t *testing.T) {
- addrs4 := []krpc.NodeAddr{
- krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747}, // 0
- krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4748}, // 1
- krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 2), Port: 4747}, // 2
- krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 2), Port: 4748}, // 3
- }
- addrs6 := []krpc.NodeAddr{
- krpc.NodeAddr{IP: net.IPv6loopback, Port: 4747}, // 0
- krpc.NodeAddr{IP: net.IPv6loopback, Port: 4748}, // 1
- krpc.NodeAddr{IP: net.IPv6loopback, Port: 4749}, // 2
- krpc.NodeAddr{IP: net.IPv6loopback, Port: 4750}, // 3
- }
- f := pp.PexPrefersEncryption | pp.PexOutgoingConn
-
t.Run("ipv4", func(t *testing.T) {
addrs := addrs4
var m pexMsgFactory
- m.Add(addrs[0], f)
- m.Drop(addrs[1])
+ m.addEvent(pexEvent{pexAdd, addrs[0], f})
+ m.addEvent(pexEvent{pexDrop, addrs[1], 0})
for _, addr := range addrs {
- m.Drop(addr)
+ m.addEvent(pexEvent{pexDrop, addr, 0})
}
targ := pp.PexMsg{
Dropped: krpc.CompactIPv4NodeAddrs{
- addrs[1],
- addrs[2],
- addrs[3],
+ mustNodeAddr(addrs[1]),
+ mustNodeAddr(addrs[2]),
+ mustNodeAddr(addrs[3]),
},
}
assertPexMsgsEqual(t, targ, m.PexMsg())
@@ -418,24 +406,24 @@ func TestPexDrop(t *testing.T) {
t.Run("ipv6", func(t *testing.T) {
addrs := addrs6
var m pexMsgFactory
- m.Add(addrs[0], f)
- m.Drop(addrs[1])
+ m.addEvent(pexEvent{pexAdd, addrs[0], f})
+ m.addEvent(pexEvent{pexDrop, addrs[1], 0})
for _, addr := range addrs {
- m.Drop(addr)
+ m.addEvent(pexEvent{pexDrop, addr, 0})
}
targ := pp.PexMsg{
Dropped6: krpc.CompactIPv6NodeAddrs{
- addrs[1],
- addrs[2],
- addrs[3],
+ mustNodeAddr(addrs[1]),
+ mustNodeAddr(addrs[2]),
+ mustNodeAddr(addrs[3]),
},
}
assertPexMsgsEqual(t, targ, m.PexMsg())
})
t.Run("empty", func(t *testing.T) {
- addr := krpc.NodeAddr{}
+ nullAddr := &net.TCPAddr{}
var xm pexMsgFactory
- require.Panics(t, func() { xm.Drop(addr) })
+ xm.addEvent(pexEvent{pexDrop, nullAddr, f})
m := xm.PexMsg()
require.EqualValues(t, 0, len(m.Dropped))
require.EqualValues(t, 0, len(m.Dropped6))
--
2.51.0