path: root/p2p/protocol_test.go

package p2p

import (
    "fmt"
    "net"
    "reflect"
    "sync"
    "testing"

    "github.com/ethereum/go-ethereum/crypto"
)

type peerId struct {
    privKey, pubkey []byte
}

func (self *peerId) String() string {
    return fmt.Sprintf("test peer %x", self.Pubkey()[:4])
}

func (self *peerId) Pubkey() (pubkey []byte) {
    pubkey = self.pubkey
    if len(pubkey) == 0 {
        pubkey = crypto.GenerateNewKeyPair().PublicKey
        self.pubkey = pubkey
    }
    return
}

func (self *peerId) PrivKey() (privKey []byte) {
    privKey = self.privKey
    if len(privKey) == 0 {
        privKey = crypto.GenerateNewKeyPair().PublicKey
        self.privKey = privKey
    }
    return
}

func newTestPeer() (peer *Peer) {
    peer = NewPeer(&peerId{}, []Cap{})
    peer.pubkeyHook = func(*peerAddr) error { return nil }
    peer.ourID = &peerId{}
    peer.listenAddr = &peerAddr{}
    peer.otherPeers = func() []*Peer { return nil }
    return
}

func TestBaseProtocolPeers(t *testing.T) {
    peerList := []*peerAddr{
        {IP: net.ParseIP("1.2.3.4"), Port: 2222, Pubkey: []byte{}},
        {IP: net.ParseIP("5.6.7.8"), Port: 3333, Pubkey: []byte{}},
    }
    listenAddr := &peerAddr{IP: net.ParseIP("1.3.5.7"), Port: 1111, Pubkey: []byte{}}
    rw1, rw2 := MsgPipe()
    defer rw1.Close()
    wg := new(sync.WaitGroup)

    // run matcher, close pipe when addresses have arrived
    numPeers := len(peerList) + 1
    addrChan := make(chan *peerAddr)
    wg.Add(1)
    go func() {
        i := 0
        for got := range addrChan {
            var want *peerAddr
            switch {
            case i < len(peerList):
                want = peerList[i]
            case i == len(peerList):
                want = listenAddr // listenAddr should be the last thing sent
            }
            t.Logf("got peer %d/%d: %v", i+1, numPeers, got)
            if !reflect.DeepEqual(want, got) {
                t.Errorf("mismatch: got %+v, want %+v", got, want)
            }
            i++
            if i == numPeers {
                break
            }
        }
        if i != numPeers {
            t.Errorf("wrong number of peers received: got %d, want %d", i, numPeers)
        }
        rw1.Close()
        wg.Done()
    }()

    // run first peer (in background)
    peer1 := newTestPeer()
    peer1.ourListenAddr = listenAddr
    peer1.otherPeers = func() []*Peer {
        pl := make([]*Peer, len(peerList))
        for i, addr := range peerList {
            pl[i] = &Peer{listenAddr: addr}
        }
        return pl
    }
    wg.Add(1)
    go func() {
        runBaseProtocol(peer1, rw1)
        wg.Done()
    }()

    // run second peer
    peer2 := newTestPeer()
    peer2.newPeerAddr = addrChan // feed peer suggestions into matcher
    if err := runBaseProtocol(peer2, rw2); err != ErrPipeClosed {
        t.Errorf("peer2 terminated with unexpected error: %v", err)
    }

    // terminate matcher
    close(addrChan)
    wg.Wait()
}

func TestBaseProtocolDisconnect(t *testing.T) {
    peer := NewPeer(&peerId{}, nil)
    peer.ourID = &peerId{}
    peer.pubkeyHook = func(*peerAddr) error { return nil }

    rw1, rw2 := MsgPipe()
    done := make(chan struct{})
    go func() {
        if err := expectMsg(rw2, handshakeMsg); err != nil {
            t.Error(err)
        }
        err := EncodeMsg(rw2, handshakeMsg,
            baseProtocolVersion,
            "",
            []interface{}{},
            0,
            make([]byte, 64),
        )
        if err != nil {
            t.Error(err)
        }
        if err := expectMsg(rw2, getPeersMsg); err != nil {
            t.Error(err)
        }
        if err := EncodeMsg(rw2, discMsg, DiscQuitting); err != nil {
            t.Error(err)
        }

        close(done)
    }()

    if err := runBaseProtocol(peer, rw1); err == nil {
        t.Errorf("base protocol returned without error")
    } else if reason, ok := err.(discRequestedError); !ok || reason != DiscQuitting {
        t.Errorf("base protocol returned wrong error: %v", err)
    }
    <-done
}

func expectMsg(r MsgReader, code uint64) error {
    msg, err := r.ReadMsg()
    if err != nil {
        return err
    }
    if err := msg.Discard(); err != nil {
        return err
    }
    if msg.Code != code {
        return fmt.Errorf("wrong message code: got %d, expected %d", msg.Code, code)
    }
    return nil
}