135 lines
2.6 KiB
Go
135 lines
2.6 KiB
Go
package client
|
|
|
|
import (
|
|
"crypto/sha1"
|
|
|
|
"storrent/bt"
|
|
"storrent/metainfo"
|
|
)
|
|
|
|
type pieceScheduler struct {
|
|
pieces []piece
|
|
incomplete []int
|
|
verified int
|
|
meta *metainfo.File
|
|
dir string
|
|
}
|
|
|
|
func newPieceScheduler(m *metainfo.File, dir string) *pieceScheduler {
|
|
n := len(m.Info.Pieces)
|
|
pieces := make([]piece, n)
|
|
incomplete := make([]int, 0, n)
|
|
verified := 0
|
|
|
|
for i := range n {
|
|
size := m.Info.PieceSize
|
|
if i == n-1 {
|
|
size = m.Size - int64(n-1)*m.Info.PieceSize
|
|
}
|
|
nblocks := (size + bt.BlockSize - 1) / bt.BlockSize
|
|
pieces[i] = piece{
|
|
hash: m.Info.Pieces[i],
|
|
size: size,
|
|
data: make([]byte, size),
|
|
have: make([]bool, nblocks),
|
|
reqs: make([]bool, nblocks),
|
|
reqPeer: make([]*bt.Peer, nblocks),
|
|
}
|
|
|
|
data := m.Info.Read(dir, i, pieces[i].size)
|
|
if data != nil && sha1.Sum(data) == pieces[i].hash {
|
|
pieces[i].data = data
|
|
pieces[i].done = true
|
|
verified++
|
|
} else {
|
|
incomplete = append(incomplete, i)
|
|
}
|
|
}
|
|
|
|
return &pieceScheduler{
|
|
pieces: pieces,
|
|
incomplete: incomplete,
|
|
verified: verified,
|
|
meta: m,
|
|
dir: dir,
|
|
}
|
|
}
|
|
|
|
func (ps *pieceScheduler) nextRequest(p *bt.Peer) *bt.Msg {
|
|
for _, idx := range ps.incomplete {
|
|
pc := &ps.pieces[idx]
|
|
if pc.done {
|
|
continue
|
|
}
|
|
if !p.HasPiece(idx) {
|
|
continue
|
|
}
|
|
if begin, length, ok := pc.next(p); ok {
|
|
return &bt.Msg{
|
|
Kind: bt.Request,
|
|
Index: uint32(idx),
|
|
Begin: uint32(begin),
|
|
Length: uint32(length),
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (ps *pieceScheduler) put(index int, begin int, block []byte) bool {
|
|
pc := &ps.pieces[index]
|
|
if pc.put(begin, block) {
|
|
ps.meta.Info.Write(ps.dir, index, pc.data)
|
|
ps.verified++
|
|
ps.removeIncomplete(index)
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (ps *pieceScheduler) removeIncomplete(index int) {
|
|
for i, idx := range ps.incomplete {
|
|
if idx == index {
|
|
ps.incomplete = append(ps.incomplete[:i], ps.incomplete[i+1:]...)
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (ps *pieceScheduler) peerDisconnected(p *bt.Peer) {
|
|
for i := range ps.pieces {
|
|
ps.pieces[i].resetPeer(p)
|
|
}
|
|
}
|
|
|
|
func (ps *pieceScheduler) isComplete() bool {
|
|
return ps.verified == len(ps.pieces)
|
|
}
|
|
|
|
func (ps *pieceScheduler) progress() (verified, total int) {
|
|
return ps.verified, len(ps.pieces)
|
|
}
|
|
|
|
func (ps *pieceScheduler) pieceData(index int) []byte {
|
|
pc := &ps.pieces[index]
|
|
if !pc.done {
|
|
return nil
|
|
}
|
|
return pc.data
|
|
}
|
|
|
|
func (ps *pieceScheduler) count() int {
|
|
return len(ps.pieces)
|
|
}
|
|
|
|
func (ps *pieceScheduler) bitfield(all bool) []byte {
|
|
n := len(ps.pieces)
|
|
bf := make([]byte, (n+7)/8)
|
|
for i := range ps.pieces {
|
|
if all || ps.pieces[i].done {
|
|
bt.SetBit(bf, i)
|
|
}
|
|
}
|
|
return bf
|
|
}
|