gorond test files

[btrtrc.git] / peer_protocol / pex_test.go

diff --git a/peer_protocol/pex_test.go b/peer_protocol/pex_test.go
index 19cd66bfd8372115c3b43f9aa1ae28d82bd2e98f..5e5e96c14a3f0854f8c366bb63f3dd3b3cb9352a 100644 (file)
--- a/peer_protocol/pex_test.go
+++ b/peer_protocol/pex_test.go
@@ -1,10 +1,15 @@
 package peer_protocol
 
 import (
+       "bufio"
+       "bytes"
+       "net"
        "testing"
 
-       "github.com/anacrolix/torrent/bencode"
+       "github.com/anacrolix/dht/v2/krpc"
        "github.com/stretchr/testify/require"
+
+       "github.com/anacrolix/torrent/bencode"
 )
 
 func TestUnmarshalPex(t *testing.T) {
@@ -23,3 +28,37 @@ func TestEmptyPexMsg(t *testing.T) {
        require.NoError(t, err)
        require.NoError(t, bencode.Unmarshal(b, &pm))
 }
+
+func TestMarshalPexMessage(t *testing.T) {
+       addr := krpc.NodeAddr{IP: net.IP{127, 0, 0, 1}, Port: 0x55aa}
+       f := PexPrefersEncryption | PexOutgoingConn
+       pm := new(PexMsg)
+       pm.Added = append(pm.Added, addr)
+       pm.AddedFlags = append(pm.AddedFlags, f)
+
+       _, 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.NewReader(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)
+}