Ztorrent/internal/torrent/peerwire.go

package torrent

import (
	"bytes"
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"time"

	"github.com/veggiedefender/torrent-client/internal/torrentfile"
)

const (
	wireProtocolString = "BitTorrent protocol"

	msgChoke         = 0
	msgUnchoke       = 1
	msgInterested    = 2
	msgNotInterested = 3
	msgHave          = 4
	msgBitfield      = 5
	msgRequest       = 6
	msgPiece         = 7

	maxWireMessageSize = 2 * 1024 * 1024
	requestBlockSize   = 16 * 1024

	peerConnectTimeout = 8 * time.Second
	peerReadTimeout    = 15 * time.Second
	peerWriteTimeout   = 10 * time.Second
)

type wireMessage struct {
	ID      int
	Payload []byte
}

type peerClient struct {
	conn net.Conn

	have         []bool
	hasPieceInfo bool
	peerIsChoked bool
}

func newPeerClient(ctx context.Context, addr string, infoHash [20]byte, peerID [20]byte, pieceCount int) (*peerClient, error) {
	dialer := net.Dialer{Timeout: peerConnectTimeout}
	conn, err := dialer.DialContext(ctx, "tcp", addr)
	if err != nil {
		return nil, err
	}

	pc := &peerClient{
		conn:         conn,
		have:         make([]bool, pieceCount),
		peerIsChoked: true,
	}

	if err := pc.sendHandshake(ctx, infoHash, peerID); err != nil {
		conn.Close()
		return nil, err
	}
	if err := pc.readHandshake(ctx, infoHash); err != nil {
		conn.Close()
		return nil, err
	}

	if err := pc.sendMessage(ctx, msgInterested, nil); err != nil {
		conn.Close()
		return nil, err
	}

	if err := pc.readInitialMessages(ctx); err != nil {
		conn.Close()
		return nil, err
	}

	return pc, nil
}

func (pc *peerClient) Close() error {
	return pc.conn.Close()
}

func (pc *peerClient) PieceAvailability() ([]bool, bool) {
	have := make([]bool, len(pc.have))
	copy(have, pc.have)
	return have, pc.hasPieceInfo
}

func (pc *peerClient) DownloadPiece(ctx context.Context, pieceIndex int, pieceLength int) ([]byte, error) {
	if pieceLength <= 0 {
		return nil, fmt.Errorf("invalid piece length %d", pieceLength)
	}

	if err := pc.waitForUnchoke(ctx); err != nil {
		return nil, err
	}

	piece := make([]byte, pieceLength)
	for offset := 0; offset < pieceLength; {
		blockLength := requestBlockSize
		if remaining := pieceLength - offset; remaining < blockLength {
			blockLength = remaining
		}

		if err := pc.sendRequest(ctx, pieceIndex, offset, blockLength); err != nil {
			return nil, err
		}

		block, err := pc.readPieceBlock(ctx, pieceIndex, offset, blockLength)
		if err != nil {
			return nil, err
		}
		copy(piece[offset:], block)
		offset += len(block)
	}

	return piece, nil
}

func (pc *peerClient) readInitialMessages(ctx context.Context) error {
	if err := pc.setReadDeadlineFromContext(ctx, 2*time.Second); err != nil {
		return err
	}
	defer pc.conn.SetReadDeadline(time.Time{})

	for {
		msg, err := readWireMessage(pc.conn)
		if err != nil {
			if isTimeout(err) {
				return nil
			}
			return err
		}
		pc.consumeMessage(msg)
	}
}

func (pc *peerClient) waitForUnchoke(ctx context.Context) error {
	if !pc.peerIsChoked {
		return nil
	}

	deadline := time.Now().Add(12 * time.Second)
	for {
		if err := ctx.Err(); err != nil {
			return err
		}
		if time.Now().After(deadline) {
			return errors.New("peer did not unchoke")
		}

		if err := pc.setReadDeadlineFromContext(ctx, peerReadTimeout); err != nil {
			return err
		}
		msg, err := readWireMessage(pc.conn)
		if err != nil {
			if isTimeout(err) {
				continue
			}
			return err
		}
		pc.consumeMessage(msg)
		if msg.ID == msgUnchoke {
			return nil
		}
	}
}

func (pc *peerClient) sendRequest(ctx context.Context, pieceIndex, begin, length int) error {
	payload := make([]byte, 12)
	binary.BigEndian.PutUint32(payload[0:4], uint32(pieceIndex))
	binary.BigEndian.PutUint32(payload[4:8], uint32(begin))
	binary.BigEndian.PutUint32(payload[8:12], uint32(length))
	return pc.sendMessage(ctx, msgRequest, payload)
}

func (pc *peerClient) readPieceBlock(ctx context.Context, pieceIndex, begin, expectedLen int) ([]byte, error) {
	for {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		if err := pc.setReadDeadlineFromContext(ctx, peerReadTimeout); err != nil {
			return nil, err
		}

		msg, err := readWireMessage(pc.conn)
		if err != nil {
			if isTimeout(err) {
				continue
			}
			return nil, err
		}

		switch msg.ID {
		case msgPiece:
			if len(msg.Payload) < 8 {
				continue
			}
			gotIndex := int(binary.BigEndian.Uint32(msg.Payload[0:4]))
			gotBegin := int(binary.BigEndian.Uint32(msg.Payload[4:8]))
			block := msg.Payload[8:]
			if gotIndex != pieceIndex || gotBegin != begin {
				pc.consumeMessage(msg)
				continue
			}
			if len(block) == 0 {
				continue
			}
			if len(block) > expectedLen {
				block = block[:expectedLen]
			}
			return block, nil
		case msgChoke:
			pc.consumeMessage(msg)
			return nil, errors.New("peer choked")
		default:
			pc.consumeMessage(msg)
		}
	}
}

func (pc *peerClient) consumeMessage(msg wireMessage) {
	switch msg.ID {
	case msgChoke:
		pc.peerIsChoked = true
	case msgUnchoke:
		pc.peerIsChoked = false
	case msgHave:
		if len(msg.Payload) < 4 {
			return
		}
		idx := int(binary.BigEndian.Uint32(msg.Payload[:4]))
		if idx >= 0 && idx < len(pc.have) {
			pc.have[idx] = true
			pc.hasPieceInfo = true
		}
	case msgBitfield:
		pc.hasPieceInfo = true
		for i := range pc.have {
			pc.have[i] = bitfieldHasPiece(msg.Payload, i)
		}
	}
}

func (pc *peerClient) sendHandshake(ctx context.Context, infoHash [20]byte, peerID [20]byte) error {
	payload := make([]byte, 49+len(wireProtocolString))
	payload[0] = byte(len(wireProtocolString))
	copy(payload[1:1+len(wireProtocolString)], wireProtocolString)
	copy(payload[1+len(wireProtocolString)+8:1+len(wireProtocolString)+8+20], infoHash[:])
	copy(payload[1+len(wireProtocolString)+8+20:], peerID[:])

	if err := pc.setWriteDeadlineFromContext(ctx, peerWriteTimeout); err != nil {
		return err
	}
	_, err := pc.conn.Write(payload)
	return err
}

func (pc *peerClient) readHandshake(ctx context.Context, expectedInfoHash [20]byte) error {
	head := make([]byte, 1)
	if err := pc.setReadDeadlineFromContext(ctx, peerReadTimeout); err != nil {
		return err
	}
	if _, err := io.ReadFull(pc.conn, head); err != nil {
		return err
	}

	pstrlen := int(head[0])
	if pstrlen <= 0 || pstrlen > 64 {
		return fmt.Errorf("invalid handshake pstrlen %d", pstrlen)
	}

	rest := make([]byte, pstrlen+48)
	if _, err := io.ReadFull(pc.conn, rest); err != nil {
		return err
	}

	if string(rest[:pstrlen]) != wireProtocolString {
		return errors.New("invalid peer protocol string")
	}

	infoHashOffset := pstrlen + 8
	if !bytes.Equal(rest[infoHashOffset:infoHashOffset+20], expectedInfoHash[:]) {
		return errors.New("peer info_hash mismatch")
	}

	return nil
}

func (pc *peerClient) sendMessage(ctx context.Context, msgID int, payload []byte) error {
	if err := pc.setWriteDeadlineFromContext(ctx, peerWriteTimeout); err != nil {
		return err
	}

	length := uint32(1 + len(payload))
	buf := make([]byte, 4+length)
	binary.BigEndian.PutUint32(buf[0:4], length)
	buf[4] = byte(msgID)
	copy(buf[5:], payload)

	_, err := pc.conn.Write(buf)
	return err
}

func readWireMessage(r io.Reader) (wireMessage, error) {
	var lengthBuf [4]byte
	if _, err := io.ReadFull(r, lengthBuf[:]); err != nil {
		return wireMessage{}, err
	}

	length := binary.BigEndian.Uint32(lengthBuf[:])
	if length == 0 {
		return wireMessage{ID: -1}, nil
	}
	if length > maxWireMessageSize {
		return wireMessage{}, fmt.Errorf("wire message too large: %d", length)
	}

	msg := make([]byte, length)
	if _, err := io.ReadFull(r, msg); err != nil {
		return wireMessage{}, err
	}
	return wireMessage{ID: int(msg[0]), Payload: msg[1:]}, nil
}

func bitfieldHasPiece(bitfield []byte, index int) bool {
	byteIndex := index / 8
	if byteIndex < 0 || byteIndex >= len(bitfield) {
		return false
	}
	bitOffset := 7 - (index % 8)
	return bitfield[byteIndex]&(1<<bitOffset) != 0
}

func pieceSizeForIndex(tf *torrentfile.TorrentFile, pieceIndex int) int {
	if pieceIndex < 0 || pieceIndex >= len(tf.PieceHashes) {
		return 0
	}
	if pieceIndex == len(tf.PieceHashes)-1 {
		used := tf.PieceLength * (len(tf.PieceHashes) - 1)
		return tf.Length - used
	}
	return tf.PieceLength
}

func (pc *peerClient) setReadDeadlineFromContext(ctx context.Context, fallback time.Duration) error {
	deadline := time.Now().Add(fallback)
	if ctxDeadline, ok := ctx.Deadline(); ok && ctxDeadline.Before(deadline) {
		deadline = ctxDeadline
	}
	return pc.conn.SetReadDeadline(deadline)
}

func (pc *peerClient) setWriteDeadlineFromContext(ctx context.Context, fallback time.Duration) error {
	deadline := time.Now().Add(fallback)
	if ctxDeadline, ok := ctx.Deadline(); ok && ctxDeadline.Before(deadline) {
		deadline = ctxDeadline
	}
	return pc.conn.SetWriteDeadline(deadline)
}

func isTimeout(err error) bool {
	if err == nil {
		return false
	}
	var netErr net.Error
	return errors.As(err, &netErr) && netErr.Timeout()
}