flit/cli/internal/transport/transport.go

// Package transport implements flit's yaw/2.1 signaling + WebRTC session for
// 1:1 ephemeral file transfer, mirroring pwa/src/transport/flit.ts. Trimmed
// from waste-go: no chat, no multi-peer mesh — exactly one peer per session.
package transport

import (
	"context"
	"crypto/hmac"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

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

	"flit/internal/crypto"
)

const (
	bindPrefix = "yaw/2 bind"
	ekeyPrefix = "yaw/2.1 ekey"
	ekeyTimeout = 2 * time.Second
	chunkSize   = 64 * 1024
)

type Config struct {
	SignalURL string
	TurnURL   string
	// TurnSecret holds the coturn use-auth-secret shared secret. Safe to keep
	// in a native config file (unlike the browser, this never ships to a
	// client that can view-source it).
	TurnSecret string
}

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

// PairRoomName derives a stable room name for two known peers — mirrors
// pairRoomName() in pwa/src/pairing/ephemeral.ts. Both sides compute it
// independently from their stored peer IDs, so no QR/invite is needed.
func PairRoomName(a, b string) string {
	if a > b {
		a, b = b, a
	}
	return "flit-pair:" + a + ":" + b
}

func iceServers(cfg Config) []webrtc.ICEServer {
	servers := []webrtc.ICEServer{{URLs: []string{"stun:stun.l.google.com:19302"}}}
	if cfg.TurnURL != "" && cfg.TurnSecret != "" {
		expiry := strconv.FormatInt(time.Now().Add(time.Hour).Unix(), 10) + ":flit"
		mac := hmac.New(sha1.New, []byte(cfg.TurnSecret))
		mac.Write([]byte(expiry))
		cred := base64.StdEncoding.EncodeToString(mac.Sum(nil))
		servers = append(servers, webrtc.ICEServer{
			URLs: []string{cfg.TurnURL}, Username: expiry, Credential: cred,
			CredentialType: webrtc.ICECredentialTypePassword,
		})
	}
	return servers
}

// ── Signaling ─────────────────────────────────────────────────────────────────

type anchorMsg struct {
	Type  string   `json:"type"`
	Nonce string   `json:"nonce,omitempty"`
	ID    string   `json:"id,omitempty"`
	Net   string   `json:"net,omitempty"`
	Sig   string   `json:"sig,omitempty"`
	Peers []string `json:"peers,omitempty"`
	From  string   `json:"from,omitempty"`
	To    string   `json:"to,omitempty"`
	Box   string   `json:"box,omitempty"`
}

type signaling struct {
	conn   *websocket.Conn
	sendCh chan anchorMsg
}

func dialSignaling(ctx context.Context, url string, id *crypto.Identity, netHash string) (*signaling, []string, error) {
	conn, _, err := websocket.Dial(ctx, url, nil)
	if err != nil {
		return nil, nil, fmt.Errorf("dial: %w", err)
	}
	s := &signaling{conn: conn, sendCh: make(chan anchorMsg, 64)}
	go func() {
		for msg := range s.sendCh {
			if err := wsjson.Write(ctx, conn, msg); err != nil {
				return
			}
		}
	}()

	for {
		var msg anchorMsg
		if err := wsjson.Read(ctx, conn, &msg); err != nil {
			return nil, nil, fmt.Errorf("read: %w", err)
		}
		if msg.Type == "challenge" {
			nonceBytes, err := hex.DecodeString(msg.Nonce)
			if err != nil {
				return nil, nil, fmt.Errorf("bad challenge nonce: %w", err)
			}
			sig := id.Sign(append(nonceBytes, []byte(netHash)...))
			s.sendCh <- anchorMsg{Type: "join", ID: id.PeerID(), Net: netHash, Sig: sig}
		} else if msg.Type == "joined" {
			return s, msg.Peers, nil
		}
	}
}

func (s *signaling) sendTo(to, box string) {
	select {
	case s.sendCh <- anchorMsg{Type: "to", To: to, Box: box}:
	default:
	}
}

// ── Session ───────────────────────────────────────────────────────────────────

type FileOffer struct {
	XID  string
	Name string
	Size int64
}

// Session is a 1:1 ephemeral pairing: join a room, connect to exactly one
// peer (optionally pinned to a specific id), exchange files.
type Session struct {
	identity *crypto.Identity
	cfg      Config

	mu        sync.Mutex
	pc        *webrtc.PeerConnection
	dc        *webrtc.DataChannel
	peerID    string
	verified  bool
	esk, epk  *[32]byte
	peerEPK   *[32]byte
	ekeySent  bool
	offered   bool

	sig         *signaling
	pendingRecv *recvState

	OnConnected    func(verified bool)
	OnDisconnected func()
	OnFileOffer    func(FileOffer)
	OnFileProgress func(xid string, received, total int64)
	OnFileDone     func(name string, path string)
}

func NewSession(dataDir string, cfg Config) (*Session, error) {
	id, err := crypto.LoadOrCreate(dataDir)
	if err != nil {
		return nil, err
	}
	return &Session{identity: id, cfg: cfg}, nil
}

func (s *Session) PeerID() string { return s.identity.PeerID() }

// Join connects to the signaling server for roomName and waits for exactly
// one peer (or the pinned trustedPeerID) to complete the handshake.
func (s *Session) Join(ctx context.Context, roomName string, trustedPeerID string) error {
	hash := NetHash(roomName)
	sig, present, err := dialSignaling(ctx, s.cfg.SignalURL, s.identity, hash)
	if err != nil {
		return err
	}
	s.sig = sig

	for _, pid := range present {
		if trustedPeerID != "" && pid != trustedPeerID {
			continue
		}
		if err := s.connectTo(ctx, pid); err != nil {
			return err
		}
		break
	}

	go func() {
		for {
			var msg anchorMsg
			if err := wsjson.Read(ctx, sig.conn, &msg); err != nil {
				return
			}
			switch msg.Type {
			case "peer-join":
				if trustedPeerID != "" && msg.ID != trustedPeerID {
					continue
				}
				_ = s.connectTo(ctx, msg.ID)
			case "from":
				if trustedPeerID != "" && msg.From != trustedPeerID {
					continue
				}
				s.onBox(msg.From, msg.Box)
			}
		}
	}()
	return nil
}

func (s *Session) connectTo(ctx context.Context, peerID string) error {
	s.mu.Lock()
	if s.pc != nil {
		s.mu.Unlock()
		return nil
	}
	s.peerID = peerID
	pc, err := webrtc.NewPeerConnection(webrtc.Configuration{ICEServers: iceServers(s.cfg)})
	if err != nil {
		s.mu.Unlock()
		return err
	}
	s.pc = pc
	kp, err := crypto.GenerateEphemeral()
	if err != nil {
		s.mu.Unlock()
		return err
	}
	s.esk, s.epk = kp.PrivateRaw(), kp.PublicRaw()
	s.mu.Unlock()

	pc.OnDataChannel(func(dc *webrtc.DataChannel) { s.wireDC(dc) })
	pc.OnConnectionStateChange(func(state webrtc.PeerConnectionState) {
		if state == webrtc.PeerConnectionStateFailed ||
			state == webrtc.PeerConnectionStateDisconnected ||
			state == webrtc.PeerConnectionStateClosed {
			s.mu.Lock()
			s.pc = nil
			s.ekeySent = false
			s.offered = false
			s.peerEPK = nil
			s.mu.Unlock()
			if s.OnDisconnected != nil {
				s.OnDisconnected()
			}
		}
	})

	s.sendEkey()
	offerByOrder := s.identity.PeerID() < peerID
	if offerByOrder {
		go func() {
			time.Sleep(ekeyTimeout)
			s.maybeOffer()
		}()
	}
	return nil
}

func (s *Session) sendEkey() {
	s.mu.Lock()
	if s.ekeySent {
		s.mu.Unlock()
		return
	}
	s.ekeySent = true
	epk := *s.epk
	s.mu.Unlock()

	signed := append([]byte(ekeyPrefix), mustHex(s.identity.PeerID())...)
	signed = append(signed, mustHex(s.peerID)...)
	signed = append(signed, epk[:]...)
	msg := map[string]string{
		"kind": "ekey", "v": "yaw/2.1",
		"epk": hex.EncodeToString(epk[:]),
		"sig": s.identity.Sign(signed),
	}
	s.sealAndSend(msg, false)
}

func (s *Session) maybeOffer() {
	s.mu.Lock()
	if s.offered || s.pc == nil {
		s.mu.Unlock()
		return
	}
	s.offered = true
	pc := s.pc
	s.mu.Unlock()

	dc, err := pc.CreateDataChannel("yaw", nil)
	if err != nil {
		log.Printf("flit: create datachannel: %v", err)
		return
	}
	s.wireDC(dc)

	offer, err := pc.CreateOffer(nil)
	if err != nil {
		return
	}
	if err := pc.SetLocalDescription(offer); err != nil {
		return
	}
	<-gatherComplete(pc)
	s.mu.Lock()
	hasEPK := s.peerEPK != nil
	s.mu.Unlock()
	s.sealAndSend(map[string]string{"kind": "offer", "sdp": pc.LocalDescription().SDP}, hasEPK)
}

func (s *Session) onBox(from, box string) {
	plain, usedEph := s.openBox(box)
	if plain == nil {
		return
	}
	var obj map[string]any
	if err := json.Unmarshal(plain, &obj); err != nil {
		return
	}
	switch obj["kind"] {
	case "ekey":
		s.onEkey(obj)
	case "offer":
		s.mu.Lock()
		pc := s.pc
		s.mu.Unlock()
		if pc == nil {
			return
		}
		_ = pc.SetRemoteDescription(webrtc.SessionDescription{Type: webrtc.SDPTypeOffer, SDP: obj["sdp"].(string)})
		answer, err := pc.CreateAnswer(nil)
		if err != nil {
			return
		}
		if err := pc.SetLocalDescription(answer); err != nil {
			return
		}
		<-gatherComplete(pc)
		s.sealAndSend(map[string]string{"kind": "answer", "sdp": pc.LocalDescription().SDP}, usedEph)
	case "answer":
		s.mu.Lock()
		pc := s.pc
		s.mu.Unlock()
		if pc == nil {
			return
		}
		_ = pc.SetRemoteDescription(webrtc.SessionDescription{Type: webrtc.SDPTypeAnswer, SDP: obj["sdp"].(string)})
	}
}

func (s *Session) onEkey(obj map[string]any) {
	s.mu.Lock()
	if s.peerEPK != nil {
		s.mu.Unlock()
		return
	}
	epkHex, _ := obj["epk"].(string)
	sigHex, _ := obj["sig"].(string)
	epkRaw, err := hex.DecodeString(epkHex)
	if err != nil || len(epkRaw) != 32 {
		s.mu.Unlock()
		return
	}
	signed := append([]byte(ekeyPrefix), mustHex(s.peerID)...)
	signed = append(signed, mustHex(s.identity.PeerID())...)
	signed = append(signed, epkRaw...)
	if err := crypto.Verify(s.peerID, signed, sigHex); err != nil {
		s.mu.Unlock()
		return
	}
	var epk [32]byte
	copy(epk[:], epkRaw)
	s.peerEPK = &epk
	s.mu.Unlock()

	s.sendEkey()
	s.maybeOffer()
}

func (s *Session) sealAndSend(obj map[string]string, preferEph bool) {
	data, _ := json.Marshal(obj)
	s.mu.Lock()
	defer s.mu.Unlock()
	var sealed string
	if preferEph && s.peerEPK != nil {
		sealed = crypto.SignalingBox(data, s.peerEPK, s.esk)
	} else {
		recip, err := crypto.CurveFromPeerID(s.peerID)
		if err != nil {
			return
		}
		sealed = crypto.SignalingBox(data, recip, s.identity.CurvePrivateKey())
	}
	s.sig.sendTo(s.peerID, sealed)
}

func (s *Session) openBox(boxB64 string) ([]byte, bool) {
	s.mu.Lock()
	peerEPK, esk := s.peerEPK, s.esk
	peerID := s.peerID
	s.mu.Unlock()

	if peerEPK != nil {
		if pt, err := crypto.SignalingOpen(boxB64, peerEPK, esk); err == nil {
			return pt, true
		}
	}
	senderCurve, err := crypto.CurveFromPeerID(peerID)
	if err != nil {
		return nil, false
	}
	pt, err := crypto.SignalingOpen(boxB64, senderCurve, s.identity.CurvePrivateKey())
	if err != nil {
		return nil, false
	}
	return pt, false
}

// ── control channel + file transfer ───────────────────────────────────────────

func (s *Session) wireDC(dc *webrtc.DataChannel) {
	if dc.Label() == "yaw" {
		s.mu.Lock()
		s.dc = dc
		s.mu.Unlock()
		dc.OnOpen(func() { s.sendHello() })
		dc.OnMessage(func(msg webrtc.DataChannelMessage) { s.onControl(msg.Data) })
		return
	}
	if strings.HasPrefix(dc.Label(), "f:") {
		s.wireFileRecv(dc)
	}
}

func (s *Session) sendHello() {
	s.controlSend(map[string]string{
		"type": "hello", "id": s.identity.PeerID(),
		"sig": s.identity.Sign([]byte(bindPrefix)),
	})
}

func (s *Session) onControl(data []byte) {
	var m map[string]any
	if err := json.Unmarshal(data, &m); err != nil {
		return
	}
	switch m["type"] {
	case "hello":
		s.mu.Lock()
		s.verified = m["id"] == s.peerID
		s.mu.Unlock()
		if s.OnConnected != nil {
			s.OnConnected(s.verified)
		}
	case "file-offer":
		if s.OnFileOffer != nil {
			size, _ := m["size"].(float64)
			s.OnFileOffer(FileOffer{XID: m["xid"].(string), Name: m["name"].(string), Size: int64(size)})
		}
	}
}

func (s *Session) controlSend(obj map[string]string) {
	s.mu.Lock()
	dc := s.dc
	s.mu.Unlock()
	if dc == nil || dc.ReadyState() != webrtc.DataChannelStateOpen {
		return
	}
	data, _ := json.Marshal(obj)
	_ = dc.SendText(string(data))
}

// SendFile offers and streams a file to the connected peer.
func (s *Session) SendFile(path string) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()
	st, err := f.Stat()
	if err != nil {
		return err
	}
	xid := fmt.Sprintf("push-%d", time.Now().UnixNano())
	name := st.Name()

	s.controlSend(map[string]string{"type": "file-offer", "name": name, "size": fmt.Sprint(st.Size()), "xid": xid})

	s.mu.Lock()
	pc := s.pc
	s.mu.Unlock()
	dc, err := pc.CreateDataChannel("f:"+xid, nil)
	if err != nil {
		return err
	}
	opened := make(chan struct{})
	dc.OnOpen(func() { close(opened) })
	select {
	case <-opened:
	case <-time.After(15 * time.Second):
		return fmt.Errorf("timed out waiting for file-accept")
	}

	buf := make([]byte, chunkSize)
	for {
		n, err := f.Read(buf)
		if n > 0 {
			for dc.BufferedAmount() > 1<<20 {
				time.Sleep(10 * time.Millisecond)
			}
			if err := dc.Send(buf[:n]); err != nil {
				return err
			}
		}
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
	}
	return dc.Close()
}

// AcceptOffer accepts an incoming file-offer and writes the transfer to downloadDir.
func (s *Session) AcceptOffer(offer FileOffer, downloadDir string) {
	s.controlSend(map[string]string{"type": "file-accept", "xid": offer.XID})
	s.mu.Lock()
	s.pendingRecv = &recvState{offer: offer, dir: downloadDir}
	s.mu.Unlock()
}

func (s *Session) RejectOffer(xid string) {
	s.controlSend(map[string]string{"type": "file-cancel", "xid": xid})
}

type recvState struct {
	offer FileOffer
	dir   string
	have  int64
	f     *os.File
}

func (s *Session) wireFileRecv(dc *webrtc.DataChannel) {
	xid := strings.TrimPrefix(dc.Label(), "f:")
	s.mu.Lock()
	rs := s.pendingRecv
	s.pendingRecv = nil
	s.mu.Unlock()
	if rs == nil || rs.offer.XID != xid {
		return
	}
	path := rs.dir + "/" + rs.offer.Name
	f, err := os.Create(path)
	if err != nil {
		log.Printf("flit: create %s: %v", path, err)
		return
	}
	rs.f = f

	dc.OnMessage(func(msg webrtc.DataChannelMessage) {
		if _, err := f.Write(msg.Data); err != nil {
			return
		}
		rs.have += int64(len(msg.Data))
		if s.OnFileProgress != nil {
			s.OnFileProgress(xid, rs.have, rs.offer.Size)
		}
	})
	dc.OnClose(func() {
		f.Close()
		if s.OnFileDone != nil {
			s.OnFileDone(rs.offer.Name, path)
		}
	})
}

func mustHex(s string) []byte {
	b, _ := hex.DecodeString(s)
	return b
}

func gatherComplete(pc *webrtc.PeerConnection) <-chan struct{} {
	ch := make(chan struct{})
	if pc.ICEGatheringState() == webrtc.ICEGatheringStateComplete {
		close(ch)
		return ch
	}
	pc.OnICEGatheringStateChange(func(s webrtc.ICEGathererState) {
		if s == webrtc.ICEGathererStateComplete {
			select {
			case <-ch:
			default:
				close(ch)
			}
		}
	})
	go func() {
		time.Sleep(6 * time.Second)
		select {
		case <-ch:
		default:
			close(ch)
		}
	}()
	return ch
}