Expose handshake stuff in peer_protocol

client.go

@@ -57,7 +57,7 @@ type Client struct {
 	dhtServers     []*dht.Server
 	ipBlockList    iplist.Ranger
 	// Our BitTorrent protocol extension bytes, sent in our BT handshakes.
-	extensionBytes peerExtensionBytes
+	extensionBytes pp.PeerExtensionBits
 
 	// Set of addresses that have our client ID. This intentionally will
 	// include ourselves if we end up trying to connect to our own address
@@ -730,12 +730,12 @@ func (cl *Client) receiveHandshakes(c *connection) (t *Torrent, err error) {
 
 // Returns !ok if handshake failed for valid reasons.
 func (cl *Client) connBTHandshake(c *connection, ih *metainfo.Hash) (ret metainfo.Hash, ok bool, err error) {
-	res, ok, err := handshake(c.rw(), ih, cl.peerID, cl.extensionBytes)
+	res, ok, err := pp.Handshake(c.rw(), ih, cl.peerID, cl.extensionBytes)
 	if err != nil || !ok {
 		return
 	}
 	ret = res.Hash
-	c.PeerExtensionBytes = res.peerExtensionBytes
+	c.PeerExtensionBytes = res.PeerExtensionBits
 	c.PeerID = res.PeerID
 	c.completedHandshake = time.Now()
 	return

connection.go

@@ -90,7 +90,7 @@ type connection struct {
 	PeerInterested     bool
 	PeerChoked         bool
 	PeerRequests       map[request]struct{}
-	PeerExtensionBytes peerExtensionBytes
+	PeerExtensionBytes pp.PeerExtensionBits
 	// The pieces the peer has claimed to have.
 	peerPieces bitmap.Bitmap
 	// The peer has everything. This can occur due to a special message, when

global.go

@@ -1,6 +1,7 @@
 package torrent
 
 import (
+	pp "github.com/anacrolix/torrent/peer_protocol"
 	"crypto"
 	"expvar"
 )
@@ -16,8 +17,8 @@ const (
 	pexExtendedId
 )
 
-func defaultPeerExtensionBytes() peerExtensionBytes {
-	return newPeerExtensionBytes(ExtensionBitDHT, ExtensionBitExtended, ExtensionBitFast)
+func defaultPeerExtensionBytes() PeerExtensionBits {
+	return pp.NewPeerExtensionBytes(pp.ExtensionBitDHT, pp.ExtensionBitExtended, pp.ExtensionBitFast)
 }
 
 // I could move a lot of these counters to their own file, but I suspect they

handshake.go

@@ -2,146 +2,15 @@ package torrent
 
 import (
 	"bytes"
-	"encoding/hex"
 	"fmt"
 	"io"
 	"net"
 	"time"
 
-	"github.com/anacrolix/missinggo"
-
-	"github.com/anacrolix/torrent/metainfo"
 	"github.com/anacrolix/torrent/mse"
 	pp "github.com/anacrolix/torrent/peer_protocol"
 )
 
-type ExtensionBit uint
-
-const (
-	ExtensionBitDHT      = 0  // http://www.bittorrent.org/beps/bep_0005.html
-	ExtensionBitExtended = 20 // http://www.bittorrent.org/beps/bep_0010.html
-	ExtensionBitFast     = 2  // http://www.bittorrent.org/beps/bep_0006.html
-)
-
-func handshakeWriter(w io.Writer, bb <-chan []byte, done chan<- error) {
-	var err error
-	for b := range bb {
-		_, err = w.Write(b)
-		if err != nil {
-			break
-		}
-	}
-	done <- err
-}
-
-type (
-	peerExtensionBytes [8]byte
-)
-
-func (me peerExtensionBytes) String() string {
-	return hex.EncodeToString(me[:])
-}
-
-func newPeerExtensionBytes(bits ...ExtensionBit) (ret peerExtensionBytes) {
-	for _, b := range bits {
-		ret.SetBit(b)
-	}
-	return
-}
-
-func (pex peerExtensionBytes) SupportsExtended() bool {
-	return pex.GetBit(ExtensionBitExtended)
-}
-
-func (pex peerExtensionBytes) SupportsDHT() bool {
-	return pex.GetBit(ExtensionBitDHT)
-}
-
-func (pex peerExtensionBytes) SupportsFast() bool {
-	return pex.GetBit(ExtensionBitFast)
-}
-
-func (pex *peerExtensionBytes) SetBit(bit ExtensionBit) {
-	pex[7-bit/8] |= 1 << (bit % 8)
-}
-
-func (pex peerExtensionBytes) GetBit(bit ExtensionBit) bool {
-	return pex[7-bit/8]&(1<<(bit%8)) != 0
-}
-
-type handshakeResult struct {
-	peerExtensionBytes
-	PeerID
-	metainfo.Hash
-}
-
-// ih is nil if we expect the peer to declare the InfoHash, such as when the
-// peer initiated the connection. Returns ok if the handshake was successful,
-// and err if there was an unexpected condition other than the peer simply
-// abandoning the handshake.
-func handshake(sock io.ReadWriter, ih *metainfo.Hash, peerID [20]byte, extensions peerExtensionBytes) (res handshakeResult, ok bool, err error) {
-	// Bytes to be sent to the peer. Should never block the sender.
-	postCh := make(chan []byte, 4)
-	// A single error value sent when the writer completes.
-	writeDone := make(chan error, 1)
-	// Performs writes to the socket and ensures posts don't block.
-	go handshakeWriter(sock, postCh, writeDone)
-
-	defer func() {
-		close(postCh) // Done writing.
-		if !ok {
-			return
-		}
-		if err != nil {
-			panic(err)
-		}
-		// Wait until writes complete before returning from handshake.
-		err = <-writeDone
-		if err != nil {
-			err = fmt.Errorf("error writing: %s", err)
-		}
-	}()
-
-	post := func(bb []byte) {
-		select {
-		case postCh <- bb:
-		default:
-			panic("mustn't block while posting")
-		}
-	}
-
-	post([]byte(pp.Protocol))
-	post(extensions[:])
-	if ih != nil { // We already know what we want.
-		post(ih[:])
-		post(peerID[:])
-	}
-	var b [68]byte
-	_, err = io.ReadFull(sock, b[:68])
-	if err != nil {
-		err = nil
-		return
-	}
-	if string(b[:20]) != pp.Protocol {
-		return
-	}
-	missinggo.CopyExact(&res.peerExtensionBytes, b[20:28])
-	missinggo.CopyExact(&res.Hash, b[28:48])
-	missinggo.CopyExact(&res.PeerID, b[48:68])
-	peerExtensions.Add(res.peerExtensionBytes.String(), 1)
-
-	// TODO: Maybe we can just drop peers here if we're not interested. This
-	// could prevent them trying to reconnect, falsely believing there was
-	// just a problem.
-	if ih == nil { // We were waiting for the peer to tell us what they wanted.
-		post(res.Hash[:])
-		post(peerID[:])
-	}
-
-	ok = true
-	return
-}
-
 // Wraps a raw connection and provides the interface we want for using the
 // connection in the message loop.
 type deadlineReader struct {
@@ -201,3 +70,5 @@ func handleEncryption(
 	})
 	return
 }
+
+type PeerExtensionBits=pp.PeerExtensionBits

peer_protocol/handshake.go

@@ -0,0 +1,137 @@
+package peer_protocol
+
+import (
+	"encoding/hex"
+	"fmt"
+	"io"
+
+	"github.com/anacrolix/missinggo"
+	"github.com/anacrolix/torrent/metainfo"
+)
+
+type ExtensionBit uint
+
+const (
+	ExtensionBitDHT      = 0  // http://www.bittorrent.org/beps/bep_0005.html
+	ExtensionBitExtended = 20 // http://www.bittorrent.org/beps/bep_0010.html
+	ExtensionBitFast     = 2  // http://www.bittorrent.org/beps/bep_0006.html
+)
+
+func handshakeWriter(w io.Writer, bb <-chan []byte, done chan<- error) {
+	var err error
+	for b := range bb {
+		_, err = w.Write(b)
+		if err != nil {
+			break
+		}
+	}
+	done <- err
+}
+
+type (
+	PeerExtensionBits [8]byte
+)
+
+func (me PeerExtensionBits) String() string {
+	return hex.EncodeToString(me[:])
+}
+
+func NewPeerExtensionBytes(bits ...ExtensionBit) (ret PeerExtensionBits) {
+	for _, b := range bits {
+		ret.SetBit(b)
+	}
+	return
+}
+
+func (pex PeerExtensionBits) SupportsExtended() bool {
+	return pex.GetBit(ExtensionBitExtended)
+}
+
+func (pex PeerExtensionBits) SupportsDHT() bool {
+	return pex.GetBit(ExtensionBitDHT)
+}
+
+func (pex PeerExtensionBits) SupportsFast() bool {
+	return pex.GetBit(ExtensionBitFast)
+}
+
+func (pex *PeerExtensionBits) SetBit(bit ExtensionBit) {
+	pex[7-bit/8] |= 1 << (bit % 8)
+}
+
+func (pex PeerExtensionBits) GetBit(bit ExtensionBit) bool {
+	return pex[7-bit/8]&(1<<(bit%8)) != 0
+}
+
+type HandshakeResult struct {
+	PeerExtensionBits
+	PeerID [20]byte
+	metainfo.Hash
+}
+
+// ih is nil if we expect the peer to declare the InfoHash, such as when the
+// peer initiated the connection. Returns ok if the Handshake was successful,
+// and err if there was an unexpected condition other than the peer simply
+// abandoning the Handshake.
+func Handshake(sock io.ReadWriter, ih *metainfo.Hash, peerID [20]byte, extensions PeerExtensionBits) (res HandshakeResult, ok bool, err error) {
+	// Bytes to be sent to the peer. Should never block the sender.
+	postCh := make(chan []byte, 4)
+	// A single error value sent when the writer completes.
+	writeDone := make(chan error, 1)
+	// Performs writes to the socket and ensures posts don't block.
+	go handshakeWriter(sock, postCh, writeDone)
+
+	defer func() {
+		close(postCh) // Done writing.
+		if !ok {
+			return
+		}
+		if err != nil {
+			panic(err)
+		}
+		// Wait until writes complete before returning from handshake.
+		err = <-writeDone
+		if err != nil {
+			err = fmt.Errorf("error writing: %s", err)
+		}
+	}()
+
+	post := func(bb []byte) {
+		select {
+		case postCh <- bb:
+		default:
+			panic("mustn't block while posting")
+		}
+	}
+
+	post([]byte(Protocol))
+	post(extensions[:])
+	if ih != nil { // We already know what we want.
+		post(ih[:])
+		post(peerID[:])
+	}
+	var b [68]byte
+	_, err = io.ReadFull(sock, b[:68])
+	if err != nil {
+		err = nil
+		return
+	}
+	if string(b[:20]) != Protocol {
+		return
+	}
+	missinggo.CopyExact(&res.PeerExtensionBits, b[20:28])
+	missinggo.CopyExact(&res.Hash, b[28:48])
+	missinggo.CopyExact(&res.PeerID, b[48:68])
+	// peerExtensions.Add(res.PeerExtensionBits.String(), 1) 
+
+	// TODO: Maybe we can just drop peers here if we're not interested. This
+	// could prevent them trying to reconnect, falsely believing there was
+	// just a problem.
+	if ih == nil { // We were waiting for the peer to tell us what they wanted.
+		post(res.Hash[:])
+		post(peerID[:])
+	}
+
+	ok = true
+	return
+}

torrent_test.go

@@ -182,12 +182,12 @@ func TestTorrentMetainfoIncompleteMetadata(t *testing.T) {
 	require.NoError(t, err)
 	defer nc.Close()
 
-	var pex peerExtensionBytes
-	pex.SetBit(ExtensionBitExtended)
-	hr, ok, err := handshake(nc, &ih, [20]byte{}, pex)
+	var pex PeerExtensionBits
+	pex.SetBit(pp.ExtensionBitExtended)
+	hr, ok, err := pp.Handshake(nc, &ih, [20]byte{}, pex)
 	require.NoError(t, err)
 	assert.True(t, ok)
-	assert.True(t, hr.peerExtensionBytes.GetBit(ExtensionBitExtended))
+	assert.True(t, hr.PeerExtensionBits.GetBit(pp.ExtensionBitExtended))
 	assert.EqualValues(t, cl.PeerID(), hr.PeerID)
 	assert.Equal(t, ih, hr.Hash)