waste-go/internal/mesh/transfer.go

package mesh

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"hash"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/pion/webrtc/v3"

	"github.com/waste-go/internal/proto"
)

const (
	fileChunkSize       = 64 * 1024  // 64 KiB per chunk
	fileBufferHighWater = 512 * 1024 // pause sending above this
	fileBufferLowWater  = 256 * 1024 // resume when it drops here
)

type outboundTransfer struct {
	peerID proto.PeerID
	path   string
	size   int64
	sha256 string
}

type inboundTransfer struct {
	from         proto.PeerID
	name         string
	size         int64
	sha256       string
	mu           sync.Mutex
	tmp          *os.File
	hasher       hash.Hash
	written      int64  // total bytes received; starts at resumeOffset when resuming
	resumePath   string // path to existing .tmp when resuming; empty for new transfers
	resumeOffset int64  // bytes already present in resumePath
	metaPath     string // path to the .tmp.meta sidecar
}

// partialMeta is written as a JSON sidecar alongside each in-progress .tmp file.
// It survives interruptions so the receiver can resume later.
type partialMeta struct {
	Name   string `json:"name"`
	SHA256 string `json:"sha256"`
	From   string `json:"from"`
	Size   int64  `json:"size"`
}

// findPartial scans dir for a .tmp.meta sidecar whose sha256 matches.
// Returns (tmpPath, metaPath, offset) — all empty/zero if no match.
func findPartial(dir, sha256hex string) (tmpPath, metaPath string, offset int64) {
	metas, _ := filepath.Glob(filepath.Join(dir, "*.tmp.meta"))
	for _, mp := range metas {
		data, err := os.ReadFile(mp)
		if err != nil {
			continue
		}
		var m partialMeta
		if err := json.Unmarshal(data, &m); err != nil {
			continue
		}
		if m.SHA256 != sha256hex {
			continue
		}
		tp := strings.TrimSuffix(mp, ".meta")
		info, err := os.Stat(tp)
		if err != nil {
			continue
		}
		return tp, mp, info.Size()
	}
	return "", "", 0
}

func writePartialMeta(path string, t *inboundTransfer) {
	data, _ := json.Marshal(partialMeta{
		Name:   t.name,
		SHA256: t.sha256,
		From:   string(t.from),
		Size:   t.size,
	})
	os.WriteFile(path, data, 0o644) //nolint:errcheck
}

// ScanResumable scans the download directory for .tmp.meta sidecars left by
// interrupted transfers and emits a resumable_transfers IPC event listing them.
// Called once after a network is joined so the UI can show pending transfers.
func (m *Mesh) ScanResumable() {
	if m.DownloadDir == "" {
		return
	}
	metas, _ := filepath.Glob(filepath.Join(m.DownloadDir, "*.tmp.meta"))
	var files []proto.ResumableFile
	for _, mp := range metas {
		data, err := os.ReadFile(mp)
		if err != nil {
			continue
		}
		var meta partialMeta
		if err := json.Unmarshal(data, &meta); err != nil {
			continue
		}
		tp := strings.TrimSuffix(mp, ".meta")
		info, err := os.Stat(tp)
		if err != nil {
			continue
		}
		files = append(files, proto.ResumableFile{
			Name:   meta.Name,
			SHA256: meta.SHA256,
			From:   meta.From,
			Size:   meta.Size,
			Offset: info.Size(),
		})
	}
	if len(files) == 0 {
		return
	}
	m.emit(proto.IpcMessage{
		Type:           proto.EvtResumableTransfers,
		ResumableFiles: files,
	})
	log.Printf("transfer: %d resumable transfer(s) found in %s", len(files), m.DownloadDir)
}

// OfferFile reads filename from ShareDir, computes its SHA-256, and sends a
// file-offer to peerID over the existing "yaw" DataChannel.
func (m *Mesh) OfferFile(peerID proto.PeerID, filename string) error {
	if m.ShareDir == "" {
		return fmt.Errorf("no share directory configured")
	}
	if strings.ContainsAny(filename, "/\\") || filename == ".." {
		return fmt.Errorf("invalid filename %q", filename)
	}
	path := filepath.Join(m.ShareDir, filename)

	f, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("open %q: %w", filename, err)
	}
	defer f.Close()

	info, err := f.Stat()
	if err != nil {
		return err
	}

	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return fmt.Errorf("hash %q: %w", filename, err)
	}
	shaHex := hex.EncodeToString(h.Sum(nil))
	xid := newXid()

	m.transferMu.Lock()
	m.outbound[xid] = &outboundTransfer{
		peerID: peerID,
		path:   path,
		size:   info.Size(),
		sha256: shaHex,
	}
	m.transferMu.Unlock()

	wire, err := json.Marshal(proto.PeerMessage{
		Type:   proto.MsgFileOffer,
		Xid:    xid,
		Name:   filename,
		Size:   info.Size(),
		SHA256: shaHex,
	})
	if err != nil {
		return err
	}
	if !m.SendTo(peerID, wire) {
		m.transferMu.Lock()
		delete(m.outbound, xid)
		m.transferMu.Unlock()
		return fmt.Errorf("peer %s not connected", peerID.Short())
	}
	log.Printf("transfer: offered %s (%d bytes) to %s xid=%s", filename, info.Size(), peerID.Short(), xid[:8])
	return nil
}

// acceptIncoming is called when we receive file-offer from a peer.
// Checks for a resumable partial download matching the offer's sha256.
// Sends file-accept immediately, with a non-zero resume_offset when resuming.
func (m *Mesh) acceptIncoming(msg proto.PeerMessage, from proto.PeerID) {
	t := &inboundTransfer{
		from:   from,
		name:   msg.Name,
		size:   msg.Size,
		sha256: msg.SHA256,
	}

	var resumeOffset int64
	if m.DownloadDir != "" {
		if tp, mp, off := findPartial(m.DownloadDir, msg.SHA256); tp != "" {
			t.resumePath = tp
			t.metaPath = mp
			t.resumeOffset = off
			t.written = off
			resumeOffset = off
			log.Printf("transfer: found partial for %s at %s (%d/%d bytes)", msg.Name, tp, off, msg.Size)
		}
	}

	m.transferMu.Lock()
	m.inbound[msg.Xid] = t
	m.transferMu.Unlock()

	accept, _ := json.Marshal(proto.PeerMessage{
		Type:         proto.MsgFileAccept,
		Xid:          msg.Xid,
		ResumeOffset: resumeOffset,
	})
	m.SendTo(from, accept)
	if resumeOffset > 0 {
		log.Printf("transfer: resuming %s from byte %d xid=%s", msg.Name, resumeOffset, msg.Xid[:8])
	} else {
		log.Printf("transfer: auto-accepted %s (%d bytes) from %s xid=%s", msg.Name, msg.Size, from.Short(), msg.Xid[:8])
	}
}

// startSend opens the binary DataChannel "f:<xid>" and streams the file.
// Called when we receive file-accept from the peer. resumeOffset is non-zero
// when the receiver is resuming a previous partial download.
func (m *Mesh) startSend(xid string, from proto.PeerID, resumeOffset int64) {
	m.transferMu.Lock()
	t, ok := m.outbound[xid]
	m.transferMu.Unlock()
	if !ok {
		log.Printf("transfer: file-accept for unknown xid %s", xid[:8])
		return
	}
	if t.peerID != from {
		log.Printf("transfer: file-accept from unexpected peer %s", from.Short())
		return
	}

	m.mu.RLock()
	conn, ok := m.peers[from]
	m.mu.RUnlock()
	if !ok {
		log.Printf("transfer: peer %s not connected for xid %s", from.Short(), xid[:8])
		return
	}

	ordered := true
	dc, err := conn.PC.CreateDataChannel("f:"+xid, &webrtc.DataChannelInit{Ordered: &ordered})
	if err != nil {
		log.Printf("transfer: create file DC xid=%s: %v", xid[:8], err)
		return
	}

	// Wait for receiver to signal "ok" before sending — this ensures the receiver
	// has its OnMessage/OnClose handlers registered before any data arrives.
	dc.OnOpen(func() {
		dc.OnMessage(func(msg webrtc.DataChannelMessage) {
			if msg.IsString && string(msg.Data) == "ok" {
				go m.sendFileChunks(dc, t, xid, resumeOffset)
			}
		})
	})
}

func (m *Mesh) sendFileChunks(dc *webrtc.DataChannel, t *outboundTransfer, xid string, offset int64) {
	defer func() {
		m.transferMu.Lock()
		delete(m.outbound, xid)
		m.transferMu.Unlock()
	}()

	f, err := os.Open(t.path)
	if err != nil {
		log.Printf("transfer: open file xid=%s: %v", xid[:8], err)
		dc.Close()
		return
	}
	defer f.Close()

	if offset > 0 {
		if _, err := f.Seek(offset, io.SeekStart); err != nil {
			log.Printf("transfer: seek to %d xid=%s: %v", offset, xid[:8], err)
			dc.Close()
			return
		}
		log.Printf("transfer: resuming send from byte %d xid=%s", offset, xid[:8])
	}

	// Backpressure: block when the DC buffer is full.
	resume := make(chan struct{}, 1)
	dc.SetBufferedAmountLowThreshold(fileBufferLowWater)
	dc.OnBufferedAmountLow(func() {
		select {
		case resume <- struct{}{}:
		default:
		}
	})

	buf := make([]byte, fileChunkSize)
	sent := offset // start progress reporting from where the receiver left off

	for {
		n, readErr := f.Read(buf)
		if n > 0 {
			for dc.BufferedAmount() > fileBufferHighWater {
				<-resume
			}
			if err := dc.Send(buf[:n]); err != nil {
				log.Printf("transfer: send chunk xid=%s: %v", xid[:8], err)
				return
			}
			sent += int64(n)
			m.Emit(proto.IpcMessage{
				Type:          proto.EvtFileProgress,
				TransferID:    xid,
				BytesReceived: sent,
				TotalBytes:    t.size,
			})
		}
		if readErr == io.EOF {
			break
		}
		if readErr != nil {
			log.Printf("transfer: read file xid=%s: %v", xid[:8], readErr)
			return
		}
	}

	dc.Close()

	done, _ := json.Marshal(proto.PeerMessage{
		Type:   proto.MsgFileDone,
		Xid:    xid,
		SHA256: t.sha256,
	})
	m.SendTo(t.peerID, done)
	log.Printf("transfer: sent %s (%d bytes) xid=%s", filepath.Base(t.path), sent, xid[:8])
}

// HandleInboundFileDC is called from the anchor when a "f:<xid>" DataChannel
// arrives on a PeerConnection. It writes chunks to a temp file (or resumes an
// existing partial), verifies the SHA-256 on close, and emits EvtFileComplete.
// Interrupted transfers keep their .tmp and .meta files for future resume.
func (m *Mesh) HandleInboundFileDC(dc *webrtc.DataChannel, xid string, from proto.PeerID) {
	m.transferMu.Lock()
	t, ok := m.inbound[xid]
	m.transferMu.Unlock()
	if !ok {
		log.Printf("transfer: no pending inbound for xid %s", xid[:8])
		dc.Close()
		return
	}

	var once sync.Once
	doOpen := func() {
		if err := os.MkdirAll(m.DownloadDir, 0o755); err != nil {
			log.Printf("transfer: mkdir %s: %v", m.DownloadDir, err)
			return
		}

		t.mu.Lock()
		defer t.mu.Unlock()

		if t.resumePath != "" {
			// Resume: open existing .tmp in read-write mode, seek to end for appending.
			f, err := os.OpenFile(t.resumePath, os.O_RDWR, 0o644)
			if err != nil {
				log.Printf("transfer: open resume file %s: %v", t.resumePath, err)
				return
			}
			// Restore the hasher state by re-hashing the bytes already on disk.
			h := sha256.New()
			if _, err := f.Seek(0, io.SeekStart); err == nil {
				io.Copy(h, f) //nolint:errcheck
			}
			if _, err := f.Seek(0, io.SeekEnd); err != nil {
				log.Printf("transfer: seek to end xid=%s: %v", xid[:8], err)
				f.Close()
				return
			}
			t.tmp = f
			t.hasher = h
			log.Printf("transfer: resuming inbound %s at byte %d xid=%s", t.name, t.resumeOffset, xid[:8])
		} else {
			tmp, err := os.CreateTemp(m.DownloadDir, "dl-*.tmp")
			if err != nil {
				log.Printf("transfer: create temp file: %v", err)
				return
			}
			t.tmp = tmp
			t.hasher = sha256.New()
			t.metaPath = tmp.Name() + ".meta"
			writePartialMeta(t.metaPath, t)
			log.Printf("transfer: receiving %s xid=%s", t.name, xid[:8])
		}

		// Signal sender that we are ready — it will not start streaming until
		// it receives this, guaranteeing our OnMessage/OnClose are registered first.
		dc.SendText("ok") //nolint:errcheck
	}

	// Register OnMessage first so pion queues any early binary chunks.
	dc.OnMessage(func(msg webrtc.DataChannelMessage) {
		if msg.IsString {
			return // "ok" echo or any other text; ignore
		}
		t.mu.Lock()
		tmp, h := t.tmp, t.hasher
		t.mu.Unlock()
		if tmp == nil {
			return
		}
		if _, err := tmp.Write(msg.Data); err != nil {
			log.Printf("transfer: write xid=%s: %v", xid[:8], err)
			return
		}
		h.Write(msg.Data)
		t.mu.Lock()
		t.written += int64(len(msg.Data))
		written := t.written
		t.mu.Unlock()
		m.Emit(proto.IpcMessage{
			Type:          proto.EvtFileProgress,
			TransferID:    xid,
			BytesReceived: written,
			TotalBytes:    t.size,
		})
	})

	dc.OnOpen(func() { once.Do(doOpen) })
	if dc.ReadyState() == webrtc.DataChannelStateOpen {
		once.Do(doOpen)
	}

	dc.OnClose(func() {
		t.mu.Lock()
		tmp := t.tmp
		metaPath := t.metaPath
		actualSha := ""
		if t.hasher != nil {
			actualSha = hex.EncodeToString(t.hasher.Sum(nil))
		}
		name, expectedSha, size, written := t.name, t.sha256, t.size, t.written
		t.mu.Unlock()

		m.transferMu.Lock()
		delete(m.inbound, xid)
		m.transferMu.Unlock()

		if tmp == nil {
			return
		}
		tmpName := tmp.Name()
		tmp.Close()

		// Incomplete — keep .tmp and .meta for future resume.
		if written < size {
			log.Printf("transfer: interrupted %s at %d/%d bytes xid=%s (resumable)", name, written, size, xid[:8])
			return
		}

		// Full receive — verify integrity.
		if actualSha != expectedSha {
			os.Remove(tmpName)
			if metaPath != "" {
				os.Remove(metaPath)
			}
			log.Printf("transfer: sha256 mismatch xid=%s got=%s want=%s", xid[:8], actualSha[:8], expectedSha[:8])
			m.Emit(proto.IpcMessage{
				Type:         proto.EvtError,
				TransferID:   xid,
				ErrorMessage: fmt.Sprintf("file %s: sha256 mismatch", name),
			})
			return
		}

		// Success — remove sidecar and move to final path.
		if metaPath != "" {
			os.Remove(metaPath)
		}
		final := filepath.Join(m.DownloadDir, name)
		if _, err := os.Stat(final); err == nil {
			ext := filepath.Ext(name)
			base := strings.TrimSuffix(name, ext)
			final = filepath.Join(m.DownloadDir, base+"-"+xid[:8]+ext)
		}
		if err := os.Rename(tmpName, final); err != nil {
			log.Printf("transfer: rename xid=%s: %v", xid[:8], err)
			return
		}
		log.Printf("transfer: saved %s -> %s", name, final)
		m.Emit(proto.IpcMessage{
			Type:       proto.EvtFileComplete,
			TransferID: xid,
			Path:       final,
		})
	})
}

func newXid() string {
	b := make([]byte, 16)
	rand.Read(b) //nolint:errcheck
	return hex.EncodeToString(b)
}