waste-go/internal/mesh/mesh.go

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

import (
	"encoding/json"
	"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"
)

// ICEServer mirrors webrtc.ICEServer so callers don't import pion directly.
type ICEServer = webrtc.ICEServer

// 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
	ICEServers []ICEServer // extra ICE servers (e.g. TURN); appended to the default STUN entry

	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

	// historyRequested tracks rooms for which we have already sent a history_request
	// this session. Reset on reconnect is intentional (new peers may have newer history).
	historyMu        sync.Mutex
	historyRequested map[string]bool // room → true
	historyFirstPeer proto.PeerID    // ID of the peer we requested history from

	// 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),
		historyRequested: make(map[string]bool),
	}
}

// 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
		}
	}
}

// RequestHistoryFrom sends history_request messages to peerID for all rooms
// we know about but haven't yet requested this session. Only contacts the first
// peer we connect to, to avoid fan-out amplification.
func (m *Mesh) RequestHistoryFrom(peerID proto.PeerID) {
	if m.Store == nil {
		return
	}
	m.historyMu.Lock()
	if m.historyFirstPeer != "" && m.historyFirstPeer != peerID {
		m.historyMu.Unlock()
		return // only request from the first peer
	}
	m.historyFirstPeer = peerID
	m.historyMu.Unlock()

	rooms, err := m.Store.Rooms()
	if err != nil {
		return
	}
	// Always include "general" even if not explicitly created.
	roomSet := map[string]bool{"general": true}
	for _, r := range rooms {
		roomSet[r] = true
	}

	m.historyMu.Lock()
	var toRequest []string
	for r := range roomSet {
		if !m.historyRequested[r] {
			m.historyRequested[r] = true
			toRequest = append(toRequest, r)
		}
	}
	m.historyMu.Unlock()

	for _, room := range toRequest {
		req, err := json.Marshal(proto.PeerMessage{
			Type:  proto.MsgHistoryRequest,
			Room:  room,
			Limit: 200,
		})
		if err != nil {
			continue
		}
		m.SendTo(peerID, req)
		log.Printf("mesh: sent history_request room=%s to %s", room, peerID.Short())
	}
}

// HandleHistoryRequest responds to a history_request from a peer.
func (m *Mesh) HandleHistoryRequest(from proto.PeerID, room string, sinceMs int64, limit int) {
	if m.Store == nil {
		return
	}
	msgs, err := m.Store.RecentMessagesSince(room, sinceMs, limit)
	if err != nil {
		log.Printf("mesh: history_request from %s room=%s: %v", from.Short(), room, err)
		return
	}

	// Look up aliases for from_peer values.
	entries := make([]proto.HistoryEntry, 0, len(msgs))
	for _, msg := range msgs {
		entries = append(entries, proto.HistoryEntry{
			Mid:       msg.Mid,
			From:      string(msg.From),
			FromAlias: m.Store.PeerAlias(msg.From),
			Text:      msg.Text,
			Ts:        msg.Ts,
		})
	}

	chunk, err := json.Marshal(proto.PeerMessage{
		Type:        proto.MsgHistoryChunk,
		Room:        room,
		History:     entries,
		HistoryDone: true,
	})
	if err != nil {
		return
	}
	m.SendTo(from, chunk)
	log.Printf("mesh: sent history_chunk room=%s to %s: %d msgs", room, from.Short(), len(entries))
}

// HandleHistoryChunk saves received history messages and emits history_loaded.
func (m *Mesh) HandleHistoryChunk(room string, entries []proto.HistoryEntry) {
	if m.Store == nil || len(entries) == 0 {
		return
	}
	var saved []proto.ChatMessage
	for _, e := range entries {
		msg := &proto.ChatMessage{
			Mid:   e.Mid,
			MsgID: e.Mid, // mid is already content-addressed for gossipped messages
			From:  proto.PeerID(e.From),
			Room:  room,
			Text:  e.Text,
			Ts:    e.Ts,
		}
		if err := m.Store.SaveMessage(msg); err != nil {
			continue
		}
		saved = append(saved, *msg)
	}
	if len(saved) == 0 {
		return
	}
	m.emit(proto.IpcMessage{
		Type:     proto.EvtHistoryLoaded,
		Room:     room,
		Messages: saved,
	})
	log.Printf("mesh: history_chunk room=%s: %d/%d new messages", room, len(saved), len(entries))
}

// 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:
		}
	}
}