117 lines
2.7 KiB
Go
117 lines
2.7 KiB
Go
package bt
|
|
|
|
// https://wiki.theory.org/BitTorrentSpecification#Messages
|
|
// Messages: <length prefix><message ID><payload>
|
|
// keep-alive: <len=0000>
|
|
// choke: <len=0001><type=0>
|
|
// unchoke: <len=0001><type=1>
|
|
// interested: <len=0001><type=2>
|
|
// not interested: <len=0001><type=3>
|
|
// have: <len=0005><type=4><piece index>
|
|
// bitfield: <len=0001+X><type=5><bitfield>
|
|
// request: <len=0013><type=6><index><begin><length>
|
|
// piece: <len=0009+X><type=7><index><begin><block>
|
|
// cancel: <len=0013><type=8><index><begin><length>
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"io"
|
|
)
|
|
|
|
type MsgKind uint8
|
|
|
|
const BlockSize = 16384
|
|
|
|
const (
|
|
Choke MsgKind = iota
|
|
Unchoke
|
|
Interested
|
|
NotInterested
|
|
Have
|
|
Bitfield
|
|
Request
|
|
Piece
|
|
Cancel
|
|
)
|
|
|
|
type Msg struct {
|
|
Kind MsgKind
|
|
Index uint32
|
|
Begin uint32
|
|
Length uint32
|
|
Bitfield []byte
|
|
Block []byte
|
|
}
|
|
|
|
// [4: len] [1: type] [len-1: payload]
|
|
func readMsg(r io.Reader) (*Msg, int, error) {
|
|
var lenBuf [4]byte
|
|
if _, err := io.ReadFull(r, lenBuf[:]); err != nil {
|
|
return nil, 0, err
|
|
}
|
|
n := binary.BigEndian.Uint32(lenBuf[:])
|
|
if n == 0 {
|
|
return nil, 4, nil
|
|
}
|
|
buf := make([]byte, n)
|
|
if _, err := io.ReadFull(r, buf); err != nil {
|
|
return nil, 0, err
|
|
}
|
|
msg := &Msg{Kind: MsgKind(buf[0])}
|
|
payload := buf[1:]
|
|
switch msg.Kind {
|
|
case Have:
|
|
msg.Index = binary.BigEndian.Uint32(payload)
|
|
case Bitfield:
|
|
msg.Bitfield = payload
|
|
case Request, Cancel:
|
|
msg.Index = binary.BigEndian.Uint32(payload[0:4])
|
|
msg.Begin = binary.BigEndian.Uint32(payload[4:8])
|
|
msg.Length = binary.BigEndian.Uint32(payload[8:12])
|
|
case Piece:
|
|
msg.Index = binary.BigEndian.Uint32(payload[0:4])
|
|
msg.Begin = binary.BigEndian.Uint32(payload[4:8])
|
|
msg.Block = payload[8:]
|
|
}
|
|
return msg, int(4 + n), nil
|
|
}
|
|
|
|
func writeMsg(w io.Writer, msg *Msg) (int, error) {
|
|
if msg == nil {
|
|
return w.Write(make([]byte, 4))
|
|
}
|
|
|
|
var payload []byte
|
|
switch msg.Kind {
|
|
case Choke, Unchoke, Interested, NotInterested:
|
|
case Have:
|
|
payload = make([]byte, 4)
|
|
binary.BigEndian.PutUint32(payload, msg.Index)
|
|
case Bitfield:
|
|
payload = msg.Bitfield
|
|
case Request, Cancel:
|
|
payload = make([]byte, 12)
|
|
binary.BigEndian.PutUint32(payload[0:4], msg.Index)
|
|
binary.BigEndian.PutUint32(payload[4:8], msg.Begin)
|
|
binary.BigEndian.PutUint32(payload[8:12], msg.Length)
|
|
case Piece:
|
|
payload = make([]byte, 8+len(msg.Block))
|
|
binary.BigEndian.PutUint32(payload[0:4], msg.Index)
|
|
binary.BigEndian.PutUint32(payload[4:8], msg.Begin)
|
|
copy(payload[8:], msg.Block)
|
|
}
|
|
buf := make([]byte, 5+len(payload))
|
|
binary.BigEndian.PutUint32(buf[0:4], uint32(1+len(payload)))
|
|
buf[4] = byte(msg.Kind)
|
|
copy(buf[5:], payload)
|
|
return w.Write(buf)
|
|
}
|
|
|
|
func SetBit(bf []byte, i int) {
|
|
bf[i/8] |= 1 << (7 - i%8)
|
|
}
|
|
|
|
func HasBit(bf []byte, i int) bool {
|
|
return bf[i/8]&(1<<(7-i%8)) != 0
|
|
}
|