share/src/swarm/client.go

package swarm

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	log "github.com/tzgyn/share/src/logger"

	"github.com/tzgyn/share/src/comm"
	"github.com/tzgyn/share/src/tcp"
	"github.com/tzgyn/share/src/utils"
)

// Options configures one participant in a swarm room. Everyone who joins
// with the same Code can both send and receive: peers who pass Files seed
// them into the room, and every peer (seeding or not) helps re-share
// whatever chunks it has already received to others who join later.
type Options struct {
	RelayAddress string
	Code         string   // room name and PAKE codephrase, both at once
	Files        []string // local files to seed; empty means pure receiver
	OutDir       string   // where received files are written
	OneTime      bool     // exit once everything currently known is complete, instead of staying to reseed
}

// Run joins the room described by opts and blocks until either an error
// occurs, or (with OneTime set) every file known in the room has been
// fully downloaded and verified. Without OneTime, Run only returns on
// error — like a real torrent peer, it stays to keep reseeding whatever
// it has to others until stopped.
func Run(opts Options) error {
	return newClient(opts).run()
}

type sharedFile struct {
	entry FileManifestEntry // immutable identity: ID, Name, Size, Hash
	store *ChunkStore
}

type peerLink struct {
	conn   *comm.Comm
	key    []byte
	remote string // remote address, for logging who's talking to whom
}

func newPeerLink(c *comm.Comm, key []byte) *peerLink {
	return &peerLink{conn: c, key: key, remote: c.Connection().RemoteAddr().String()}
}

type client struct {
	opts   Options
	selfID string

	mu      sync.Mutex
	files   map[string]*sharedFile
	pending map[string]map[int64]*peerLink // fileID -> offset -> link we asked
	links   map[*peerLink]bool

	dialedMu sync.Mutex
	dialed   map[string]bool // relay peer IDs already dialed, so we don't redial every peer-list update
}

func newClient(opts Options) *client {
	return &client{
		opts:    opts,
		selfID:  randomID(),
		files:   make(map[string]*sharedFile),
		pending: make(map[string]map[int64]*peerLink),
		links:   make(map[*peerLink]bool),
		dialed:  make(map[string]bool),
	}
}

func randomID() string {
	b := make([]byte, 8)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	return hex.EncodeToString(b)
}

func (cl *client) run() error {
	for _, path := range cl.opts.Files {
		if err := cl.addSeedFile(path); err != nil {
			return err
		}
	}

	listener, listenAddr, err := startListener()
	if err != nil {
		return fmt.Errorf("could not open a listener for peers: %w", err)
	}
	defer listener.Close()
	go cl.acceptLoop(listener)

	relayConn, peerUpdates, err := tcp.JoinRoom(cl.opts.RelayAddress, cl.opts.Code, tcp.PeerInfo{ID: cl.selfID, Addr: listenAddr})
	if err != nil {
		return fmt.Errorf("could not join room: %w", err)
	}
	defer relayConn.Close()

	for {
		select {
		case peers, ok := <-peerUpdates:
			if !ok {
				return fmt.Errorf("lost connection to relay")
			}
			for _, p := range peers {
				if p.ID == cl.selfID || p.Addr == "" {
					continue
				}
				cl.dialedMu.Lock()
				already := cl.dialed[p.ID]
				cl.dialed[p.ID] = true
				cl.dialedMu.Unlock()
				if !already {
					go cl.dialPeer(p)
				}
			}
		case <-time.After(500 * time.Millisecond):
			if cl.opts.OneTime && cl.allKnownFilesComplete() {
				return nil
			}
		}
	}
}

func (cl *client) addSeedFile(path string) error {
	info, err := os.Stat(path)
	if err != nil {
		return err
	}
	if info.IsDir() {
		return fmt.Errorf("%s is a directory; folder sharing isn't supported yet", path)
	}
	hash, err := utils.HashFile(path, "xxhash")
	if err != nil {
		return err
	}
	name := filepath.Base(path)
	id := utils.SHA256(fmt.Sprintf("%s:%d", name, info.Size()))
	cl.files[id] = &sharedFile{
		entry: FileManifestEntry{ID: id, Name: name, Size: info.Size(), Hash: hash},
		store: NewSeedChunkStore(path, info.Size()),
	}
	log.Infof("seeding %s (%s)", name, utils.ByteCountDecimal(info.Size()))
	return nil
}

func (cl *client) allKnownFilesComplete() bool {
	cl.mu.Lock()
	defer cl.mu.Unlock()
	if len(cl.files) == 0 {
		return false
	}
	for _, f := range cl.files {
		if !f.store.Complete() {
			return false
		}
	}
	return true
}

func startListener() (net.Listener, string, error) {
	l, err := net.Listen("tcp", ":0")
	if err != nil {
		return nil, "", err
	}
	port := l.Addr().(*net.TCPAddr).Port
	return l, net.JoinHostPort(utils.LocalIP(), strconv.Itoa(port)), nil
}

func (cl *client) acceptLoop(l net.Listener) {
	for {
		conn, err := l.Accept()
		if err != nil {
			return
		}
		go func(conn net.Conn) {
			c := comm.New(conn)
			key, err := handshake(c, cl.opts.Code, false)
			if err != nil {
				log.Debugf("peer handshake failed: %v", err)
				conn.Close()
				return
			}
			cl.handlePeerLink(newPeerLink(c, key))
		}(conn)
	}
}

func (cl *client) dialPeer(p tcp.PeerInfo) {
	conn, err := net.DialTimeout("tcp", p.Addr, 5*time.Second)
	if err != nil {
		log.Debugf("could not dial peer %s at %s: %v", p.ID, p.Addr, err)
		return
	}
	c := comm.New(conn)
	key, err := handshake(c, cl.opts.Code, true)
	if err != nil {
		log.Debugf("handshake with peer %s failed: %v", p.ID, err)
		conn.Close()
		return
	}
	cl.handlePeerLink(newPeerLink(c, key))
}

func (cl *client) handlePeerLink(link *peerLink) {
	cl.mu.Lock()
	cl.links[link] = true
	cl.mu.Unlock()
	log.Infof("connected to peer %s", link.remote)
	defer func() {
		cl.mu.Lock()
		delete(cl.links, link)
		// free up anything only this link was asked for, so others can serve it
		for _, byOffset := range cl.pending {
			for offset, l := range byOffset {
				if l == link {
					delete(byOffset, offset)
				}
			}
		}
		cl.mu.Unlock()
		link.conn.Close()
		log.Infof("disconnected from peer %s", link.remote)
	}()

	if err := cl.sendManifest(link); err != nil {
		log.Debugf("send manifest failed: %v", err)
		return
	}

	for {
		msg, err := receiveMessage(link.conn, link.key)
		if err != nil {
			log.Debugf("peer link closed: %v", err)
			return
		}
		if err := cl.handleMessage(link, msg); err != nil {
			log.Debugf("handling %s message: %v", msg.Type, err)
		}
	}
}

func (cl *client) sendManifest(link *peerLink) error {
	cl.mu.Lock()
	files := make([]FileManifestEntry, 0, len(cl.files))
	for _, f := range cl.files {
		e := f.entry
		e.Have = f.store.HaveOffsets()
		files = append(files, e)
	}
	cl.mu.Unlock()
	return sendMessage(link.conn, link.key, Message{Type: TypeManifest, Files: files})
}

func (cl *client) handleMessage(link *peerLink, msg Message) error {
	switch msg.Type {
	case TypeManifest:
		for _, entry := range msg.Files {
			cl.learnFile(entry)
			cl.considerRequests(entry.ID, entry.Have, link)
		}
	case TypeHave:
		cl.considerRequests(msg.FileID, []int64{msg.Offset}, link)
	case TypeRequest:
		return cl.serveRequest(link, msg.FileID, msg.Offset)
	case TypeChunk:
		return cl.receiveChunk(link, msg.FileID, msg.Offset, msg.Data)
	}
	return nil
}

func (cl *client) learnFile(entry FileManifestEntry) {
	cl.mu.Lock()
	_, known := cl.files[entry.ID]
	cl.mu.Unlock()
	if known {
		return
	}
	path := filepath.Join(cl.opts.OutDir, entry.Name)
	store, err := OpenChunkStore(path, entry.Size)
	if err != nil {
		log.Errorf("could not open %s for receiving: %v", path, err)
		return
	}
	cl.mu.Lock()
	cl.files[entry.ID] = &sharedFile{
		entry: FileManifestEntry{ID: entry.ID, Name: entry.Name, Size: entry.Size, Hash: entry.Hash},
		store: store,
	}
	cl.mu.Unlock()
	log.Infof("new file in room: %s (%s)", entry.Name, utils.ByteCountDecimal(entry.Size))
}

// considerRequests asks link for any of the given offsets of fileID that
// we're still missing and haven't already asked someone else for.
func (cl *client) considerRequests(fileID string, offsets []int64, link *peerLink) {
	cl.mu.Lock()
	f, ok := cl.files[fileID]
	if !ok {
		cl.mu.Unlock()
		return
	}
	if cl.pending[fileID] == nil {
		cl.pending[fileID] = make(map[int64]*peerLink)
	}
	toRequest := make([]int64, 0, len(offsets))
	for _, off := range offsets {
		if f.store.HasChunk(off) {
			continue
		}
		if _, already := cl.pending[fileID][off]; already {
			continue
		}
		cl.pending[fileID][off] = link
		toRequest = append(toRequest, off)
	}
	cl.mu.Unlock()

	for _, off := range toRequest {
		if err := sendMessage(link.conn, link.key, Message{Type: TypeRequest, FileID: fileID, Offset: off}); err != nil {
			log.Debugf("request failed: %v", err)
			return
		}
	}
}

func (cl *client) serveRequest(link *peerLink, fileID string, offset int64) error {
	cl.mu.Lock()
	f, ok := cl.files[fileID]
	cl.mu.Unlock()
	if !ok || !f.store.HasChunk(offset) {
		return nil
	}
	data, err := f.store.ReadChunk(offset)
	if err != nil {
		return err
	}
	log.Infof("sending %s chunk %d/%d to %s", f.entry.Name, offset/ChunkSize+1, NumChunks(f.entry.Size), link.remote)
	return sendMessage(link.conn, link.key, Message{Type: TypeChunk, FileID: fileID, Offset: offset, Data: data})
}

func (cl *client) receiveChunk(link *peerLink, fileID string, offset int64, data []byte) error {
	cl.mu.Lock()
	f, ok := cl.files[fileID]
	if ok {
		delete(cl.pending[fileID], offset)
	}
	links := make([]*peerLink, 0, len(cl.links))
	for l := range cl.links {
		links = append(links, l)
	}
	cl.mu.Unlock()
	if !ok {
		return nil
	}

	log.Infof("received %s chunk %d/%d from %s", f.entry.Name, offset/ChunkSize+1, NumChunks(f.entry.Size), link.remote)
	if err := f.store.WriteChunk(offset, data); err != nil {
		return err
	}

	// let everyone else in the swarm know we can now serve this chunk too
	for _, l := range links {
		go func(l *peerLink) {
			_ = sendMessage(l.conn, l.key, Message{Type: TypeHave, FileID: fileID, Offset: offset})
		}(l)
	}

	if f.store.Complete() {
		verified, err := f.store.VerifyHash(f.entry.Hash)
		if err != nil {
			log.Errorf("could not verify %s: %v", f.entry.Name, err)
		} else if !verified {
			log.Errorf("%s failed hash verification after transfer", f.entry.Name)
		} else {
			log.Infof("%s complete and verified", f.entry.Name)
		}
	}
	return nil
}