Share current connections with peers over PEX (anacrolix#341)

[btrtrc.git] / peer_protocol / pex_test.go

diff --git a/peer_protocol/pex_test.go b/peer_protocol/pex_test.go
index 0c2b9f9ed1489008984571bd09502da3902033b0..30e0d363a8cf381046ab662c63650bdc43652ad7 100644 (file)
--- a/peer_protocol/pex_test.go
+++ b/peer_protocol/pex_test.go
@@ -1,10 +1,14 @@
 package peer_protocol
 
 import (
+       "bufio"
+       "bytes"
+       "net"
        "testing"
 
        "github.com/stretchr/testify/require"
 
+       "github.com/anacrolix/dht/v2/krpc"
        "github.com/anacrolix/torrent/bencode"
 )
 
@@ -24,3 +28,102 @@ func TestEmptyPexMsg(t *testing.T) {
        require.NoError(t, err)
        require.NoError(t, bencode.Unmarshal(b, &pm))
 }
+
+func TestPexAppendAdded(t *testing.T) {
+       t.Run("ipv4", func(t *testing.T) {
+               addr := krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747}
+               f := PexPrefersEncryption | PexOutgoingConn
+               xm := PexMsg{}
+               xm.AppendAdded(addr, f)
+               require.EqualValues(t, len(xm.Added), 1)
+               require.EqualValues(t, len(xm.AddedFlags), 1)
+               require.EqualValues(t, len(xm.Added6), 0)
+               require.EqualValues(t, len(xm.Added6Flags), 0)
+               require.True(t, xm.Added[0].IP.Equal(addr.IP), "IPs should match")
+               require.EqualValues(t, xm.Added[0].Port, addr.Port)
+               require.EqualValues(t, xm.AddedFlags[0], f)
+       })
+       t.Run("ipv6", func(t *testing.T) {
+               addr := krpc.NodeAddr{IP: net.IPv6loopback, Port: 4747}
+               f := PexPrefersEncryption | PexOutgoingConn
+               xm := PexMsg{}
+               xm.AppendAdded(addr, f)
+               require.EqualValues(t, len(xm.Added), 0)
+               require.EqualValues(t, len(xm.AddedFlags), 0)
+               require.EqualValues(t, len(xm.Added6), 1)
+               require.EqualValues(t, len(xm.Added6Flags), 1)
+               require.True(t, xm.Added6[0].IP.Equal(addr.IP), "IPs should match")
+               require.EqualValues(t, xm.Added6[0].Port, addr.Port)
+               require.EqualValues(t, xm.Added6Flags[0], f)
+       })
+       t.Run("unspecified", func(t *testing.T) {
+               addr := krpc.NodeAddr{}
+               xm := PexMsg{}
+               xm.AppendAdded(addr, 0)
+               require.EqualValues(t, len(xm.Added), 0)
+               require.EqualValues(t, len(xm.AddedFlags), 0)
+               require.EqualValues(t, len(xm.Added6), 0)
+               require.EqualValues(t, len(xm.Added6Flags), 0)
+       })
+}
+
+func TestPexAppendDropped(t *testing.T) {
+       t.Run("ipv4", func(t *testing.T) {
+               addr := krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747}
+               xm := PexMsg{}
+               xm.AppendDropped(addr)
+               require.EqualValues(t, len(xm.Dropped), 1)
+               require.EqualValues(t, len(xm.Dropped6), 0)
+               require.True(t, xm.Dropped[0].IP.Equal(addr.IP), "IPs should match")
+               require.EqualValues(t, xm.Dropped[0].Port, addr.Port)
+       })
+       t.Run("ipv6", func(t *testing.T) {
+               addr := krpc.NodeAddr{IP: net.IPv6loopback, Port: 4747}
+               xm := PexMsg{}
+               xm.AppendDropped(addr)
+               require.EqualValues(t, len(xm.Dropped), 0)
+               require.EqualValues(t, len(xm.Dropped6), 1)
+               require.True(t, xm.Dropped6[0].IP.Equal(addr.IP), "IPs should match")
+               require.EqualValues(t, xm.Dropped6[0].Port, addr.Port)
+       })
+       t.Run("unspecified", func(t *testing.T) {
+               addr := krpc.NodeAddr{}
+               xm := PexMsg{}
+               xm.AppendDropped(addr)
+               require.EqualValues(t, len(xm.Dropped), 0)
+               require.EqualValues(t, len(xm.Dropped6), 0)
+       })
+}
+
+func TestMarshalPexMessage(t *testing.T) {
+       addr := krpc.NodeAddr{IP: net.IP{127, 0, 0, 1}, Port: 0x55aa}
+       f := PexPrefersEncryption | PexOutgoingConn
+       pm := PexMsg{}
+       pm.AppendAdded(addr, f)
+
+       b, err := bencode.Marshal(pm)
+       require.NoError(t, err)
+
+       pexExtendedId := ExtensionNumber(7)
+       msg := pm.Message(pexExtendedId)
+       expected := []byte("\x00\x00\x00\x4c\x14\x07d5:added6:\x7f\x00\x00\x01\x55\xaa7:added.f1:\x116:added60:8:added6.f0:7:dropped0:8:dropped60:e")
+       b, err = msg.MarshalBinary()
+       require.NoError(t, err)
+       require.EqualValues(t, b, expected)
+
+       msg = Message{}
+       dec := Decoder{
+               R:         bufio.NewReader(bytes.NewBuffer(b)),
+               MaxLength: 128,
+       }
+       pmOut := PexMsg{}
+       err = dec.Decode(&msg)
+       require.NoError(t, err)
+       require.EqualValues(t, Extended, msg.Type)
+       require.EqualValues(t, pexExtendedId, msg.ExtendedID)
+       err = bencode.Unmarshal(msg.ExtendedPayload, &pmOut)
+       require.NoError(t, err)
+       require.EqualValues(t, len(pm.Added), len(pmOut.Added))
+       require.EqualValues(t, pm.Added[0].IP, pmOut.Added[0].IP)
+       require.EqualValues(t, pm.Added[0].Port, pmOut.Added[0].Port)
+}