427 lines
11 KiB
Go
427 lines
11 KiB
Go
/**
|
||
* package mediaserver
|
||
*
|
||
* WebRTC SFU (Selective Forwarding Unit) 服务
|
||
* 功能：
|
||
* 1. 接收 WebRTC 流
|
||
* 2. 转发给同房间的其他用户
|
||
* 3. 支持音视频分离转发
|
||
*/
|
||
package mediaserver
|
||
|
||
import (
|
||
"crypto/hmac"
|
||
"crypto/sha1"
|
||
"encoding/base64"
|
||
"encoding/json"
|
||
"fmt"
|
||
"log"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/pion/interceptor"
|
||
"github.com/pion/interceptor/pkg/intervalpli"
|
||
"github.com/pion/webrtc/v3"
|
||
"github.com/spf13/viper"
|
||
)
|
||
|
||
// SFUServer WebRTC SFU 服务器
|
||
type SFUServer struct {
|
||
config *MediaServerConfig
|
||
api *webrtc.API
|
||
mu sync.RWMutex
|
||
running bool
|
||
|
||
// Track 管理
|
||
trackLocals map[string]*webrtc.TrackLocalStaticRTP // trackID -> localTrack
|
||
}
|
||
|
||
// NewSFUServer 创建 SFU 服务器
|
||
func NewSFUServer(config *MediaServerConfig) *SFUServer {
|
||
return &SFUServer{
|
||
config: config,
|
||
trackLocals: make(map[string]*webrtc.TrackLocalStaticRTP),
|
||
}
|
||
}
|
||
|
||
// Start 启动 SFU 服务器
|
||
func (s *SFUServer) Start() {
|
||
s.mu.Lock()
|
||
if s.running {
|
||
s.mu.Unlock()
|
||
return
|
||
}
|
||
s.running = true
|
||
s.mu.Unlock()
|
||
|
||
// 初始化失败时必须复位 running：
|
||
// 原实现失败直接 return，running 仍为 true 但 api 为 nil，
|
||
// SFU 处于"假启动"状态且因为 running=true 无法再次 Start，只能重启进程
|
||
fail := func(format string, args ...interface{}) {
|
||
log.Printf(format, args...)
|
||
s.mu.Lock()
|
||
s.running = false
|
||
s.mu.Unlock()
|
||
}
|
||
|
||
// 创建 MediaEngine
|
||
m := &webrtc.MediaEngine{}
|
||
|
||
// 注册默认编解码器
|
||
if err := m.RegisterDefaultCodecs(); err != nil {
|
||
fail("❌ [SFU] 注册编解码器失败: %v", err)
|
||
return
|
||
}
|
||
|
||
// 创建拦截器注册表
|
||
i := &interceptor.Registry{}
|
||
|
||
// 注册 PLI 拦截器（用于请求关键帧）
|
||
intervalPliFactory, err := intervalpli.NewReceiverInterceptor()
|
||
if err != nil {
|
||
fail("❌ [SFU] 创建 PLI 拦截器失败: %v", err)
|
||
return
|
||
}
|
||
i.Add(intervalPliFactory)
|
||
|
||
// 使用拦截器
|
||
if err := webrtc.RegisterDefaultInterceptors(m, i); err != nil {
|
||
fail("❌ [SFU] 注册拦截器失败: %v", err)
|
||
return
|
||
}
|
||
|
||
// 创建 API
|
||
s.api = webrtc.NewAPI(webrtc.WithMediaEngine(m), webrtc.WithInterceptorRegistry(i))
|
||
|
||
log.Printf("🚀 [SFU] WebRTC SFU 已启动 | 端口: %d", s.config.WebRTCPort)
|
||
}
|
||
|
||
// Stop 停止 SFU 服务器
|
||
func (s *SFUServer) Stop() {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
|
||
s.running = false
|
||
s.trackLocals = make(map[string]*webrtc.TrackLocalStaticRTP)
|
||
log.Println("🛑 [SFU] 已停止")
|
||
}
|
||
|
||
// CreatePeerConnection 为用户创建 PeerConnection
|
||
func (s *SFUServer) CreatePeerConnection(roomID, userID string) (*webrtc.PeerConnection, error) {
|
||
s.mu.RLock()
|
||
if !s.running || s.api == nil {
|
||
s.mu.RUnlock()
|
||
return nil, ErrSFUNotReady
|
||
}
|
||
s.mu.RUnlock()
|
||
|
||
// 获取 ICE 服务器配置
|
||
iceServers := s.getICEServers(userID)
|
||
|
||
// 创建 PeerConnection 配置
|
||
config := webrtc.Configuration{
|
||
ICEServers: iceServers,
|
||
}
|
||
|
||
// 创建 PeerConnection
|
||
pc, err := s.api.NewPeerConnection(config)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("create peer connection failed: %w", err)
|
||
}
|
||
|
||
// 监听 ICE 连接状态
|
||
pc.OnICEConnectionStateChange(func(state webrtc.ICEConnectionState) {
|
||
log.Printf("🔗 [SFU] Room:%s User:%s ICE状态: %s", roomID, userID, state.String())
|
||
|
||
// 仅在 Failed / Closed 时清理资源。
|
||
// Disconnected 是可自愈的瞬时状态（网络抖动、WiFi切换等，ICE 会自动恢复），
|
||
// 原实现把 Disconnected 也当成断线立即踢出房间，导致弱网用户被误踢
|
||
if state == webrtc.ICEConnectionStateFailed || state == webrtc.ICEConnectionStateClosed {
|
||
s.handleDisconnect(roomID, userID)
|
||
}
|
||
})
|
||
|
||
// 监听轨道
|
||
pc.OnTrack(func(remoteTrack *webrtc.TrackRemote, receiver *webrtc.RTPReceiver) {
|
||
s.handleTrack(roomID, userID, remoteTrack, receiver)
|
||
})
|
||
|
||
return pc, nil
|
||
}
|
||
|
||
// getICEServers 获取 ICE 服务器配置
|
||
func (s *SFUServer) getICEServers(userID string) []webrtc.ICEServer {
|
||
servers := []webrtc.ICEServer{}
|
||
|
||
// 添加 STUN 服务器
|
||
turnPublicIP := viper.GetString("turn.public_ip")
|
||
turnPort := viper.GetInt("turn.listen_port")
|
||
if turnPublicIP != "" && turnPort > 0 {
|
||
stunURL := fmt.Sprintf("stun:%s:%d", turnPublicIP, turnPort)
|
||
servers = append(servers, webrtc.ICEServer{
|
||
URLs: []string{stunURL},
|
||
})
|
||
|
||
// 添加 TURN 服务器（如果启用）
|
||
if viper.GetBool("turn.enabled") {
|
||
turnURL := fmt.Sprintf("turn:%s:%d", turnPublicIP, turnPort)
|
||
// 生成 TURN 凭证
|
||
username, credential := generateTURNCredentials(userID)
|
||
servers = append(servers, webrtc.ICEServer{
|
||
URLs: []string{turnURL},
|
||
Username: username,
|
||
Credential: credential,
|
||
})
|
||
}
|
||
}
|
||
|
||
return servers
|
||
}
|
||
|
||
// handleTrack 处理远程轨道
|
||
func (s *SFUServer) handleTrack(roomID, userID string, remoteTrack *webrtc.TrackRemote, receiver *webrtc.RTPReceiver) {
|
||
trackID := fmt.Sprintf("%s_%s_%s", roomID, userID, remoteTrack.Kind().String())
|
||
|
||
log.Printf("🎬 [SFU] 收到轨道 | Room:%s User:%s Kind:%s ID:%s",
|
||
roomID, userID, remoteTrack.Kind().String(), trackID)
|
||
|
||
// 创建本地轨道用于转发
|
||
localTrack, err := webrtc.NewTrackLocalStaticRTP(
|
||
remoteTrack.Codec().RTPCodecCapability,
|
||
trackID,
|
||
fmt.Sprintf("stream_%s", userID),
|
||
)
|
||
if err != nil {
|
||
log.Printf("❌ [SFU] 创建本地轨道失败: %v", err)
|
||
return
|
||
}
|
||
|
||
// 保存轨道
|
||
s.mu.Lock()
|
||
s.trackLocals[trackID] = localTrack
|
||
s.mu.Unlock()
|
||
|
||
// 转发 RTP 包
|
||
go func() {
|
||
buf := make([]byte, 1500)
|
||
for {
|
||
n, _, readErr := remoteTrack.Read(buf)
|
||
if readErr != nil {
|
||
log.Printf("⚠️ [SFU] 读取轨道失败 | Track:%s Error:%v", trackID, readErr)
|
||
break
|
||
}
|
||
|
||
// 写入本地轨道（会自动转发给所有订阅者）
|
||
if _, writeErr := localTrack.Write(buf[:n]); writeErr != nil {
|
||
log.Printf("⚠️ [SFU] 写入轨道失败 | Track:%s Error:%v", trackID, writeErr)
|
||
break
|
||
}
|
||
}
|
||
|
||
// 清理轨道
|
||
s.mu.Lock()
|
||
delete(s.trackLocals, trackID)
|
||
s.mu.Unlock()
|
||
}()
|
||
|
||
// 通知房间内其他用户有新轨道
|
||
s.notifyNewTrack(roomID, userID, localTrack)
|
||
}
|
||
|
||
// notifyNewTrack 通知房间内其他用户有新轨道
|
||
func (s *SFUServer) notifyNewTrack(roomID, senderUserID string, track *webrtc.TrackLocalStaticRTP) {
|
||
ms := GetServer()
|
||
room := ms.GetRoom(roomID)
|
||
if room == nil {
|
||
return
|
||
}
|
||
|
||
// 获取其他参与者
|
||
for _, p := range room.GetOtherParticipants(senderUserID) {
|
||
if p.Type != ParticipantTypeWebRTC {
|
||
continue // 只处理 WebRTC 用户
|
||
}
|
||
|
||
if p.PeerConnection == nil {
|
||
continue
|
||
}
|
||
|
||
pc, ok := p.PeerConnection.(*webrtc.PeerConnection)
|
||
if !ok {
|
||
continue
|
||
}
|
||
|
||
// 添加轨道到对方的 PeerConnection
|
||
if _, err := pc.AddTrack(track); err != nil {
|
||
log.Printf("⚠️ [SFU] 添加轨道到用户 %s 失败: %v", p.UserID, err)
|
||
}
|
||
}
|
||
}
|
||
|
||
// handleDisconnect 处理断开连接
|
||
func (s *SFUServer) handleDisconnect(roomID, userID string) {
|
||
ms := GetServer()
|
||
room := ms.GetRoom(roomID)
|
||
if room == nil {
|
||
return
|
||
}
|
||
|
||
room.RemoveParticipant(userID)
|
||
|
||
// 清理该用户的所有轨道
|
||
s.mu.Lock()
|
||
for trackID := range s.trackLocals {
|
||
// 检查 trackID 是否属于该用户
|
||
prefix := fmt.Sprintf("%s_%s_", roomID, userID)
|
||
if strings.HasPrefix(trackID, prefix) {
|
||
delete(s.trackLocals, trackID)
|
||
}
|
||
}
|
||
s.mu.Unlock()
|
||
|
||
// 如果房间空了，移除房间
|
||
if room.IsEmpty() {
|
||
ms.RemoveRoom(roomID)
|
||
}
|
||
}
|
||
|
||
// GetTrackCount 获取轨道数量
|
||
func (s *SFUServer) GetTrackCount() int {
|
||
s.mu.RLock()
|
||
defer s.mu.RUnlock()
|
||
return len(s.trackLocals)
|
||
}
|
||
|
||
// HandleOffer 处理 SDP Offer
|
||
func (s *SFUServer) HandleOffer(roomID, userID string, offerSDP string) (string, error) {
|
||
ms := GetServer()
|
||
room := ms.GetOrCreateRoom(roomID)
|
||
|
||
// 添加参与者（重复加入时 AddParticipant 内部会先 Close 掉旧的 PeerConnection，防止泄漏）
|
||
participant, err := room.AddParticipant(userID, ParticipantTypeWebRTC)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
|
||
// 创建 PeerConnection
|
||
pc, err := s.CreatePeerConnection(roomID, userID)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
participant.PeerConnection = pc
|
||
|
||
// 协商中途失败时关闭刚建的 PeerConnection，避免半初始化的连接泄漏
|
||
cleanup := func() {
|
||
_ = pc.Close()
|
||
participant.PeerConnection = nil
|
||
}
|
||
|
||
// 设置远程描述
|
||
offer := webrtc.SessionDescription{
|
||
Type: webrtc.SDPTypeOffer,
|
||
SDP: offerSDP,
|
||
}
|
||
if err := pc.SetRemoteDescription(offer); err != nil {
|
||
cleanup()
|
||
return "", fmt.Errorf("set remote description failed: %w", err)
|
||
}
|
||
|
||
// 添加房间内其他用户的轨道
|
||
s.addExistingTracks(pc, roomID, userID)
|
||
|
||
// 创建 Answer
|
||
answer, err := pc.CreateAnswer(nil)
|
||
if err != nil {
|
||
cleanup()
|
||
return "", fmt.Errorf("create answer failed: %w", err)
|
||
}
|
||
|
||
// 在 SetLocalDescription 之前创建"收集完成"信号（SetLocalDescription 才会触发 ICE 收集）
|
||
gatherComplete := webrtc.GatheringCompletePromise(pc)
|
||
|
||
// 设置本地描述
|
||
if err := pc.SetLocalDescription(answer); err != nil {
|
||
cleanup()
|
||
return "", fmt.Errorf("set local description failed: %w", err)
|
||
}
|
||
|
||
// 等待 ICE 收集完成，最多等 10 秒。
|
||
// 原实现裸 <-channel 无超时：网络异常时收集可能永远不结束，
|
||
// 处理该请求的 HTTP worker 会被永久阻塞，积累多了服务就没有可用协程了。
|
||
// 超时后用当前已收集到的候选返回，客户端仍可通过 trickle ICE 补充。
|
||
select {
|
||
case <-gatherComplete:
|
||
case <-time.After(10 * time.Second):
|
||
log.Printf("⚠️ [SFU] ICE 收集超时(10s)，返回已收集候选 | Room:%s User:%s", roomID, userID)
|
||
}
|
||
|
||
return pc.LocalDescription().SDP, nil
|
||
}
|
||
|
||
// addExistingTracks 添加房间内已有的轨道
|
||
func (s *SFUServer) addExistingTracks(pc *webrtc.PeerConnection, roomID, excludeUserID string) {
|
||
s.mu.RLock()
|
||
defer s.mu.RUnlock()
|
||
|
||
prefix := fmt.Sprintf("%s_", roomID)
|
||
excludePrefix := fmt.Sprintf("%s_%s_", roomID, excludeUserID)
|
||
|
||
for trackID, track := range s.trackLocals {
|
||
// 检查是否是同房间的轨道，且不是自己的
|
||
if strings.HasPrefix(trackID, prefix) && !strings.HasPrefix(trackID, excludePrefix) {
|
||
if _, err := pc.AddTrack(track); err != nil {
|
||
log.Printf("⚠️ [SFU] 添加已有轨道失败: %v", err)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// HandleICECandidate 处理 ICE 候选
|
||
func (s *SFUServer) HandleICECandidate(roomID, userID string, candidateJSON string) error {
|
||
ms := GetServer()
|
||
room := ms.GetRoom(roomID)
|
||
if room == nil {
|
||
return ErrRoomNotFound
|
||
}
|
||
|
||
participant := room.GetParticipant(userID)
|
||
if participant == nil {
|
||
return ErrParticipantNotFound
|
||
}
|
||
|
||
pc, ok := participant.PeerConnection.(*webrtc.PeerConnection)
|
||
if !ok || pc == nil {
|
||
return ErrSFUNotReady
|
||
}
|
||
|
||
var candidate webrtc.ICECandidateInit
|
||
if err := json.Unmarshal([]byte(candidateJSON), &candidate); err != nil {
|
||
return fmt.Errorf("parse ICE candidate failed: %w", err)
|
||
}
|
||
|
||
if err := pc.AddICECandidate(candidate); err != nil {
|
||
return fmt.Errorf("add ICE candidate failed: %w", err)
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// generateTURNCredentials 生成 TURN 凭证
|
||
func generateTURNCredentials(userID string) (string, string) {
|
||
timestamp := time.Now().Add(24 * time.Hour).Unix()
|
||
username := fmt.Sprintf("%d:%s", timestamp, userID)
|
||
secret := viper.GetString("turn.shared_secret")
|
||
|
||
mac := hmac.New(sha1.New, []byte(secret))
|
||
mac.Write([]byte(username))
|
||
password := base64.StdEncoding.EncodeToString(mac.Sum(nil))
|
||
|
||
return username, password
|
||
}
|
||
|
||
// GetTURNCredentials 获取 TURN 凭证（供外部调用）
|
||
func GetTURNCredentials(userID string) (string, string) {
|
||
return generateTURNCredentials(userID)
|
||
}
|