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"
	"log"
	"strings"
	"sync"
	"time"

	"github.com/pion/webrtc/v3"

	"github.com/waste-go/internal/crypto"
	"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),
		}
		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,
		}
		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()
	})
}

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

	case proto.MsgPeerGossip:
		if msg.Gossip != nil {
			log.Printf("mesh: gossip from %s: %d peer hints", from.Short(), len(msg.Gossip.Peers))
		}
	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())
	}
}

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

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