waste-go/internal/mesh/mesh.go

// Package mesh manages the set of live peer connections and broadcasts events.
package mesh

import (
	"log"
	"os"
	"sync"

	"github.com/pion/webrtc/v3"

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

// PeerConn is a live connection to one peer.
type PeerConn struct {
	Info proto.PeerInfo
	// Send a line of JSON to this peer (pre-encrypted by the sender goroutine).
	Send chan<- []byte
	// PC is the underlying PeerConnection, used to open additional DataChannels.
	PC *webrtc.PeerConnection
}

// Mesh is the shared state of the local node.
// All methods are safe to call from multiple goroutines.
type Mesh struct {
	Identity    *crypto.Identity
	Store       *store.Store // may be nil if persistence is disabled
	ShareDir      string // directory whose contents are shared with peers; "" = no sharing
	DownloadDir   string // directory where received files are saved
	RequireInvite bool   // waste-go ext: reject peers that present no valid signed invite
	InviteString  string // the invite this peer used to join (sent in hello to other peers)
	// ScanFiles overrides ScanShareDir when set — allows the manager to inject
	// multi-share scanning without the mesh needing to know about shares.json.
	ScanFiles func() []proto.FileEntry

	mu    sync.RWMutex
	peers map[proto.PeerID]*PeerConn

	// file transfer state
	transferMu sync.Mutex
	outbound   map[string]*outboundTransfer // xid → pending outbound
	inbound    map[string]*inboundTransfer  // xid → pending inbound

	// PendingConnect receives peer IDs discovered via gossip that we should
	// attempt to connect to. Drained by the anchor client's runOnce loop.
	PendingConnect chan proto.PeerID

	// subscribers receive a copy of every event (fan-out to IPC clients)
	subMu sync.Mutex
	subs  []chan proto.IpcMessage
}

// New creates an empty mesh with the given identity.
// Pass a non-nil store to enable message and peer persistence.
func New(id *crypto.Identity, st *store.Store) *Mesh {
	return &Mesh{
		Identity:       id,
		Store:          st,
		peers:          make(map[proto.PeerID]*PeerConn),
		outbound:       make(map[string]*outboundTransfer),
		inbound:        make(map[string]*inboundTransfer),
		PendingConnect: make(chan proto.PeerID, 32),
	}
}

// trustedPeerIDs returns a set of peer IDs trusted on this network:
// all currently connected peers plus all peers in the persistent store.
func (m *Mesh) trustedPeerIDs() map[string]bool {
	trusted := map[string]bool{}
	// Own identity is always trusted.
	trusted[string(m.Identity.PeerID())] = true
	// Connected peers.
	m.mu.RLock()
	for id := range m.peers {
		trusted[string(id)] = true
	}
	m.mu.RUnlock()
	// Previously seen peers from the store.
	if m.Store != nil {
		if known, err := m.Store.KnownPeers(); err == nil {
			for id := range known {
				trusted[string(id)] = true
			}
		}
	}
	return trusted
}

// ScanShareDir returns the list of files in the local share directory.
// Returns an empty slice if ShareDir is unset or the directory is empty.
func (m *Mesh) ScanShareDir() []proto.FileEntry {
	if m.ShareDir == "" {
		return nil
	}
	entries, err := os.ReadDir(m.ShareDir)
	if err != nil {
		log.Printf("mesh: scan share dir %s: %v", m.ShareDir, err)
		return nil
	}
	var files []proto.FileEntry
	for _, e := range entries {
		if e.IsDir() {
			continue
		}
		info, err := e.Info()
		if err != nil {
			continue
		}
		files = append(files, proto.FileEntry{Name: e.Name(), SizeBytes: info.Size()})
	}
	return files
}

// ── Peer management ───────────────────────────────────────────────────────────

// AddPeer registers a connected peer and notifies subscribers.
func (m *Mesh) AddPeer(conn *PeerConn) {
	// Seed alias from cache so returning peers resolve immediately (before hello).
	if m.Store != nil {
		if cached := m.Store.PeerAlias(conn.Info.ID); cached != "" {
			conn.Info.Alias = cached
		}
	}

	m.mu.Lock()
	m.peers[conn.Info.ID] = conn
	m.mu.Unlock()

	m.emit(proto.IpcMessage{
		Type: proto.EvtPeerConnected,
		Peer: &conn.Info,
	})
}

// SaveMessage persists a chat message if a store is configured.
// Duplicate mids are silently dropped.
func (m *Mesh) SaveMessage(msg *proto.ChatMessage) {
	if m.Store == nil || msg == nil {
		return
	}
	if err := m.Store.SaveMessage(msg); err != nil {
		log.Printf("mesh: store message %s: %v", msg.Mid, err)
	}
}

// UpdatePeerAlias updates the cached alias for a peer after hello verification.
func (m *Mesh) UpdatePeerAlias(id proto.PeerID, alias string) {
	if m.Store != nil && alias != "" {
		if err := m.Store.SavePeer(id, alias); err != nil {
			log.Printf("mesh: update peer alias %s: %v", id.Short(), err)
		}
	}
}

// RemovePeer unregisters a peer and notifies subscribers.
func (m *Mesh) RemovePeer(id proto.PeerID) {
	m.mu.Lock()
	delete(m.peers, id)
	m.mu.Unlock()

	m.emit(proto.IpcMessage{
		Type:   proto.EvtPeerDisconnected,
		PeerID: &id,
	})
}

// ConnectedPeers returns a snapshot of current peer infos.
func (m *Mesh) ConnectedPeers() []proto.PeerInfo {
	m.mu.RLock()
	defer m.mu.RUnlock()
	out := make([]proto.PeerInfo, 0, len(m.peers))
	for _, c := range m.peers {
		out = append(out, c.Info)
	}
	return out
}

// SendTo delivers a raw JSON payload to a specific peer.
// Returns false if the peer isn't connected.
func (m *Mesh) SendTo(id proto.PeerID, payload []byte) bool {
	m.mu.RLock()
	conn, ok := m.peers[id]
	m.mu.RUnlock()
	if !ok {
		return false
	}
	select {
	case conn.Send <- payload:
		return true
	default:
		return false // channel full — peer is slow
	}
}

// Broadcast delivers a raw JSON payload to every connected peer.
func (m *Mesh) Broadcast(payload []byte) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	for _, conn := range m.peers {
		select {
		case conn.Send <- payload:
		default:
		}
	}
}

// ── Event fan-out ─────────────────────────────────────────────────────────────

// Subscribe returns a channel that receives every IPC event.
// The caller must drain it; a full channel is silently dropped.
func (m *Mesh) Subscribe() <-chan proto.IpcMessage {
	ch := make(chan proto.IpcMessage, 64)
	m.subMu.Lock()
	m.subs = append(m.subs, ch)
	m.subMu.Unlock()
	return ch
}

// Unsubscribe removes and closes a subscription channel.
func (m *Mesh) Unsubscribe(ch <-chan proto.IpcMessage) {
	m.subMu.Lock()
	defer m.subMu.Unlock()
	for i, s := range m.subs {
		if s == ch {
			m.subs = append(m.subs[:i], m.subs[i+1:]...)
			close(s)
			return
		}
	}
}

// Emit sends an event to all IPC subscribers (exported for ipc/nat packages).
func (m *Mesh) Emit(msg proto.IpcMessage) {
	m.emit(msg)
}

func (m *Mesh) emit(msg proto.IpcMessage) {
	m.subMu.Lock()
	defer m.subMu.Unlock()
	for _, ch := range m.subs {
		select {
		case ch <- msg:
		default:
		}
	}
}