Commilitia-Drop/desktop/engine/session.go

package engine

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pion/ice/v4"
	"github.com/pion/webrtc/v4"
)

const (
	roleSender   = "sender"
	roleReceiver = "receiver"
)

// session 是一条 P2P 传输的全部状态：pion 连接 + 数据通道 + 角色相关的收发循环。
type session struct {
	eng  *Engine
	id   string
	peer string
	role string

	pc *webrtc.PeerConnection
	dc *webrtc.DataChannel

	// 发送端
	filePath string
	fileSize int64

	// 计数器（跨 goroutine：发送循环 / pion 回调）。
	bytesSent atomic.Int64 // 已塞进本地 SCTP 缓冲的字节
	received  atomic.Int64 // 接收端已写盘字节
	acked     atomic.Int64 // 接收端经 ack 回传的已收字节（发送端进度/完成依据）

	// 接收端落盘
	out     *os.File
	outName string

	bufLow    chan struct{} // OnBufferedAmountLow 唤醒发送循环
	done      chan struct{} // 会话终结信号
	closeOnce sync.Once

	stateMu sync.Mutex
	state   string
}

func newSession(e *Engine, id, peer, role string, pc *webrtc.PeerConnection) *session {
	return &session{
		eng:    e,
		id:     id,
		peer:   peer,
		role:   role,
		pc:     pc,
		bufLow: make(chan struct{}, 1),
		done:   make(chan struct{}),
	}
}

// newPeerConnection 建一个配了大 SCTP 接收窗的 pion 连接。SetSCTPMaxReceiveBufferSize 把
// 接收窗顶到 4MB（远高于 WebKit 写死的 256KB），是吞吐专项的核心——使本端作接收端时
// `吞吐 ≈ rwnd/RTT` 的天花板抬高一个量级。pion 默认用真实 IP 的 host 候选（不做 mDNS
// 混淆），故同内网天然落到 host↔host 直连，避开 mac WKWebView 的 srflx 退化。
func newPeerConnection(iceServersJSON string) (*webrtc.PeerConnection, error) {
	se := webrtc.SettingEngine{}
	se.SetSCTPMaxReceiveBufferSize(sctpReceiveBuffer)
	// 解析对端的 mDNS（.local）候选。iOS / Safari 出于隐私只播 mDNS host 候选；pion 默认不
	// 解析，便无法与之成 host↔host、退而走 srflx/relay 慢路径（实测 mac→iOS 落到 relay↔srflx、
	// 1MB/s 且停滞）。QueryOnly＝解析对端 mDNS、本端仍播真实 IP host 候选（对端可直接用），
	// 双向打通同内网直连——这是吞吐的前提（host↔host 低 RTT 下 iOS 的 256KB rwnd 才非约束）。
	se.SetICEMulticastDNSMode(ice.MulticastDNSModeQueryOnly)
	api := webrtc.NewAPI(webrtc.WithSettingEngine(se))
	pc, err := api.NewPeerConnection(webrtc.Configuration{ICEServers: parseICEServers(iceServersJSON)})
	if err != nil {
		return nil, fmt.Errorf("new peerconnection: %w", err)
	}
	return pc, nil
}

// wireCommon 绑两端共用的回调：trickle ICE 候选外发、连接状态变化。
func (s *session) wireCommon() {
	s.pc.OnICECandidate(func(c *webrtc.ICECandidate) {
		if c == nil {
			s.eng.ev.OnLog(fmt.Sprintf("session %s local gathering done", s.id))
			return
		}
		s.eng.ev.OnLog(fmt.Sprintf("session %s cand local: %s/%s", s.id, c.Typ, c.Protocol))
		init := c.ToJSON()
		s.emitSignal(s.peer, nil, &init)
	})
	// ICE 连接态变化日志：诊断「卡在等待对方接受 → 退中继」时连到哪一步（checking/connected/
	// failed/disconnected）。
	s.pc.OnICEConnectionStateChange(func(st webrtc.ICEConnectionState) {
		s.eng.ev.OnLog(fmt.Sprintf("session %s ice: %s", s.id, st))
	})
	s.pc.OnConnectionStateChange(func(st webrtc.PeerConnectionState) {
		switch st {
		case webrtc.PeerConnectionStateConnected:
			s.setState("connected")
			s.logSelectedPair()
		case webrtc.PeerConnectionStateFailed:
			s.fail(errors.New("peerconnection failed"))
		case webrtc.PeerConnectionStateClosed:
			s.setState("closed")
		}
	})
}

// --- 发送端 ---

func (s *session) bindSenderDC(dc *webrtc.DataChannel) {
	s.dc = dc
	dc.SetBufferedAmountLowThreshold(lowWatermark)
	dc.OnBufferedAmountLow(func() {
		select {
		case s.bufLow <- struct{}{}:
		default:
		}
	})
	dc.OnMessage(func(msg webrtc.DataChannelMessage) {
		if !msg.IsString {
			return
		}
		var a ackFrame
		if json.Unmarshal(msg.Data, &a) == nil && a.Type == "ack" {
			if a.Bytes > s.acked.Load() {
				s.acked.Store(a.Bytes)
				s.emitProgress()
			}
		}
	})
	dc.OnOpen(func() { go s.streamFile() })
}

// streamFile 把文件按 64KB 分片经数据通道发出，水位回压与 web 引擎一致：缓冲将越 HIGH 前
// 等抽干到 LOW。发完 done 帧后先抽干本地缓冲、再等接收端 ack 追平总量才宣告 completed
// （对齐接收端真实收程，避免发送端抢先完成）。
func (s *session) streamFile() {
	f, err := os.Open(s.filePath)
	if err != nil {
		s.fail(fmt.Errorf("open %s: %w", s.filePath, err))
		return
	}
	defer f.Close()

	meta, _ := json.Marshal(metaFrame{Type: "meta", Name: filepath.Base(s.filePath), Size: s.fileSize})
	if err := s.dc.SendText(string(meta)); err != nil {
		s.fail(fmt.Errorf("send meta: %w", err))
		return
	}

	buf := make([]byte, chunkSize)
	for {
		select {
		case <-s.done:
			return
		default:
		}
		n, rerr := f.Read(buf)
		if n > 0 {
			if s.dc.BufferedAmount()+uint64(n) > highWatermark {
				s.waitBufferLow(lowWatermark)
			}
			if err := s.dc.Send(buf[:n]); err != nil {
				s.fail(fmt.Errorf("send chunk: %w", err))
				return
			}
			s.bytesSent.Add(int64(n))
			s.emitProgress()
		}
		if rerr == io.EOF {
			break
		}
		if rerr != nil {
			s.fail(fmt.Errorf("read %s: %w", s.filePath, rerr))
			return
		}
	}

	doneB, _ := json.Marshal(doneFrame{Type: "done"})
	_ = s.dc.SendText(string(doneB))
	s.waitDrained()
	s.waitAck(s.fileSize)
	s.setState("completed")
}

// --- 接收端 ---

func (s *session) bindReceiverDC(dc *webrtc.DataChannel) {
	s.dc = dc
	dc.OnMessage(func(msg webrtc.DataChannelMessage) {
		if msg.IsString {
			s.handleControl(msg.Data)
			return
		}
		s.writeChunk(msg.Data)
	})
	go s.ackLoop()
}

func (s *session) handleControl(data []byte) {
	var probe struct {
		Type string `json:"type"`
	}
	if json.Unmarshal(data, &probe) != nil {
		return
	}
	switch probe.Type {
	case "meta":
		var m metaFrame
		if json.Unmarshal(data, &m) == nil {
			s.openOutput(m.Name)
		}
	case "done":
		s.finalize()
	}
}

func (s *session) openOutput(name string) {
	s.outName = name
	tmp := filepath.Join(s.eng.downloadDir(), "."+s.id+".part")
	f, err := os.Create(tmp)
	if err != nil {
		s.fail(fmt.Errorf("create %s: %w", tmp, err))
		return
	}
	s.out = f
}

// writeChunk 直接写盘——无 base64 桥、无 OPFS、无 WebView 内存。pion 串行投递同一通道的
// OnMessage，故写入天然有序。
func (s *session) writeChunk(b []byte) {
	if s.out != nil {
		if _, err := s.out.Write(b); err != nil {
			s.fail(fmt.Errorf("write chunk: %w", err))
			return
		}
	}
	s.received.Add(int64(len(b)))
	s.emitProgress()
}

// ackLoop 每 200ms 把最新已收字节回传发送端（节流，避免每片都回与下行争用通道）。
func (s *session) ackLoop() {
	t := time.NewTicker(ackIntervalMs * time.Millisecond)
	defer t.Stop()
	var lastSent int64 = -1
	for {
		select {
		case <-s.done:
			return
		case <-t.C:
			if r := s.received.Load(); r != lastSent {
				s.sendAck(r)
				lastSent = r
			}
		}
	}
}

func (s *session) sendAck(bytes int64) {
	if s.dc == nil || s.dc.ReadyState() != webrtc.DataChannelStateOpen {
		return
	}
	b, _ := json.Marshal(ackFrame{Type: "ack", Bytes: bytes})
	_ = s.dc.SendText(string(b))
}

// finalize 在收到 done 帧时调用：done 意味所有分片已到（ordered），received 已等于总量。
// 立即回最终 ack、关闭文件、改名落到不冲突路径，报告路径与 completed。
func (s *session) finalize() {
	s.sendAck(s.received.Load())
	if s.out != nil {
		s.out.Close()
		dir := s.eng.downloadDir()
		final := nonCollidingPath(dir, s.outName)
		tmp := filepath.Join(dir, "."+s.id+".part")
		if err := os.Rename(tmp, final); err != nil {
			s.fail(fmt.Errorf("rename to %s: %w", final, err))
			return
		}
		s.out = nil
		s.eng.ev.OnSaved(s.id, final)
	}
	s.setState("completed")
}

// --- 信令 ---

func (s *session) handleSignal(data []byte) error {
	var p signalPayload
	if err := json.Unmarshal(data, &p); err != nil {
		return fmt.Errorf("decode signal: %w", err)
	}
	s.eng.ev.OnLog(fmt.Sprintf("session %s signal in: %s", s.id, p.Type))
	switch p.Type {
	case "offer":
		if p.SDP == nil {
			return errors.New("offer missing sdp")
		}
		if err := s.pc.SetRemoteDescription(*p.SDP); err != nil {
			return fmt.Errorf("set remote (offer): %w", err)
		}
		ans, err := s.pc.CreateAnswer(nil)
		if err != nil {
			return fmt.Errorf("create answer: %w", err)
		}
		if err := s.pc.SetLocalDescription(ans); err != nil {
			return fmt.Errorf("set local (answer): %w", err)
		}
		s.emitSignal(s.peer, &ans, nil)
	case "answer":
		if p.SDP == nil {
			return errors.New("answer missing sdp")
		}
		if err := s.pc.SetRemoteDescription(*p.SDP); err != nil {
			return fmt.Errorf("set remote (answer): %w", err)
		}
	case "ice":
		if p.Candidate == nil {
			return nil
		}
		if err := s.pc.AddICECandidate(*p.Candidate); err != nil {
			return fmt.Errorf("add ice candidate: %w", err)
		}
	}
	return nil
}

// emitSignal 把一条出站信令交给宿主 POST。type 由载荷推导：有 candidate 即 ice，否则取
// sdp 自身的类型（offer/answer），与 p2p.ts 的 payload.type 语义一致。
func (s *session) emitSignal(toPeer string, sdp *webrtc.SessionDescription, cand *webrtc.ICECandidateInit) {
	p := signalPayload{SDP: sdp, Candidate: cand}
	switch {
	case cand != nil:
		p.Type = "ice"
	case sdp != nil:
		p.Type = sdp.Type.String()
	}
	b, err := json.Marshal(p)
	if err != nil {
		return
	}
	s.eng.ev.OnSignal(s.id, toPeer, string(b))
}

// --- 进度 / 状态 / 收尾 ---

func (s *session) emitProgress() {
	var b int64
	if s.role == roleSender {
		if a := s.acked.Load(); a > 0 {
			b = min(s.fileSize, a)
		} else {
			b = s.bytesSent.Load() - int64(s.dc.BufferedAmount())
			if b < 0 {
				b = 0
			}
		}
	} else {
		b = s.received.Load()
	}
	s.eng.ev.OnProgress(s.id, b)
}

// waitBufferLow 等本地 SCTP 缓冲落到 threshold 以下：靠 OnBufferedAmountLow 事件唤醒，
// 兼一个 50ms 轮询兜底（事件阈值设在 lowWatermark，其它 threshold 靠轮询）。
func (s *session) waitBufferLow(threshold uint64) {
	start := time.Now()
	warned := false
	for {
		if s.dc.BufferedAmount() <= threshold {
			return
		}
		select {
		case <-s.bufLow:
		case <-time.After(50 * time.Millisecond):
		case <-s.done:
			return
		}
		// 停滞看门狗：本地缓冲长时间不抽干＝SCTP 不再向网络交付，多为接收端 rwnd=0
		// （主线程/落盘卡住）或中继路径拥塞窗口塌缩。打一条诊断日志（不中断，polling 接力）。
		if !warned && time.Since(start) > 5*time.Second {
			warned = true
			s.eng.ev.OnLog(fmt.Sprintf(
				"session %s send stalled >5s: buffered=%d sent=%d acked=%d",
				s.id, s.dc.BufferedAmount(), s.bytesSent.Load(), s.acked.Load()))
		}
	}
}

// logSelectedPair 在连接建立后异步取 pion 选中的候选对类型（host/srflx/prflx/relay）并打日志，
// 用于确认是否走上同内网 host↔host 直连。nominated pair 在 connected 时常未定，轮询几次。
func (s *session) logSelectedPair() {
	go func() {
		for i := 0; i < 12; i += 1 {
			if local, remote, ok := s.selectedPair(); ok {
				s.eng.ev.OnLog(fmt.Sprintf("session %s ice pair: local=%s remote=%s", s.id, local, remote))
				s.eng.ev.OnIcePair(s.id, local, remote)
				return
			}
			select {
			case <-time.After(300 * time.Millisecond):
			case <-s.done:
				return
			}
		}
	}()
}

func (s *session) selectedPair() (local, remote string, ok bool) {
	defer func() { _ = recover() }() // 早期连接链路可能尚 nil，吞 panic
	sctp := s.pc.SCTP()
	if sctp == nil {
		return "", "", false
	}
	dtls := sctp.Transport()
	if dtls == nil {
		return "", "", false
	}
	it := dtls.ICETransport()
	if it == nil {
		return "", "", false
	}
	pair, err := it.GetSelectedCandidatePair()
	if err != nil || pair == nil || pair.Local == nil || pair.Remote == nil {
		return "", "", false
	}
	local = fmt.Sprintf("%s/%s", pair.Local.Typ, pair.Local.Protocol)
	remote = fmt.Sprintf("%s/%s", pair.Remote.Typ, pair.Remote.Protocol)
	return local, remote, true
}

// waitDrained 等本地缓冲彻底抽干（done 帧也已离开本端 SCTP）。
func (s *session) waitDrained() {
	for {
		if s.dc.BufferedAmount() == 0 {
			return
		}
		select {
		case <-time.After(20 * time.Millisecond):
		case <-s.done:
			return
		}
	}
}

// waitAck 等接收端 ack 追平 total（确认已全收）。超时兜底；无任何 ack（不应发生，因对端
// 皆回 ack）则不空等。
func (s *session) waitAck(total int64) {
	if s.acked.Load() == 0 {
		return
	}
	deadline := time.Now().Add(ackCompleteTimeoutSec * time.Second)
	for s.acked.Load() < total {
		if time.Now().After(deadline) {
			return
		}
		select {
		case <-time.After(50 * time.Millisecond):
		case <-s.done:
			return
		}
	}
}

func (s *session) setState(st string) {
	s.stateMu.Lock()
	// 终态不可被覆盖：completed/failed 之后即便连接随后 closed/failed 也不回退——否则接收端
	// 完成后、发送端关连接时本端会把已完成会话误翻成 closed/failed。
	if s.state == st || s.state == "completed" || s.state == "failed" {
		s.stateMu.Unlock()
		return
	}
	s.state = st
	s.stateMu.Unlock()

	s.eng.ev.OnState(s.id, st)
	switch st {
	case "failed", "closed":
		s.cleanup()
	case "completed":
		// 不立即拆连接。接收端在 finalize 里刚发出最终 ack；若此刻就关通道，发送端（尤其在
		// 前向数据把链路打满、反向 ack 被饿的高吞吐场景）可能收不到最终 ack → 误判失败（实测
		// iOS→mac 实际成功却被 iOS 标失败、进度停在 210/271）。改为留连接：ackLoop 每 200ms 续发
		// 最终 ack，待前向洪流结束、反向腾出后送达发送端。真正拆连接交给 SSE 终态驱动的 Cancel
		// （hub→p2pCleanup→P2PCancel）；60s 兜底防泄漏。
		go func() {
			select {
			case <-time.After(60 * time.Second):
				s.cleanup()
			case <-s.done:
			}
		}()
	}
}

func (s *session) fail(err error) {
	s.eng.ev.OnLog(fmt.Sprintf("session %s fail: %v", s.id, err))
	s.setState("failed")
}

func (s *session) cleanup() {
	s.closeOnce.Do(func() {
		close(s.done)
		if s.out != nil {
			s.out.Close()
			// 半截 .part 留作诊断证据由上层决定清理；此处仅关句柄。
		}
		if s.dc != nil {
			_ = s.dc.Close()
		}
		if s.pc != nil {
			_ = s.pc.Close()
		}
		s.eng.remove(s.peer, s)
	})
}