waste-go/internal/anchor/client.go

// Package anchor implements the YAW/2 anchor client (with YAW/2.1 forward-secret signaling).
// It connects to the anchor WebSocket, handles challenge/join, and routes
// sealed signaling payloads. It manages PeerConnection lifecycle and delegates
// DataChannel handling to internal/mesh.
package anchor

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

	"filippo.io/edwards25519"
	"github.com/pion/webrtc/v3"
	"nhooyr.io/websocket"
	"nhooyr.io/websocket/wsjson"

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

const ekeyTimeout = 2 * time.Second

// peerSession holds per-peer state for one (potential or live) connection.
type peerSession struct {
	pc      *webrtc.PeerConnection
	ekey    *crypto.EphemeralKey // our ephemeral keypair for this session
	peerEPK *[32]byte            // peer's ephemeral pubkey (nil until ekey received)
	fs      bool                 // forward-secret session (both sides exchanged ekey)
	ekeyRx  chan struct{}         // closed when peerEPK is set
}

func newSession(pc *webrtc.PeerConnection) (*peerSession, error) {
	ek, err := crypto.GenerateEphemeral()
	if err != nil {
		return nil, err
	}
	return &peerSession{
		pc:     pc,
		ekey:   ek,
		ekeyRx: make(chan struct{}),
	}, nil
}

func (s *peerSession) close() {
	if s.ekey != nil {
		s.ekey.Wipe()
	}
	if s.pc != nil {
		s.pc.Close()
	}
}

// Run connects to anchorURL, joins networkName, and blocks handling signaling.
// Reconnects automatically on disconnect. Cancel ctx to stop.
func Run(ctx context.Context, anchorURL, networkName string, id *crypto.Identity, m *mesh.Mesh) {
	netHash := hashNetName(networkName)
	for {
		if err := runOnce(ctx, anchorURL, netHash, id, m); err != nil {
			if ctx.Err() != nil {
				return
			}
			log.Printf("anchor: %v — reconnecting in 5s", err)
			select {
			case <-ctx.Done():
				return
			case <-time.After(5 * time.Second):
			}
		}
	}
}

func runOnce(ctx context.Context, anchorURL, netHash string, id *crypto.Identity, m *mesh.Mesh) error {
	conn, _, err := websocket.Dial(ctx, anchorURL, nil)
	if err != nil {
		return fmt.Errorf("dial: %w", err)
	}
	defer conn.Close(websocket.StatusNormalClosure, "bye")
	log.Printf("anchor: connected to %s", anchorURL)

	sendCh := make(chan proto.AnchorMessage, 64)
	go func() {
		for msg := range sendCh {
			if err := wsjson.Write(ctx, conn, msg); err != nil {
				return
			}
		}
	}()

	var (
		mu       sync.RWMutex
		sessions = make(map[proto.PeerID]*peerSession)
	)

	s := &sender{id: id, sendCh: sendCh}

	// Drain gossip-discovered peers and initiate offers. Stopped when runOnce
	// returns (via drainDone) so stale sessions from a dead connection are never reused.
	drainDone := make(chan struct{})
	defer close(drainDone)
	go func() {
		for {
			select {
			case pid, ok := <-m.PendingConnect:
				if !ok {
					return
				}
				mu.RLock()
				_, already := sessions[pid]
				mu.RUnlock()
				if already {
					continue
				}
				// Use the same lexicographic tiebreak as anchor join to avoid
				// both sides trying to offer simultaneously.
				if strings.Compare(string(id.PeerID()), string(pid)) <= 0 {
					continue
				}
				go func(pid proto.PeerID) {
					sess, err := startOffer(ctx, pid, id, m, s)
					if err != nil {
						log.Printf("anchor: gossip offer to %s: %v", pid.Short(), err)
						return
					}
					mu.Lock()
					sessions[pid] = sess
					mu.Unlock()
				}(pid)
			case <-drainDone:
				return
			case <-ctx.Done():
				return
			}
		}
	}()

	for {
		var msg proto.AnchorMessage
		if err := wsjson.Read(ctx, conn, &msg); err != nil {
			return fmt.Errorf("read: %w", err)
		}

		switch msg.Type {

		case proto.AnchorChallenge:
			nonceBytes, err := hex.DecodeString(msg.Nonce)
			if err != nil {
				return fmt.Errorf("bad challenge nonce: %w", err)
			}
			sig := id.Sign(append(nonceBytes, []byte(netHash)...))
			sendCh <- proto.AnchorMessage{
				Type: proto.AnchorJoin,
				ID:   string(id.PeerID()),
				Net:  netHash,
				Sig:  sig,
			}

		case proto.AnchorJoined:
			log.Printf("anchor: joined network, %d peer(s) present", len(msg.Peers))
			for _, peerHex := range msg.Peers {
				pid := proto.PeerID(peerHex)
				if strings.Compare(string(id.PeerID()), peerHex) > 0 {
					go func(pid proto.PeerID) {
						sess, err := startOffer(ctx, pid, id, m, s)
						if err != nil {
							log.Printf("anchor: offer to %s: %v", pid.Short(), err)
							return
						}
						mu.Lock()
						sessions[pid] = sess
						mu.Unlock()
					}(pid)
				}
			}

		case proto.AnchorPeerJoin:
			pid := proto.PeerID(msg.ID)
			log.Printf("anchor: peer joined: %s", pid.Short())
			if strings.Compare(string(id.PeerID()), msg.ID) > 0 {
				go func(pid proto.PeerID) {
					sess, err := startOffer(ctx, pid, id, m, s)
					if err != nil {
						log.Printf("anchor: offer to %s: %v", pid.Short(), err)
						return
					}
					mu.Lock()
					sessions[pid] = sess
					mu.Unlock()
				}(pid)
			}

		case proto.AnchorPeerLeave:
			pid := proto.PeerID(msg.ID)
			mu.Lock()
			if sess, ok := sessions[pid]; ok {
				sess.close()
				delete(sessions, pid)
			}
			mu.Unlock()
			log.Printf("anchor: peer left: %s", pid.Short())

		case proto.AnchorFrom:
			fromID := proto.PeerID(msg.From)

			mu.RLock()
			sess := sessions[fromID]
			mu.RUnlock()

			// Determine which key to try for opening the box.
			// If we already have the peer's ephemeral key, try ephemeral first.
			payload, usedEph, err := openBoxAuto(msg.Box, fromID, id, sess)
			if err != nil {
				log.Printf("anchor: open box from %s: %v", fromID.Short(), err)
				continue
			}

			if payload.Kind == proto.SigEkey {
				// Process ekey under static keys (we opened it correctly above).
				mu.Lock()
				if sess == nil {
					// Answerer: we haven't created a session yet, do it now.
					pc, err := newPC()
					if err != nil {
						mu.Unlock()
						log.Printf("anchor: new PC for answerer: %v", err)
						continue
					}
					sess, err = newSession(pc)
					if err != nil {
						pc.Close()
						mu.Unlock()
						log.Printf("anchor: ephemeral key gen: %v", err)
						continue
					}
					sessions[fromID] = sess
					// Send our ekey back immediately.
					go sendEkey(fromID, sess, id, s)
				}
				if err := receiveEkey(payload, fromID, id, sess); err != nil {
					mu.Unlock()
					log.Printf("anchor: bad ekey from %s: %v", fromID.Short(), err)
					continue
				}
				mu.Unlock()
				_ = usedEph
				continue
			}

			dispatchSignaling(ctx, payload, usedEph, fromID, id, m, s, sessions, &mu)

		case proto.AnchorNoPeer:
			log.Printf("anchor: no such peer: %s", proto.PeerID(msg.ID).Short())
		}
	}
}

// ── ekey helpers ─────────────────────────────────────────────────────────────

// sendEkey seals and transmits our ephemeral public key to peerID.
func sendEkey(peerID proto.PeerID, sess *peerSession, id *crypto.Identity, s *sender) {
	epkHex := hex.EncodeToString(sess.ekey.PublicRaw()[:])
	sig := signEkey(id, peerID, sess.ekey.PublicRaw())
	payload := proto.SignalingPayload{
		Kind:    proto.SigEkey,
		V:       "yaw/2.1",
		EPK:     epkHex,
		EkeySig: sig,
	}
	// ekey is always sealed under static keys (§5.4′ (a)).
	if err := s.SendTo(peerID, payload); err != nil {
		log.Printf("anchor: send ekey to %s: %v", peerID.Short(), err)
	}
}

// receiveEkey validates and stores the peer's ephemeral public key.
func receiveEkey(payload proto.SignalingPayload, from proto.PeerID, id *crypto.Identity, sess *peerSession) error {
	epkBytes, err := hex.DecodeString(payload.EPK)
	if err != nil || len(epkBytes) != 32 {
		return fmt.Errorf("bad epk hex")
	}
	// Verify sig: "yaw/2.1 ekey" || from_id_raw(32) || our_id_raw(32) || epk_raw(32)
	fromRaw, err := hex.DecodeString(string(from))
	if err != nil {
		return fmt.Errorf("bad from id")
	}
	ourRaw, err := hex.DecodeString(string(id.PeerID()))
	if err != nil {
		return fmt.Errorf("bad own id")
	}
	msg := append([]byte("yaw/2.1 ekey"), fromRaw...)
	msg = append(msg, ourRaw...)
	msg = append(msg, epkBytes...)
	if err := crypto.Verify(string(from), msg, payload.EkeySig); err != nil {
		return fmt.Errorf("ekey sig: %w", err)
	}
	var epk [32]byte
	copy(epk[:], epkBytes)
	sess.peerEPK = &epk
	sess.fs = true
	close(sess.ekeyRx) // signal waiters
	return nil
}

// signEkey produces the Ed25519 signature for our ekey message.
// Input: "yaw/2.1 ekey" || our_id_raw(32) || peer_id_raw(32) || epk_raw(32)
func signEkey(id *crypto.Identity, peerID proto.PeerID, epk *[32]byte) string {
	ourRaw, _ := hex.DecodeString(string(id.PeerID()))
	peerRaw, _ := hex.DecodeString(string(peerID))
	msg := append([]byte("yaw/2.1 ekey"), ourRaw...)
	msg = append(msg, peerRaw...)
	msg = append(msg, epk[:]...)
	return id.Sign(msg)
}

// ── signaling dispatch ────────────────────────────────────────────────────────

func dispatchSignaling(
	ctx context.Context,
	payload proto.SignalingPayload,
	usedEph bool,
	fromID proto.PeerID,
	id *crypto.Identity,
	m *mesh.Mesh,
	s *sender,
	sessions map[proto.PeerID]*peerSession,
	mu *sync.RWMutex,
) {
	switch payload.Kind {

	case proto.SigOffer:
		go func() {
			mu.Lock()
			sess := sessions[fromID]
			mu.Unlock()

			pc, err := answerOffer(ctx, payload, fromID, id, m, s, sess)
			if err != nil {
				log.Printf("anchor: answer to %s: %v", fromID.Short(), err)
				return
			}
			mu.Lock()
			if existing, ok := sessions[fromID]; ok && existing != sess {
				// Session was already replaced; close the new PC.
				pc.Close()
			} else {
				if sess == nil {
					// No session yet (2.0 peer — no ekey): create a minimal one.
					sess2, _ := newSession(pc)
					if sess2 != nil {
						sess2.fs = false
						sessions[fromID] = sess2
					}
				} else {
					sess.pc = pc
				}
			}
			mu.Unlock()
		}()

	case proto.SigAnswer:
		mu.RLock()
		sess, ok := sessions[fromID]
		mu.RUnlock()
		if !ok || sess.pc == nil {
			log.Printf("anchor: answer from %s but no PeerConnection", fromID.Short())
			return
		}
		if err := sess.pc.SetRemoteDescription(webrtc.SessionDescription{
			Type: webrtc.SDPTypeAnswer,
			SDP:  payload.SDP,
		}); err != nil {
			log.Printf("anchor: set remote answer from %s: %v", fromID.Short(), err)
		}

	case proto.SigCandidate:
		mu.RLock()
		sess, ok := sessions[fromID]
		mu.RUnlock()
		if !ok || sess.pc == nil {
			return
		}
		if err := sess.pc.AddICECandidate(webrtc.ICECandidateInit{
			Candidate:     payload.Cand,
			SDPMid:        strPtr(payload.Mid),
			SDPMLineIndex: uint16Ptr(uint16(payload.MLine)),
		}); err != nil {
			log.Printf("anchor: add candidate from %s: %v", fromID.Short(), err)
		}

	case proto.SigBye:
		mu.Lock()
		if sess, ok := sessions[fromID]; ok {
			sess.close()
			delete(sessions, fromID)
		}
		mu.Unlock()
	}
}

// ── offer / answer helpers ────────────────────────────────────────────────────

// startOffer creates a session, sends our ekey, waits up to ekeyTimeout for
// the peer's ekey, then sends the offer (ephemeral or static).
func startOffer(ctx context.Context, peerID proto.PeerID, id *crypto.Identity, m *mesh.Mesh, s *sender) (*peerSession, error) {
	pc, err := newPC()
	if err != nil {
		return nil, err
	}
	sess, err := newSession(pc)
	if err != nil {
		pc.Close()
		return nil, err
	}

	dc, err := pc.CreateDataChannel("yaw", &webrtc.DataChannelInit{Ordered: boolPtr(true)})
	if err != nil {
		sess.close()
		return nil, err
	}
	mesh.WireDataChannel(dc, pc, peerID, id, m)
	mesh.WireCandidateTrickle(pc, peerID, s)

	// Handle file DataChannels opened by the remote peer.
	pc.OnDataChannel(func(dc *webrtc.DataChannel) {
		if strings.HasPrefix(dc.Label(), "f:") {
			xid := strings.TrimPrefix(dc.Label(), "f:")
			m.HandleInboundFileDC(dc, xid, peerID)
		}
	})

	// Send our ekey immediately.
	sendEkey(peerID, sess, id, s)

	// Build the offer in a goroutine so we don't block the read loop.
	go func() {
		// Wait for peer's ekey or fall back after timeout.
		select {
		case <-sess.ekeyRx:
			log.Printf("anchor: 2.1 FS offer to %s", peerID.Short())
		case <-time.After(ekeyTimeout):
			log.Printf("anchor: 2.0 fallback offer to %s (no ekey received)", peerID.Short())
		case <-ctx.Done():
			return
		}

		sdpOffer, err := pc.CreateOffer(nil)
		if err != nil {
			log.Printf("anchor: create offer to %s: %v", peerID.Short(), err)
			return
		}
		if err := pc.SetLocalDescription(sdpOffer); err != nil {
			log.Printf("anchor: set local offer to %s: %v", peerID.Short(), err)
			return
		}
		payload := proto.SignalingPayload{Kind: proto.SigOffer, SDP: sdpOffer.SDP}
		if err := s.sealAndSend(peerID, payload, sess); err != nil {
			log.Printf("anchor: send offer to %s: %v", peerID.Short(), err)
		}
	}()

	return sess, nil
}

// answerOffer processes an incoming offer and returns the PeerConnection.
func answerOffer(ctx context.Context, payload proto.SignalingPayload, fromID proto.PeerID, id *crypto.Identity, m *mesh.Mesh, s *sender, sess *peerSession) (*webrtc.PeerConnection, error) {
	pc, err := newPC()
	if err != nil {
		return nil, err
	}

	pc.OnDataChannel(func(dc *webrtc.DataChannel) {
		switch {
		case dc.Label() == "yaw":
			mesh.WireDataChannel(dc, pc, fromID, id, m)
		case strings.HasPrefix(dc.Label(), "f:"):
			xid := strings.TrimPrefix(dc.Label(), "f:")
			m.HandleInboundFileDC(dc, xid, fromID)
		}
	})
	mesh.WireCandidateTrickle(pc, fromID, s)

	if err := pc.SetRemoteDescription(webrtc.SessionDescription{
		Type: webrtc.SDPTypeOffer,
		SDP:  payload.SDP,
	}); err != nil {
		pc.Close()
		return nil, err
	}
	sdpAnswer, err := pc.CreateAnswer(nil)
	if err != nil {
		pc.Close()
		return nil, err
	}
	if err := pc.SetLocalDescription(sdpAnswer); err != nil {
		pc.Close()
		return nil, err
	}

	answerPayload := proto.SignalingPayload{Kind: proto.SigAnswer, SDP: sdpAnswer.SDP}
	if err := s.sealAndSend(fromID, answerPayload, sess); err != nil {
		pc.Close()
		return nil, err
	}
	return pc, nil
}

// ── sender implements mesh.Anchor ────────────────────────────────────────────

type sender struct {
	id     *crypto.Identity
	sendCh chan proto.AnchorMessage
}

// SendTo seals with STATIC keys (used for ekey and 2.0 fallback).
func (s *sender) SendTo(peerID proto.PeerID, payload proto.SignalingPayload) error {
	plaintext, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	recipientCurve, err := curveFromPeerID(peerID)
	if err != nil {
		return fmt.Errorf("derive curve key for %s: %w", peerID.Short(), err)
	}
	sealed := crypto.SignalingBox(plaintext, recipientCurve, s.id.CurvePrivateKey())
	return s.enqueue(peerID, sealed)
}

// sealAndSend seals with ephemeral keys if available, static otherwise.
func (s *sender) sealAndSend(peerID proto.PeerID, payload proto.SignalingPayload, sess *peerSession) error {
	plaintext, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	var sealed string
	if sess != nil && sess.peerEPK != nil {
		// Ephemeral seal (2.1 FS).
		sealed = crypto.SignalingBox(plaintext, sess.peerEPK, sess.ekey.PrivateRaw())
	} else {
		// Static seal (2.0 compatible).
		recipientCurve, err := curveFromPeerID(peerID)
		if err != nil {
			return fmt.Errorf("derive curve key for %s: %w", peerID.Short(), err)
		}
		sealed = crypto.SignalingBox(plaintext, recipientCurve, s.id.CurvePrivateKey())
	}
	return s.enqueue(peerID, sealed)
}

func (s *sender) enqueue(peerID proto.PeerID, box string) error {
	select {
	case s.sendCh <- proto.AnchorMessage{Type: proto.AnchorTo, To: string(peerID), Box: box}:
		return nil
	default:
		return fmt.Errorf("send queue full")
	}
}

func (s *sender) LocalID() proto.PeerID { return s.id.PeerID() }

// ── box opening ───────────────────────────────────────────────────────────────

// openBoxAuto opens an incoming box, trying ephemeral keys first (if available),
// then falling back to static. Returns the payload and whether ephemeral was used.
func openBoxAuto(b64box string, fromID proto.PeerID, localID *crypto.Identity, sess *peerSession) (proto.SignalingPayload, bool, error) {
	senderCurve, err := curveFromPeerID(fromID)
	if err != nil {
		return proto.SignalingPayload{}, false, err
	}

	// Try ephemeral first if we have the peer's epk.
	if sess != nil && sess.peerEPK != nil {
		if pt, err := crypto.SignalingOpen(b64box, sess.peerEPK, sess.ekey.PrivateRaw()); err == nil {
			var p proto.SignalingPayload
			if err := json.Unmarshal(pt, &p); err == nil {
				return p, true, nil
			}
		}
	}

	// Fall back to static keys.
	plaintext, err := crypto.SignalingOpen(b64box, senderCurve, localID.CurvePrivateKey())
	if err != nil {
		return proto.SignalingPayload{}, false, fmt.Errorf("open box: %w", err)
	}
	var p proto.SignalingPayload
	return p, false, json.Unmarshal(plaintext, &p)
}

// ── helpers ───────────────────────────────────────────────────────────────────

func curveFromPeerID(id proto.PeerID) (*[32]byte, error) {
	pubBytes, err := hex.DecodeString(string(id))
	if err != nil || len(pubBytes) != 32 {
		return nil, fmt.Errorf("invalid peer id %q", id)
	}
	edPoint, err := new(edwards25519.Point).SetBytes(pubBytes)
	if err != nil {
		return nil, fmt.Errorf("ed25519 point: %w", err)
	}
	mont := edPoint.BytesMontgomery()
	var out [32]byte
	copy(out[:], mont)
	return &out, nil
}

func hashNetName(name string) string {
	h := sha256.Sum256([]byte("yaw2-net:" + name))
	return hex.EncodeToString(h[:])
}

func newPC() (*webrtc.PeerConnection, error) {
	return webrtc.NewPeerConnection(webrtc.Configuration{
		ICEServers: []webrtc.ICEServer{{URLs: []string{"stun:stun.l.google.com:19302"}}},
	})
}

func boolPtr(b bool) *bool       { return &b }
func strPtr(s string) *string    { return &s }
func uint16Ptr(v uint16) *uint16 { return &v }