-
Notifications
You must be signed in to change notification settings - Fork 73
/
extension.go
152 lines (137 loc) · 4.16 KB
/
extension.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
package peerprotocol
import (
"bytes"
"encoding/binary"
"fmt"
"io"
"net"
"github.com/zeebo/bencode"
)
const (
// ExtensionIDHandshake is ID for extension handshake message.
ExtensionIDHandshake = iota
// ExtensionIDMetadata is ID for metadata extension messages.
ExtensionIDMetadata
// ExtensionIDPEX is ID for PEX extension messages.
ExtensionIDPEX
)
const (
// ExtensionKeyMetadata is the key for the metadata extension.
ExtensionKeyMetadata = "ut_metadata"
// ExtensionKeyPEX is the key for the PEX extension.
ExtensionKeyPEX = "ut_pex"
)
const (
// ExtensionMetadataMessageTypeRequest is the id of metadata message when requesting a piece.
ExtensionMetadataMessageTypeRequest = iota
// ExtensionMetadataMessageTypeData is the id of metadata message when sending the piece data.
ExtensionMetadataMessageTypeData
// ExtensionMetadataMessageTypeReject is the id of metadata message when rejecting a piece.
ExtensionMetadataMessageTypeReject
)
// ExtensionMessage is extension to BitTorrent protocol.
type ExtensionMessage struct {
ExtendedMessageID uint8
Payload interface{}
}
// ID returns the type of a peer message.
func (m ExtensionMessage) ID() MessageID { return Extension }
// Read extension message bytes.
func (m ExtensionMessage) Read([]byte) (int, error) {
panic("Read must not be called, use WriteTo")
}
// WriteTo writes the bytes into io.Writer.
func (m ExtensionMessage) WriteTo(w io.Writer) (n int64, err error) {
nn, err := w.Write([]byte{m.ExtendedMessageID})
n += int64(nn)
if err != nil {
return
}
wc := newWriterCounter(w)
err = bencode.NewEncoder(wc).Encode(m.Payload)
n += wc.Count()
if err != nil {
return
}
if mm, ok := m.Payload.(ExtensionMetadataMessage); ok {
nn, err = w.Write(mm.Data)
n += int64(nn)
}
return
}
// UnmarshalBinary parses extension message.
func (m *ExtensionMessage) UnmarshalBinary(data []byte) error {
var extID uint8
r := bytes.NewReader(data)
err := binary.Read(r, binary.BigEndian, &extID)
if err != nil {
return err
}
m.ExtendedMessageID = extID
payload := data[1:]
dec := bencode.NewDecoder(bytes.NewReader(payload))
switch m.ExtendedMessageID {
case ExtensionIDHandshake:
var extMsg ExtensionHandshakeMessage
err = dec.Decode(&extMsg)
m.Payload = extMsg
if extMsg.MetadataSize < 0 {
extMsg.MetadataSize = 0
}
if extMsg.RequestQueue < 0 {
extMsg.RequestQueue = 0
}
case ExtensionIDMetadata:
var extMsg ExtensionMetadataMessage
err = dec.Decode(&extMsg)
extMsg.Data = payload[dec.BytesParsed():]
m.Payload = extMsg
case ExtensionIDPEX:
var extMsg ExtensionPEXMessage
err = dec.Decode(&extMsg)
m.Payload = extMsg
default:
return fmt.Errorf("peer sent invalid extension message id: %d", m.ExtendedMessageID)
}
return err
}
// ExtensionHandshakeMessage contains the information to do the extension handshake.
type ExtensionHandshakeMessage struct {
M map[string]uint8 `bencode:"m"`
V string `bencode:"v"`
YourIP string `bencode:"yourip,omitempty"`
MetadataSize int `bencode:"metadata_size,omitempty"`
RequestQueue int `bencode:"reqq"`
}
// NewExtensionHandshake returns a new ExtensionHandshakeMessage by filling the struct with given values.
func NewExtensionHandshake(metadataSize uint32, version string, yourip net.IP, requestQueueLength int) ExtensionHandshakeMessage {
return ExtensionHandshakeMessage{
M: map[string]uint8{
ExtensionKeyMetadata: ExtensionIDMetadata,
ExtensionKeyPEX: ExtensionIDPEX,
},
V: version,
YourIP: string(truncateIP(yourip)),
MetadataSize: int(metadataSize),
RequestQueue: requestQueueLength,
}
}
// ExtensionMetadataMessage is the message for the Metadata extension.
type ExtensionMetadataMessage struct {
Type int `bencode:"msg_type"`
Piece uint32 `bencode:"piece"`
TotalSize int `bencode:"total_size,omitempty"`
Data []byte `bencode:"-"`
}
// ExtensionPEXMessage is the message for the PEX extension.
type ExtensionPEXMessage struct {
Added string `bencode:"added"`
Dropped string `bencode:"dropped"`
}
func truncateIP(ip net.IP) net.IP {
ip4 := ip.To4()
if ip4 != nil {
return ip4
}
return ip
}