storrent/client/manager.go

package client

import (
	"errors"
	"fmt"
	"net"
	"strconv"
	"sync"

	"storrent/bt"
	"storrent/dht"
	"storrent/metainfo"
	"storrent/utp"
)

var errNotFound = errors.New("torrent not found")

func (m *Manager) getTorrent(id int) (*torrent, error) {
	m.mu.Lock()
	t, ok := m.torrents[id]
	m.mu.Unlock()
	if !ok {
		return nil, errNotFound
	}
	return t, nil
}

type Manager struct {
	mu       sync.Mutex
	torrents map[int]*torrent
	nextID   int
	dir      string
	dht      *dht.DHT
	ln       net.Listener
	utp      *utp.Socket
}

func NewManager(dir string, d *dht.DHT, btPort int, utpSock *utp.Socket) (*Manager, error) {
	var ln net.Listener
	if btPort > 0 {
		var err error
		ln, err = net.Listen("tcp", fmt.Sprintf(":%d", btPort))
		if err != nil {
			return nil, err
		}
	}
	return &Manager{
		torrents: make(map[int]*torrent),
		dir:      dir,
		dht:      d,
		ln:       ln,
		utp:      utpSock,
	}, nil
}

func (m *Manager) Run() {
	if m.ln == nil {
		return
	}
	for {
		p, err := bt.Accept(m.ln)
		if err != nil {
			return
		}
		m.dispatch(p)
	}
}

func (m *Manager) dispatch(p *bt.Peer) {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, t := range m.torrents {
		if t.meta.InfoHash == p.InfoHash {
			select {
			case t.incoming <- p:
			default:
				p.Close()
			}
			return
		}
	}
	p.Close()
}

func (m *Manager) Add(path string) ([]byte, error) {
	mi, err := metainfo.Parse(path)
	if err != nil {
		return nil, err
	}
	m.mu.Lock()
	id := m.nextID
	m.nextID++
	t := newTorrent(id, mi, m.dir, m.dht, m.utp)
	complete := t.scheduler.isComplete()
	go t.run()
	m.torrents[id] = t
	m.mu.Unlock()
	if complete {
		t.startSeed()
	} else {
		t.startDownload()
	}
	return []byte(strconv.Itoa(id)), nil
}

func (m *Manager) Remove(id int) error {
	m.mu.Lock()
	t, ok := m.torrents[id]
	if !ok {
		m.mu.Unlock()
		return errNotFound
	}
	delete(m.torrents, id)
	m.mu.Unlock()
	t.stopTorrent()
	return nil
}

func (m *Manager) Start(id int) error {
	t, err := m.getTorrent(id)
	if err != nil {
		return err
	}
	t.startDownload()
	return nil
}

func (m *Manager) Stop(id int) error {
	t, err := m.getTorrent(id)
	if err != nil {
		return err
	}
	t.stopTorrent()
	return nil
}

func (m *Manager) Seed(id int) error {
	t, err := m.getTorrent(id)
	if err != nil {
		return err
	}
	t.startSeed()
	return nil
}

func (m *Manager) AddPeer(id int, addr string) error {
	tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		return err
	}
	t, err := m.getTorrent(id)
	if err != nil {
		return err
	}
	t.enqueuePeer(tcpAddr)
	return nil
}

func (m *Manager) Status(id int, field string) ([]byte, error) {
	t, err := m.getTorrent(id)
	if err != nil {
		return nil, err
	}
	return []byte(t.getStatus(field)), nil
}

func (m *Manager) List() []byte {
	m.mu.Lock()
	defer m.mu.Unlock()
	var data []byte
	for id := range m.torrents {
		if len(data) > 0 {
			data = append(data, '\n')
		}
		data = append(data, []byte(strconv.Itoa(id))...)
	}
	return data
}