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
}

// 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.
// It stores the inbound state and immediately sends file-accept (auto-accept).
func (m *Mesh) acceptIncoming(msg proto.PeerMessage, from proto.PeerID) {
	m.transferMu.Lock()
	m.inbound[msg.Xid] = &inboundTransfer{
		from:   from,
		name:   msg.Name,
		size:   msg.Size,
		sha256: msg.SHA256,
	}
	m.transferMu.Unlock()

	accept, _ := json.Marshal(proto.PeerMessage{
		Type: proto.MsgFileAccept,
		Xid:  msg.Xid,
	})
	m.SendTo(from, accept)
	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.
func (m *Mesh) startSend(xid string, from proto.PeerID) {
	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)
			}
		})
	})
}

func (m *Mesh) sendFileChunks(dc *webrtc.DataChannel, t *outboundTransfer, xid string) {
	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()

	// 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)
	var sent int64

	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, verifies the
// SHA-256 on close, and emits EvtFileComplete.
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
		}
		tmp, err := os.CreateTemp(m.DownloadDir, "dl-*.tmp")
		if err != nil {
			log.Printf("transfer: create temp file: %v", err)
			return
		}
		t.mu.Lock()
		t.tmp = tmp
		t.hasher = sha256.New()
		t.mu.Unlock()
		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
		actualSha := ""
		if t.hasher != nil {
			actualSha = hex.EncodeToString(t.hasher.Sum(nil))
		}
		name, expectedSha := t.name, t.sha256
		t.mu.Unlock()

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

		if tmp == nil {
			return
		}
		tmp.Close()

		if actualSha != expectedSha {
			os.Remove(tmp.Name())
			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
		}

		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(tmp.Name(), 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)
}