Add explicit metadata extension types

client.go

@@ -1020,25 +1020,21 @@ func (cl *Client) haveDhtServer() (ret bool) {
 
 // Process incoming ut_metadata message.
 func (cl *Client) gotMetadataExtensionMsg(payload []byte, t *Torrent, c *PeerConn) error {
-	var d map[string]int
+	var d pp.ExtendedMetadataRequestMsg
 	err := bencode.Unmarshal(payload, &d)
 	if _, ok := err.(bencode.ErrUnusedTrailingBytes); ok {
 	} else if err != nil {
 		return fmt.Errorf("error unmarshalling bencode: %s", err)
 	}
-	msgType, ok := d["msg_type"]
+	piece := d.Piece
-	if !ok {
+	switch d.Type {
-		return errors.New("missing msg_type field")
-	}
-	piece := d["piece"]
-	switch msgType {
 	case pp.DataMetadataExtensionMsgType:
 		c.allStats(add(1, func(cs *ConnStats) *Count { return &cs.MetadataChunksRead }))
 		if !c.requestedMetadataPiece(piece) {
 			return fmt.Errorf("got unexpected piece %d", piece)
 		}
 		c.metadataRequests[piece] = false
-		begin := len(payload) - metadataPieceSize(d["total_size"], piece)
+		begin := len(payload) - d.PieceSize()
 		if begin < 0 || begin >= len(payload) {
 			return fmt.Errorf("data has bad offset in payload: %d", begin)
 		}
@@ -1055,7 +1051,7 @@ func (cl *Client) gotMetadataExtensionMsg(payload []byte, t *Torrent, c *PeerCon
 		return err
 	case pp.RequestMetadataExtensionMsgType:
 		if !t.haveMetadataPiece(piece) {
-			c.write(t.newMetadataExtensionMessage(c, pp.RejectMetadataExtensionMsgType, d["piece"], nil))
+			c.write(t.newMetadataExtensionMessage(c, pp.RejectMetadataExtensionMsgType, d.Piece, nil))
 			return nil
 		}
 		start := (1 << 14) * piece

peer_protocol/extended.go

@@ -1,6 +1,8 @@
 package peer_protocol
 
-import "net"
+import (
+	"net"
+)
 
 // http://www.bittorrent.org/beps/bep_0010.html
 type (
@@ -26,7 +28,6 @@ type (
 const (
 	// http://www.bittorrent.org/beps/bep_0011.html
 	ExtensionNamePex ExtensionName = "ut_pex"
-	// http://bittorrent.org/beps/bep_0009.html. Note that there's an
+
-	// LT_metadata, but I've never implemented it.
+	ExtensionDeleteNumber ExtensionNumber = 0
-	ExtensionNameMetadata = "ut_metadata"
 )

peer_protocol/metadata.go

@@ -0,0 +1,42 @@
+package peer_protocol
+
+import (
+	"github.com/anacrolix/torrent/bencode"
+)
+
+const (
+	// http://bittorrent.org/beps/bep_0009.html. Note that there's an
+	// LT_metadata, but I've never implemented it.
+	ExtensionNameMetadata = "ut_metadata"
+)
+
+type (
+	ExtendedMetadataRequestMsg struct {
+		Piece     int                            `bencode:"piece"`
+		TotalSize int                            `bencode:"total_size"`
+		Type      ExtendedMetadataRequestMsgType `bencode:"msg_type"`
+	}
+
+	ExtendedMetadataRequestMsgType int
+)
+
+func MetadataExtensionRequestMsg(peerMetadataExtensionId ExtensionNumber, piece int) Message {
+	return Message{
+		Type:       Extended,
+		ExtendedID: peerMetadataExtensionId,
+		ExtendedPayload: bencode.MustMarshal(ExtendedMetadataRequestMsg{
+			Piece: piece,
+			Type:  RequestMetadataExtensionMsgType,
+		}),
+	}
+}
+
+// Returns the expected piece size for this request message. This is needed to determine the offset
+// into an extension message payload that the request metadata piece data starts.
+func (me ExtendedMetadataRequestMsg) PieceSize() int {
+	ret := me.TotalSize - me.Piece*(1<<14)
+	if ret > 1<<14 {
+		ret = 1 << 14
+	}
+	return ret
+}

peer_protocol/protocol.go

@@ -41,7 +41,7 @@ const (
 const (
 	HandshakeExtendedID = 0
 
-	RequestMetadataExtensionMsgType = 0
+	RequestMetadataExtensionMsgType ExtendedMetadataRequestMsgType = 0
-	DataMetadataExtensionMsgType    = 1
+	DataMetadataExtensionMsgType                                   = 1
-	RejectMetadataExtensionMsgType  = 2
+	RejectMetadataExtensionMsgType                                 = 2
 )

peerconn.go

@@ -443,27 +443,14 @@ func (cn *PeerConn) write(msg pp.Message) bool {
 
 func (cn *PeerConn) requestMetadataPiece(index int) {
 	eID := cn.PeerExtensionIDs[pp.ExtensionNameMetadata]
-	if eID == 0 {
+	if eID == pp.ExtensionDeleteNumber {
 		return
 	}
 	if index < len(cn.metadataRequests) && cn.metadataRequests[index] {
 		return
 	}
 	cn.logger.WithDefaultLevel(log.Debug).Printf("requesting metadata piece %d", index)
-	cn.write(pp.Message{
+	cn.write(pp.MetadataExtensionRequestMsg(eID, index))
-		Type:       pp.Extended,
-		ExtendedID: eID,
-		ExtendedPayload: func() []byte {
-			b, err := bencode.Marshal(map[string]int{
-				"msg_type": pp.RequestMetadataExtensionMsgType,
-				"piece":    index,
-			})
-			if err != nil {
-				panic(err)
-			}
-			return b
-		}(),
-	})
 	for index >= len(cn.metadataRequests) {
 		cn.metadataRequests = append(cn.metadataRequests, false)
 	}

torrent.go

@@ -544,19 +544,15 @@ func (t *Torrent) metadataPieceSize(piece int) int {
 	return metadataPieceSize(len(t.metadataBytes), piece)
 }
 
-func (t *Torrent) newMetadataExtensionMessage(c *PeerConn, msgType int, piece int, data []byte) pp.Message {
+func (t *Torrent) newMetadataExtensionMessage(c *PeerConn, msgType pp.ExtendedMetadataRequestMsgType, piece int, data []byte) pp.Message {
-	d := map[string]int{
-		"msg_type": msgType,
-		"piece":    piece,
-	}
-	if data != nil {
-		d["total_size"] = len(t.metadataBytes)
-	}
-	p := bencode.MustMarshal(d)
 	return pp.Message{
-		Type:            pp.Extended,
+		Type:       pp.Extended,
-		ExtendedID:      c.PeerExtensionIDs[pp.ExtensionNameMetadata],
+		ExtendedID: c.PeerExtensionIDs[pp.ExtensionNameMetadata],
-		ExtendedPayload: append(p, data...),
+		ExtendedPayload: append(bencode.MustMarshal(pp.ExtendedMetadataRequestMsg{
+			Piece:     piece,
+			TotalSize: len(t.metadataBytes),
+			Type:      msgType,
+		}), data...),
 	}
 }