waste-go/internal/mesh/peer.go

// Package mesh/peer exports WebRTC DataChannel helpers used by the anchor client.
package mesh

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log"
	"strings"
	"sync"
	"time"

	"github.com/pion/webrtc/v3"

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

// Anchor is the signaling channel used to exchange sealed offers/answers/candidates.
// Implemented by internal/anchor.Client.
type Anchor interface {
	SendTo(peerID proto.PeerID, payload proto.SignalingPayload) error
	LocalID() proto.PeerID
}

// WireDataChannel sets up open/message/close handlers on a "yaw" DataChannel.
// Safe to call whether the channel is already open or not (§6 open-race).
func WireDataChannel(
	dc *webrtc.DataChannel,
	pc *webrtc.PeerConnection,
	peerID proto.PeerID,
	id *crypto.Identity,
	m *Mesh,
) {
	sendCh := make(chan []byte, 64)

	var once sync.Once
	doOpen := func() {
		log.Printf("peer: DataChannel open with %s", peerID.Short())

		// Send hello — bind our identity to this DTLS session.
		localFP, remoteFP := dtlsFingerprints(pc)
		bindBytes := proto.HelloBindString(localFP, remoteFP)
		hello := proto.HelloMessage{
			Type:   "hello",
			ID:     string(id.PeerID()),
			Nick:   id.Alias,
			Caps:   []string{"chat", "file"},
			Sig:    id.Sign(bindBytes),
			Invite: m.InviteString,
		}
		helloJSON, _ := json.Marshal(hello)
		if err := dc.SendText(string(helloJSON)); err != nil {
			log.Printf("peer: send hello to %s: %v", peerID.Short(), err)
			return
		}

		peerConn := &PeerConn{
			Info: proto.PeerInfo{ID: peerID, Alias: string(peerID.Short())},
			Send: sendCh,
			PC:   pc,
		}
		m.AddPeer(peerConn)

		// Request the peer's file list immediately after connect.
		if req, err := json.Marshal(proto.PeerMessage{Type: proto.MsgFileListReq}); err == nil {
			sendCh <- req
		}

		go func() {
			for payload := range sendCh {
				if err := dc.SendText(string(payload)); err != nil {
					log.Printf("peer: send to %s: %v", peerID.Short(), err)
					return
				}
			}
		}()
	}

	dc.OnOpen(func() { once.Do(doOpen) })
	// §6 gotcha: answerer's DC may already be open when OnDataChannel fires.
	if dc.ReadyState() == webrtc.DataChannelStateOpen {
		once.Do(doOpen)
	}

	dc.OnMessage(func(msg webrtc.DataChannelMessage) {
		if msg.IsString {
			handleDCMessage(msg.Data, peerID, id, m)
		}
	})

	dc.OnClose(func() {
		log.Printf("peer: DataChannel closed with %s", peerID.Short())
		close(sendCh)
		m.RemovePeer(peerID)
		pc.Close()
	})

	pc.OnConnectionStateChange(func(state webrtc.PeerConnectionState) {
		m.Emit(proto.IpcMessage{
			Type:      proto.EvtPeerStatus,
			PeerID:    &peerID,
			ConnState: state.String(),
		})
		if state == webrtc.PeerConnectionStateConnected {
			go emitICEStats(pc, peerID, m)
		}
	})
}

// WireCandidateTrickle seals and forwards each ICE candidate via the anchor as it arrives.
func WireCandidateTrickle(pc *webrtc.PeerConnection, peerID proto.PeerID, anchor Anchor) {
	pc.OnICECandidate(func(c *webrtc.ICECandidate) {
		if c == nil {
			return
		}
		init := c.ToJSON()
		mid := ""
		if init.SDPMid != nil {
			mid = *init.SDPMid
		}
		mline := 0
		if init.SDPMLineIndex != nil {
			mline = int(*init.SDPMLineIndex)
		}
		if err := anchor.SendTo(peerID, proto.SignalingPayload{
			Kind:  proto.SigCandidate,
			Cand:  init.Candidate,
			Mid:   mid,
			MLine: mline,
		}); err != nil {
			log.Printf("peer: trickle candidate to %s: %v", peerID.Short(), err)
		}
	})
}

// handleDCMessage dispatches a raw DataChannel text message.
func handleDCMessage(data []byte, from proto.PeerID, localID *crypto.Identity, m *Mesh) {
	var probe struct {
		Type string `json:"type"`
	}
	if err := json.Unmarshal(data, &probe); err != nil {
		return
	}

	if probe.Type == "hello" {
		var hello proto.HelloMessage
		if err := json.Unmarshal(data, &hello); err != nil {
			log.Printf("peer: bad hello from %s: %v", from.Short(), err)
			return
		}

		// Invite enforcement (waste-go extension).
		if m.RequireInvite {
			if hello.Invite == "" {
				log.Printf("peer: rejecting %s — no invite presented (RequireInvite=true)", from.Short())
				m.Emit(proto.IpcMessage{
					Type:         proto.EvtError,
					ErrorMessage: fmt.Sprintf("peer %s rejected: no invite", from.Short()),
				})
				return
			}
			inv, err := invite.Decode(hello.Invite)
			if err != nil || !inv.IsSigned() {
				log.Printf("peer: rejecting %s — invite not signed: %v", from.Short(), err)
				m.Emit(proto.IpcMessage{
					Type:         proto.EvtError,
					ErrorMessage: fmt.Sprintf("peer %s rejected: invite not signed", from.Short()),
				})
				return
			}
			trusted := m.trustedPeerIDs()
			if err := invite.Verify(inv, trusted, crypto.Verify); err != nil {
				log.Printf("peer: rejecting %s — %v", from.Short(), err)
				m.Emit(proto.IpcMessage{
					Type:         proto.EvtError,
					ErrorMessage: fmt.Sprintf("peer %s rejected: %v", from.Short(), err),
				})
				return
			}
		}

		// Update alias once we have the verified nick.
		m.mu.Lock()
		if conn, ok := m.peers[from]; ok {
			conn.Info.Alias = hello.Nick
			conn.Info.PublicKey = hello.ID
		}
		m.mu.Unlock()
		m.UpdatePeerAlias(from, hello.Nick)
		m.Emit(proto.IpcMessage{
			Type:   proto.EvtSessionReady,
			PeerID: peerIDPtr(from),
			Nick:   hello.Nick,
		})
		// Tell the new peer about everyone we can currently see.
		go m.sendGossipTo(from)
		return
	}

	var msg proto.PeerMessage
	if err := json.Unmarshal(data, &msg); err != nil {
		log.Printf("peer: bad message from %s: %v", from.Short(), err)
		return
	}
	dispatchPeerMessage(msg, from, m)
}

func dispatchPeerMessage(msg proto.PeerMessage, from proto.PeerID, m *Mesh) {
	switch msg.Type {
	case proto.MsgChat:
		chat := &proto.ChatMessage{
			Mid:  midOrRandom(msg.Mid),
			From: from,
			Room: msg.Room,
			Text: msg.Text,
			Ts:   msg.Ts,
		}
		m.SaveMessage(chat)
		m.Emit(proto.IpcMessage{Type: proto.EvtMessageReceived, Message: chat})

	case proto.MsgPm:
		// Private message — reconstruct as ChatMessage for IPC/storage using dm:<short-id> room.
		chat := &proto.ChatMessage{
			Mid:  midOrRandom(msg.Mid),
			From: from,
			Room: "dm:" + from.Short(),
			Text: msg.Text,
			Ts:   msg.Ts,
		}
		m.SaveMessage(chat)
		m.Emit(proto.IpcMessage{Type: proto.EvtMessageReceived, Message: chat})

	case proto.MsgFileListReq:
		var files []proto.FileEntry
		if m.ScanFiles != nil {
			files = m.ScanFiles()
		} else {
			files = m.ScanShareDir()
		}
		resp, err := json.Marshal(proto.PeerMessage{
			Type:         proto.MsgFileListResp,
			FileListResp: &proto.FileListResp{Files: files},
		})
		if err != nil {
			return
		}
		m.SendTo(from, resp)

	case proto.MsgFileListResp:
		if msg.FileListResp != nil {
			m.Emit(proto.IpcMessage{
				Type:   proto.EvtFileList,
				PeerID: peerIDPtr(from),
				Files:  msg.FileListResp.Files,
			})
		}

	case proto.MsgFileOffer:
		m.Emit(proto.IpcMessage{
			Type:   proto.EvtIncomingFile,
			PeerID: peerIDPtr(from),
			Offer:  &proto.FileOffer{Xid: msg.Xid, Name: msg.Name, Size: msg.Size, SHA256: msg.SHA256},
		})
		m.acceptIncoming(msg, from)

	case proto.MsgFileAccept:
		m.startSend(msg.Xid, from, msg.ResumeOffset)

	case proto.MsgFileCancel:
		log.Printf("mesh: file-cancel from %s xid=%s reason=%s", from.Short(), msg.Xid, msg.Reason)

	case proto.MsgFileDone:
		log.Printf("mesh: file-done from %s xid=%s", from.Short(), msg.Xid)

	case proto.MsgPeerGossip:
		if msg.Gossip == nil {
			return
		}
		// Build set of peers we already know about (connected + self).
		connected := m.ConnectedPeers()
		known := make(map[proto.PeerID]bool, len(connected)+1)
		known[m.Identity.PeerID()] = true
		for _, p := range connected {
			known[p.ID] = true
		}
		newPeers := 0
		for _, entry := range msg.Gossip.Peers {
			if known[entry.Peer.ID] {
				continue
			}
			select {
			case m.PendingConnect <- entry.Peer.ID:
				newPeers++
			default:
			}
		}
		log.Printf("mesh: gossip from %s: %d hints, %d new", from.Short(), len(msg.Gossip.Peers), newPeers)
	case proto.MsgPing:
		log.Printf("mesh: ping from %s", from.Short())
	case proto.MsgPong:
		log.Printf("mesh: pong from %s", from.Short())
	default:
		log.Printf("mesh: unknown message type %q from %s", msg.Type, from.Short())
	}
}

// emitICEStats reads the active (nominated) ICE candidate pair from pion's stats
// and emits a peer_status event with the candidate type and remote address.
func emitICEStats(pc *webrtc.PeerConnection, peerID proto.PeerID, m *Mesh) {
	stats := pc.GetStats()
	candidateType := "unknown"
	remoteAddress := ""

	// Find the nominated candidate pair and its remote candidate.
	for _, s := range stats {
		pair, ok := s.(webrtc.ICECandidatePairStats)
		if !ok || !pair.Nominated {
			continue
		}
		// Look up remote candidate by its ID.
		if rc, ok2 := stats[pair.RemoteCandidateID]; ok2 {
			remote, ok3 := rc.(webrtc.ICECandidateStats)
			if ok3 {
				candidateType = remote.CandidateType.String()
				if remote.Port > 0 {
					remoteAddress = fmt.Sprintf("%s:%d", remote.IP, remote.Port)
				} else {
					remoteAddress = remote.IP
				}
			}
		}
		break
	}

	m.Emit(proto.IpcMessage{
		Type:          proto.EvtPeerStatus,
		PeerID:        &peerID,
		CandidateType: candidateType,
		RemoteAddress: remoteAddress,
	})
}

// midOrRandom returns mid if non-empty, otherwise generates a random 16-byte hex string.
// Ensures every stored message has a unique mid even from peers that don't send one.
func midOrRandom(mid string) string {
	if mid != "" {
		return mid
	}
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return hex.EncodeToString([]byte(time.Now().String()))
	}
	return hex.EncodeToString(b)
}

func dtlsFingerprints(pc *webrtc.PeerConnection) (local, remote []byte) {
	if ld := pc.LocalDescription(); ld != nil {
		local = fingerprintFromSDP(ld.SDP)
	}
	if rd := pc.RemoteDescription(); rd != nil {
		remote = fingerprintFromSDP(rd.SDP)
	}
	return
}

func fingerprintFromSDP(sdp string) []byte {
	for _, line := range strings.Split(sdp, "\n") {
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, "a=fingerprint:sha-256 ") {
			hexStr := strings.ReplaceAll(strings.TrimPrefix(line, "a=fingerprint:sha-256 "), ":", "")
			if b, err := hex.DecodeString(hexStr); err == nil {
				return b
			}
		}
	}
	return nil
}

// sendGossipTo sends the current live peer list to peerID so they can discover
// and connect to peers they haven't met yet. Called after hello is verified.
func (m *Mesh) sendGossipTo(peerID proto.PeerID) {
	peers := m.ConnectedPeers()
	entries := make([]proto.GossipEntry, 0, len(peers))
	for _, p := range peers {
		if p.ID == peerID {
			continue // don't include the recipient in their own gossip
		}
		entries = append(entries, proto.GossipEntry{Peer: p})
	}
	if len(entries) == 0 {
		return
	}
	wire, err := json.Marshal(proto.PeerMessage{
		Type:   proto.MsgPeerGossip,
		Gossip: &proto.PeerGossip{Peers: entries},
	})
	if err != nil {
		return
	}
	m.SendTo(peerID, wire)
	log.Printf("mesh: sent gossip to %s: %d peer(s)", peerID.Short(), len(entries))
}

func peerIDPtr(p proto.PeerID) *proto.PeerID { return &p }