github.com/uber/kraken@v0.1.4/lib/torrent/scheduler/conn/fake_peer.go (about)

     1  package conn
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"strconv"
     7  	"time"
     8  
     9  	"github.com/uber/kraken/core"
    10  	"github.com/uber/kraken/utils/log"
    11  
    12  	"github.com/willf/bitset"
    13  )
    14  
    15  // FakePeer is a testing utility which reciprocates handshakes against
    16  // arbitrary incoming connections, parroting back the requested torrent but
    17  // with an empty bitfield (so no pieces are requested).
    18  //
    19  // Useful for initializing real Conns against a motionless peer.
    20  type FakePeer struct {
    21  	listener net.Listener
    22  
    23  	id   core.PeerID
    24  	ip   string
    25  	port int
    26  
    27  	msgTimeout time.Duration
    28  }
    29  
    30  // NewFakePeer creates and starts a new FakePeer.
    31  func NewFakePeer() (*FakePeer, error) {
    32  	l, err := net.Listen("tcp", "localhost:0")
    33  	if err != nil {
    34  		return nil, err
    35  	}
    36  	ip, portStr, err := net.SplitHostPort(l.Addr().String())
    37  	if err != nil {
    38  		return nil, err
    39  	}
    40  	port, err := strconv.Atoi(portStr)
    41  	if err != nil {
    42  		return nil, err
    43  	}
    44  	p := &FakePeer{
    45  		listener:   l,
    46  		id:         core.PeerIDFixture(),
    47  		ip:         ip,
    48  		port:       port,
    49  		msgTimeout: 5 * time.Second,
    50  	}
    51  	go func() {
    52  		err := p.handshakeConns()
    53  		log.Infof("Fake peer exiting: %s", err)
    54  	}()
    55  	return p, nil
    56  }
    57  
    58  // PeerID returns the peer's PeerID.
    59  func (p *FakePeer) PeerID() core.PeerID {
    60  	return p.id
    61  }
    62  
    63  // Addr returns the ip:port of the peer.
    64  func (p *FakePeer) Addr() string {
    65  	return fmt.Sprintf("%s:%d", p.ip, p.port)
    66  }
    67  
    68  // PeerInfo returns the peers' PeerInfo.
    69  func (p *FakePeer) PeerInfo() *core.PeerInfo {
    70  	return core.NewPeerInfo(p.id, p.ip, p.port, false, false)
    71  }
    72  
    73  // Close shuts down the peer.
    74  func (p *FakePeer) Close() {
    75  	p.listener.Close()
    76  }
    77  
    78  func (p *FakePeer) handshakeConns() error {
    79  	for {
    80  		nc, err := p.listener.Accept()
    81  		if err != nil {
    82  			return err
    83  		}
    84  		reqMsg, err := readMessageWithTimeout(nc, p.msgTimeout)
    85  		if err != nil {
    86  			return err
    87  		}
    88  		req, err := handshakeFromP2PMessage(reqMsg)
    89  		if err != nil {
    90  			return err
    91  		}
    92  		resp := &handshake{
    93  			peerID:   p.id,
    94  			digest:   req.digest,
    95  			infoHash: req.infoHash,
    96  			// Oh darn, we have no pieces!
    97  			bitfield:  bitset.New(req.bitfield.Len()),
    98  			namespace: req.namespace,
    99  		}
   100  		respMsg, err := resp.toP2PMessage()
   101  		if err != nil {
   102  			return err
   103  		}
   104  		if err := sendMessageWithTimeout(nc, respMsg, p.msgTimeout); err != nil {
   105  			return err
   106  		}
   107  	}
   108  }