975 lines
26 KiB
Go
975 lines
26 KiB
Go
/**
|
||
* package mediaserver
|
||
*
|
||
* WebSocket to RTMP 代理服务
|
||
* 接收 Web端 通过 WebSocket 发送的 WebM 数据，转换为 FLV 并推送到 RTMP
|
||
*
|
||
* 技术栈：
|
||
* - WebSocket: gorilla/websocket
|
||
* - WebM 解析: github.com/at-wat/ebml-go
|
||
* - FLV 封装: github.com/yutopp/go-flv
|
||
* - RTMP 推送: 通过内部流管道
|
||
*/
|
||
package mediaserver
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/binary"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/at-wat/ebml-go"
|
||
"github.com/gorilla/websocket"
|
||
)
|
||
|
||
// WebMToRTMPProxy WebSocket to RTMP 代理
|
||
type WebMToRTMPProxy struct {
|
||
rtmpServer *RTMPServer
|
||
upgrader websocket.Upgrader
|
||
sessions map[string]*ProxySession
|
||
sessionsMu sync.RWMutex
|
||
}
|
||
|
||
// CodecInfo 编解码器信息（从前端发送的 JSON 消息解析）
|
||
type CodecInfo struct {
|
||
Type string `json:"type"` // "codec_info"
|
||
VideoCodec string `json:"video_codec"` // "h264", "vp8", "vp9"
|
||
AudioCodec string `json:"audio_codec"` // "opus", "aac"
|
||
MimeType string `json:"mime_type"` // 完整的 MIME 类型
|
||
}
|
||
|
||
// ProxySession 代理会话
|
||
type ProxySession struct {
|
||
ID string
|
||
StreamID string
|
||
UserID string
|
||
RoomID string
|
||
Conn *websocket.Conn
|
||
Stream *RTMPStream
|
||
StopChan chan struct{}
|
||
StartTime time.Time
|
||
|
||
// WebM 解析状态
|
||
webmBuffer *bytes.Buffer
|
||
headerParsed bool
|
||
videoTrackNum uint64
|
||
audioTrackNum uint64
|
||
|
||
// 时间戳
|
||
baseTimestamp uint32
|
||
lastTimestamp uint32
|
||
|
||
// 编解码器信息
|
||
codecInfo *CodecInfo
|
||
codecDetected bool
|
||
|
||
// FFmpeg 转码器（VP8/VP9 需要转码）
|
||
transcoder *FFmpegTranscoder
|
||
needTranscode bool
|
||
}
|
||
|
||
// WebMHeader WebM 文件头信息
|
||
type WebMHeader struct {
|
||
EBMLVersion uint64 `ebml:"EBMLVersion"`
|
||
EBMLReadVersion uint64 `ebml:"EBMLReadVersion"`
|
||
EBMLMaxIDLength uint64 `ebml:"EBMLMaxIDLength"`
|
||
EBMLMaxSizeLength uint64 `ebml:"EBMLMaxSizeLength"`
|
||
DocType string `ebml:"DocType"`
|
||
DocTypeVersion uint64 `ebml:"DocTypeVersion"`
|
||
DocTypeReadVersion uint64 `ebml:"DocTypeReadVersion"`
|
||
}
|
||
|
||
// WebMSegment WebM Segment
|
||
type WebMSegment struct {
|
||
Info WebMSegmentInfo `ebml:"Info"`
|
||
Tracks WebMTracks `ebml:"Tracks"`
|
||
Cluster []WebMCluster `ebml:"Cluster"`
|
||
}
|
||
|
||
// WebMSegmentInfo Segment 信息
|
||
type WebMSegmentInfo struct {
|
||
TimecodeScale uint64 `ebml:"TimecodeScale"`
|
||
Duration float64 `ebml:"Duration,omitempty"`
|
||
MuxingApp string `ebml:"MuxingApp,omitempty"`
|
||
WritingApp string `ebml:"WritingApp,omitempty"`
|
||
}
|
||
|
||
// WebMTracks 轨道信息
|
||
type WebMTracks struct {
|
||
TrackEntry []WebMTrackEntry `ebml:"TrackEntry"`
|
||
}
|
||
|
||
// WebMTrackEntry 轨道条目
|
||
type WebMTrackEntry struct {
|
||
TrackNumber uint64 `ebml:"TrackNumber"`
|
||
TrackType uint64 `ebml:"TrackType"` // 1=video, 2=audio
|
||
CodecID string `ebml:"CodecID"`
|
||
Video *WebMVideoTrack `ebml:"Video,omitempty"`
|
||
Audio *WebMAudioTrack `ebml:"Audio,omitempty"`
|
||
}
|
||
|
||
// WebMVideoTrack 视频轨道
|
||
type WebMVideoTrack struct {
|
||
PixelWidth uint64 `ebml:"PixelWidth"`
|
||
PixelHeight uint64 `ebml:"PixelHeight"`
|
||
}
|
||
|
||
// WebMAudioTrack 音频轨道
|
||
type WebMAudioTrack struct {
|
||
SamplingFrequency float64 `ebml:"SamplingFrequency"`
|
||
Channels uint64 `ebml:"Channels"`
|
||
BitDepth uint64 `ebml:"BitDepth,omitempty"`
|
||
}
|
||
|
||
// WebMCluster WebM Cluster
|
||
type WebMCluster struct {
|
||
Timecode uint64 `ebml:"Timecode"`
|
||
SimpleBlock []ebml.Block `ebml:"SimpleBlock,omitempty"`
|
||
}
|
||
|
||
// NewWebMToRTMPProxy 创建代理服务
|
||
func NewWebMToRTMPProxy(rtmpServer *RTMPServer) *WebMToRTMPProxy {
|
||
return &WebMToRTMPProxy{
|
||
rtmpServer: rtmpServer,
|
||
upgrader: websocket.Upgrader{
|
||
CheckOrigin: func(r *http.Request) bool {
|
||
return true // 允许跨域
|
||
},
|
||
ReadBufferSize: 1024 * 1024,
|
||
WriteBufferSize: 1024 * 1024,
|
||
},
|
||
sessions: make(map[string]*ProxySession),
|
||
}
|
||
}
|
||
|
||
// HandleWebSocket 处理 WebSocket 连接
|
||
func (p *WebMToRTMPProxy) HandleWebSocket(w http.ResponseWriter, r *http.Request) {
|
||
log.Printf("🔌 [WSProxy] 收到 WebSocket 连接请求: %s from %s", r.URL.String(), r.RemoteAddr)
|
||
|
||
// 检查 WebSocket 升级头
|
||
if r.Header.Get("Upgrade") != "websocket" {
|
||
log.Printf("❌ [WSProxy] 不是 WebSocket 请求: Upgrade=%s", r.Header.Get("Upgrade"))
|
||
http.Error(w, "Not a WebSocket request", http.StatusBadRequest)
|
||
return
|
||
}
|
||
|
||
// 获取参数
|
||
streamID := r.URL.Query().Get("stream_id")
|
||
userID := r.URL.Query().Get("user_id")
|
||
roomID := r.URL.Query().Get("room_id")
|
||
token := r.URL.Query().Get("token")
|
||
|
||
log.Printf("📋 [WSProxy] 参数: stream_id=%s user_id=%s room_id=%s token=%v",
|
||
streamID, userID, roomID, token != "")
|
||
|
||
if streamID == "" || userID == "" {
|
||
log.Printf("❌ [WSProxy] 缺少必要参数")
|
||
http.Error(w, "Missing stream_id or user_id", http.StatusBadRequest)
|
||
return
|
||
}
|
||
|
||
// 验证流 Token：签名+过期时间校验，防止任意人拿到 stream_id 就能推流
|
||
if !ValidateStreamToken(streamID, token) {
|
||
log.Printf("❌ [WSProxy] token 校验失败: stream=%s", streamID)
|
||
http.Error(w, "Invalid stream token", http.StatusUnauthorized)
|
||
return
|
||
}
|
||
|
||
// 升级到 WebSocket
|
||
log.Printf("🔄 [WSProxy] 尝试 WebSocket 升级...")
|
||
conn, err := p.upgrader.Upgrade(w, r, nil)
|
||
if err != nil {
|
||
log.Printf("❌ [WSProxy] WebSocket 升级失败: %v (可能原因: 中间件干扰、响应已写入)", err)
|
||
return
|
||
}
|
||
|
||
log.Printf("✅ [WSProxy] WebSocket 升级成功 | Stream:%s User:%s Room:%s", streamID, userID, roomID)
|
||
|
||
// 获取或创建 RTMP 流
|
||
stream := p.rtmpServer.GetStream(streamID)
|
||
if stream == nil {
|
||
// 自动创建流
|
||
var err error
|
||
stream, err = p.rtmpServer.GenerateStreamURLs(roomID, userID)
|
||
if err != nil {
|
||
log.Printf("❌ [WSProxy] 创建流失败: %v", err)
|
||
conn.Close()
|
||
return
|
||
}
|
||
// 更新 streamID 为实际生成的
|
||
streamID = stream.ID
|
||
}
|
||
|
||
// 创建会话
|
||
session := &ProxySession{
|
||
ID: fmt.Sprintf("%s_%d", userID, time.Now().UnixNano()),
|
||
StreamID: streamID,
|
||
UserID: userID,
|
||
RoomID: roomID,
|
||
Conn: conn,
|
||
Stream: stream,
|
||
StopChan: make(chan struct{}),
|
||
StartTime: time.Now(),
|
||
webmBuffer: bytes.NewBuffer(nil),
|
||
headerParsed: false,
|
||
}
|
||
|
||
// 注册会话
|
||
p.sessionsMu.Lock()
|
||
p.sessions[session.ID] = session
|
||
p.sessionsMu.Unlock()
|
||
|
||
// 激活流
|
||
stream.mu.Lock()
|
||
stream.IsActive = true
|
||
stream.mu.Unlock()
|
||
|
||
// 处理连接
|
||
go p.handleSession(session)
|
||
}
|
||
|
||
// handleSession 处理会话
|
||
func (p *WebMToRTMPProxy) handleSession(session *ProxySession) {
|
||
log.Printf("▶️ [WSProxy] 开始处理会话 | Stream:%s User:%s", session.StreamID, session.UserID)
|
||
|
||
messageCount := 0
|
||
totalBytes := int64(0)
|
||
|
||
defer func() {
|
||
// 清理转码器
|
||
if session.transcoder != nil {
|
||
session.transcoder.Stop()
|
||
}
|
||
|
||
// 清理
|
||
p.sessionsMu.Lock()
|
||
delete(p.sessions, session.ID)
|
||
p.sessionsMu.Unlock()
|
||
|
||
session.Conn.Close()
|
||
close(session.StopChan)
|
||
|
||
// 标记流为非活动
|
||
if session.Stream != nil {
|
||
session.Stream.mu.Lock()
|
||
session.Stream.IsActive = false
|
||
session.Stream.mu.Unlock()
|
||
}
|
||
|
||
codecStr := "unknown"
|
||
if session.codecInfo != nil {
|
||
codecStr = session.codecInfo.VideoCodec
|
||
}
|
||
log.Printf("🔌 [WSProxy] 连接关闭 | Stream:%s User:%s | 收到消息:%d 总字节:%d 编解码器:%s 需要转码:%v",
|
||
session.StreamID, session.UserID, messageCount, totalBytes, codecStr, session.needTranscode)
|
||
}()
|
||
|
||
for {
|
||
select {
|
||
case <-session.StopChan:
|
||
log.Printf("⏹️ [WSProxy] 收到停止信号 | Stream:%s", session.StreamID)
|
||
return
|
||
default:
|
||
}
|
||
|
||
// 读取 WebSocket 消息
|
||
messageType, data, err := session.Conn.ReadMessage()
|
||
if err != nil {
|
||
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
||
log.Printf("⚠️ [WSProxy] 读取错误: %v | Stream:%s", err, session.StreamID)
|
||
} else {
|
||
log.Printf("ℹ️ [WSProxy] 连接关闭: %v | Stream:%s", err, session.StreamID)
|
||
}
|
||
return
|
||
}
|
||
|
||
messageCount++
|
||
totalBytes += int64(len(data))
|
||
|
||
// 记录首次收到消息
|
||
if messageCount == 1 {
|
||
log.Printf("📥 [WSProxy] 首次收到消息 | Stream:%s Type:%d Size:%d", session.StreamID, messageType, len(data))
|
||
}
|
||
|
||
// 每 100 条消息记录一次统计
|
||
if messageCount%100 == 0 {
|
||
log.Printf("📊 [WSProxy] 消息统计 | Stream:%s Count:%d TotalBytes:%d NeedTranscode:%v",
|
||
session.StreamID, messageCount, totalBytes, session.needTranscode)
|
||
}
|
||
|
||
// 处理文本消息（可能是编解码器信息）
|
||
if messageType == websocket.TextMessage {
|
||
if err := p.handleTextMessage(session, data); err != nil {
|
||
log.Printf("⚠️ [WSProxy] 处理文本消息失败: %v", err)
|
||
}
|
||
continue
|
||
}
|
||
|
||
if messageType != websocket.BinaryMessage {
|
||
log.Printf("⚠️ [WSProxy] 忽略未知消息类型 | Type:%d", messageType)
|
||
continue
|
||
}
|
||
|
||
// 处理二进制数据（WebM）
|
||
if session.needTranscode && session.transcoder != nil {
|
||
// VP8/VP9: 通过 FFmpeg 转码
|
||
if err := p.processWithTranscoder(session, data); err != nil {
|
||
log.Printf("⚠️ [WSProxy] 转码处理失败: %v", err)
|
||
}
|
||
} else {
|
||
// H.264 或未检测到编解码器: 直接处理
|
||
if err := p.processWebMData(session, data); err != nil {
|
||
log.Printf("⚠️ [WSProxy] 处理 WebM 数据失败: %v", err)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// handleTextMessage 处理文本消息（编解码器信息）
|
||
func (p *WebMToRTMPProxy) handleTextMessage(session *ProxySession, data []byte) error {
|
||
var codecInfo CodecInfo
|
||
if err := json.Unmarshal(data, &codecInfo); err != nil {
|
||
return fmt.Errorf("JSON 解析失败: %w", err)
|
||
}
|
||
|
||
if codecInfo.Type != "codec_info" {
|
||
log.Printf("ℹ️ [WSProxy] 收到非 codec_info 消息: %s", codecInfo.Type)
|
||
return nil
|
||
}
|
||
|
||
session.codecInfo = &codecInfo
|
||
session.codecDetected = true
|
||
|
||
log.Printf("🎬 [WSProxy] 收到编解码器信息 | Stream:%s Video:%s Audio:%s MIME:%s",
|
||
session.StreamID, codecInfo.VideoCodec, codecInfo.AudioCodec, codecInfo.MimeType)
|
||
|
||
// 判断编解码器类型并初始化对应的转码器
|
||
// 所有格式都通过 FFmpeg 处理，但 H.264 使用 copy 模式（只重封装），VP8/VP9 需要转码
|
||
switch codecInfo.VideoCodec {
|
||
case "h264", "avc1":
|
||
// H.264: 使用 FFmpeg copy 模式，只重封装为 FLV，不重新编码
|
||
session.needTranscode = true
|
||
log.Printf("✅ [WSProxy] H.264 编码，使用 FFmpeg copy 模式重封装为 FLV")
|
||
|
||
if err := p.initTranscoderForH264(session); err != nil {
|
||
log.Printf("⚠️ [WSProxy] 初始化 H.264 转码器失败: %v，尝试直接解析", err)
|
||
session.needTranscode = false
|
||
}
|
||
case "vp8", "vp9":
|
||
// VP8/VP9: 需要 FFmpeg 转码为 H.264
|
||
session.needTranscode = true
|
||
log.Printf("🔄 [WSProxy] %s 编码，需要 FFmpeg 转码为 H.264", codecInfo.VideoCodec)
|
||
|
||
if err := p.initTranscoder(session); err != nil {
|
||
log.Printf("⚠️ [WSProxy] 初始化转码器失败: %v，将使用非标准 FLV 格式", err)
|
||
session.needTranscode = false
|
||
}
|
||
default:
|
||
log.Printf("⚠️ [WSProxy] 未知视频编解码器: %s，尝试直接处理", codecInfo.VideoCodec)
|
||
session.needTranscode = false
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// initTranscoder 初始化 FFmpeg 转码器（VP8/VP9 转码模式）
|
||
func (p *WebMToRTMPProxy) initTranscoder(session *ProxySession) error {
|
||
transcoder, err := NewFFmpegTranscoder(nil)
|
||
if err != nil {
|
||
return fmt.Errorf("创建转码器失败: %w", err)
|
||
}
|
||
|
||
if err := transcoder.Start(); err != nil {
|
||
return fmt.Errorf("启动转码器失败: %w", err)
|
||
}
|
||
|
||
session.transcoder = transcoder
|
||
|
||
// 启动协程读取转码后的 FLV 数据并广播
|
||
go p.readTranscodedOutput(session)
|
||
|
||
log.Printf("✅ [WSProxy] FFmpeg 转码器已初始化（转码模式）| Stream:%s", session.StreamID)
|
||
return nil
|
||
}
|
||
|
||
// initTranscoderForH264 初始化 FFmpeg 转码器（H.264 copy 模式）
|
||
// H.264 视频只需要重封装，不需要重新编码，使用 copy 模式可以大大减少延迟
|
||
func (p *WebMToRTMPProxy) initTranscoderForH264(session *ProxySession) error {
|
||
config := &TranscoderConfig{
|
||
CopyMode: true, // 使用 copy 模式
|
||
AudioCodec: "aac",
|
||
AudioBitrate: "64k",
|
||
AudioSampleRate: 44100,
|
||
AudioChannels: 2,
|
||
}
|
||
|
||
transcoder, err := NewFFmpegTranscoder(config)
|
||
if err != nil {
|
||
return fmt.Errorf("创建 H.264 转码器失败: %w", err)
|
||
}
|
||
|
||
if err := transcoder.Start(); err != nil {
|
||
return fmt.Errorf("启动 H.264 转码器失败: %w", err)
|
||
}
|
||
|
||
session.transcoder = transcoder
|
||
|
||
// 启动协程读取转码后的 FLV 数据并广播
|
||
go p.readTranscodedOutput(session)
|
||
|
||
log.Printf("✅ [WSProxy] FFmpeg 转码器已初始化（H.264 copy 模式）| Stream:%s", session.StreamID)
|
||
return nil
|
||
}
|
||
|
||
// readTranscodedOutput 读取转码后的 FLV 输出并广播
|
||
func (p *WebMToRTMPProxy) readTranscodedOutput(session *ProxySession) {
|
||
log.Printf("▶️ [WSProxy] 开始读取转码输出 | Stream:%s", session.StreamID)
|
||
|
||
flvBuffer := bytes.NewBuffer(nil)
|
||
flvHeaderParsed := false
|
||
|
||
for data := range session.transcoder.Output() {
|
||
flvBuffer.Write(data)
|
||
|
||
// 解析 FLV 数据
|
||
for {
|
||
bufData := flvBuffer.Bytes()
|
||
|
||
// 首先解析 FLV 头（9 字节）+ PreviousTagSize（4 字节）
|
||
if !flvHeaderParsed {
|
||
if len(bufData) < 13 {
|
||
break
|
||
}
|
||
// 跳过 FLV 头
|
||
flvBuffer.Next(13)
|
||
flvHeaderParsed = true
|
||
bufData = flvBuffer.Bytes()
|
||
}
|
||
|
||
// 解析 FLV Tag
|
||
if len(bufData) < 11 {
|
||
break
|
||
}
|
||
|
||
// Tag 头
|
||
tagType := bufData[0]
|
||
dataSize := int(bufData[1])<<16 | int(bufData[2])<<8 | int(bufData[3])
|
||
|
||
// 检查是否有完整的 tag（11 字节头 + dataSize + 4 字节 PreviousTagSize）
|
||
totalTagSize := 11 + dataSize + 4
|
||
if len(bufData) < totalTagSize {
|
||
break
|
||
}
|
||
|
||
// 提取完整的 FLV tag（不包括 PreviousTagSize）
|
||
flvTag := make([]byte, 11+dataSize+4)
|
||
copy(flvTag, bufData[:totalTagSize])
|
||
|
||
// 广播 FLV tag
|
||
if tagType == 8 || tagType == 9 { // 音频或视频
|
||
p.broadcastFLVTag(session, flvTag)
|
||
}
|
||
|
||
// 从缓冲区移除已处理的数据
|
||
flvBuffer.Next(totalTagSize)
|
||
}
|
||
}
|
||
|
||
log.Printf("⏹️ [WSProxy] 转码输出读取结束 | Stream:%s", session.StreamID)
|
||
}
|
||
|
||
// processWithTranscoder 使用转码器处理数据
|
||
func (p *WebMToRTMPProxy) processWithTranscoder(session *ProxySession, data []byte) error {
|
||
if session.transcoder == nil {
|
||
return fmt.Errorf("转码器未初始化")
|
||
}
|
||
|
||
_, err := session.transcoder.Write(data)
|
||
return err
|
||
}
|
||
|
||
// processWebMData 处理 WebM 数据
|
||
func (p *WebMToRTMPProxy) processWebMData(session *ProxySession, data []byte) error {
|
||
// 将数据追加到缓冲区
|
||
session.webmBuffer.Write(data)
|
||
|
||
// 尝试解析 WebM 数据
|
||
return p.parseAndConvert(session)
|
||
}
|
||
|
||
// parseAndConvert 解析 WebM 并转换为 FLV
|
||
func (p *WebMToRTMPProxy) parseAndConvert(session *ProxySession) error {
|
||
bufData := session.webmBuffer.Bytes()
|
||
if len(bufData) < 4 {
|
||
return nil // 数据不足
|
||
}
|
||
|
||
// 检查 EBML 头部
|
||
if !session.headerParsed {
|
||
// 尝试解析头部
|
||
if err := p.parseWebMHeader(session, bufData); err != nil {
|
||
// 头部不完整，等待更多数据
|
||
return nil
|
||
}
|
||
session.headerParsed = true
|
||
log.Printf("📦 [WSProxy] WebM 头部解析完成 | Stream:%s", session.StreamID)
|
||
}
|
||
|
||
// 解析 Cluster 并转换
|
||
return p.parseWebMClusters(session)
|
||
}
|
||
|
||
// parseWebMHeader 解析 WebM 头部
|
||
func (p *WebMToRTMPProxy) parseWebMHeader(session *ProxySession, data []byte) error {
|
||
reader := bytes.NewReader(data)
|
||
|
||
// 解析 EBML 头
|
||
var header struct {
|
||
EBML struct {
|
||
EBMLVersion uint64 `ebml:"EBMLVersion"`
|
||
DocType string `ebml:"DocType"`
|
||
} `ebml:"EBML"`
|
||
}
|
||
|
||
if err := ebml.Unmarshal(reader, &header); err != nil {
|
||
return err
|
||
}
|
||
|
||
log.Printf("📦 [WSProxy] WebM DocType: %s", header.EBML.DocType)
|
||
return nil
|
||
}
|
||
|
||
// parseWebMClusters 解析 WebM Clusters
|
||
func (p *WebMToRTMPProxy) parseWebMClusters(session *ProxySession) error {
|
||
// 简化处理：直接将 WebM 数据转换为 FLV
|
||
// 实际实现需要完整解析 WebM 的 Cluster/SimpleBlock
|
||
|
||
bufData := session.webmBuffer.Bytes()
|
||
if len(bufData) < 100 {
|
||
return nil
|
||
}
|
||
|
||
// 寻找 Cluster 标记 (0x1F43B675)
|
||
clusterMarker := []byte{0x1F, 0x43, 0xB6, 0x75}
|
||
|
||
for {
|
||
idx := bytes.Index(bufData, clusterMarker)
|
||
if idx == -1 || idx+12 > len(bufData) {
|
||
break
|
||
}
|
||
|
||
// 解析 Cluster 大小
|
||
sizeStart := idx + 4
|
||
clusterSize, bytesRead := readVarInt(bufData[sizeStart:])
|
||
if bytesRead == 0 {
|
||
break
|
||
}
|
||
|
||
totalSize := idx + 4 + bytesRead + int(clusterSize)
|
||
if totalSize > len(bufData) {
|
||
// Cluster 不完整
|
||
break
|
||
}
|
||
|
||
// 提取 Cluster 数据
|
||
clusterData := bufData[idx:totalSize]
|
||
|
||
// 转换为 FLV 并广播
|
||
if err := p.convertClusterToFLV(session, clusterData); err != nil {
|
||
log.Printf("⚠️ [WSProxy] 转换 Cluster 失败: %v", err)
|
||
}
|
||
|
||
// 从缓冲区移除已处理的数据
|
||
bufData = bufData[totalSize:]
|
||
}
|
||
|
||
// 更新缓冲区
|
||
session.webmBuffer.Reset()
|
||
session.webmBuffer.Write(bufData)
|
||
|
||
return nil
|
||
}
|
||
|
||
// readVarInt 读取 EBML 变长整数
|
||
func readVarInt(data []byte) (uint64, int) {
|
||
if len(data) == 0 {
|
||
return 0, 0
|
||
}
|
||
|
||
first := data[0]
|
||
var length int
|
||
var mask byte
|
||
|
||
switch {
|
||
case first&0x80 != 0:
|
||
length = 1
|
||
mask = 0x7F
|
||
case first&0x40 != 0:
|
||
length = 2
|
||
mask = 0x3F
|
||
case first&0x20 != 0:
|
||
length = 3
|
||
mask = 0x1F
|
||
case first&0x10 != 0:
|
||
length = 4
|
||
mask = 0x0F
|
||
case first&0x08 != 0:
|
||
length = 5
|
||
mask = 0x07
|
||
case first&0x04 != 0:
|
||
length = 6
|
||
mask = 0x03
|
||
case first&0x02 != 0:
|
||
length = 7
|
||
mask = 0x01
|
||
case first&0x01 != 0:
|
||
length = 8
|
||
mask = 0x00
|
||
default:
|
||
return 0, 0
|
||
}
|
||
|
||
if len(data) < length {
|
||
return 0, 0
|
||
}
|
||
|
||
value := uint64(data[0] & mask)
|
||
for i := 1; i < length; i++ {
|
||
value = (value << 8) | uint64(data[i])
|
||
}
|
||
|
||
return value, length
|
||
}
|
||
|
||
// convertClusterToFLV 将 WebM Cluster 转换为 FLV 并广播
|
||
func (p *WebMToRTMPProxy) convertClusterToFLV(session *ProxySession, clusterData []byte) error {
|
||
// 解析 Cluster 时间戳
|
||
timecodeMarker := []byte{0xE7} // Timecode element ID
|
||
timecodeIdx := bytes.Index(clusterData, timecodeMarker)
|
||
|
||
var timestamp uint32 = session.lastTimestamp
|
||
if timecodeIdx != -1 && timecodeIdx+1 < len(clusterData) {
|
||
tcSize, bytesRead := readVarInt(clusterData[timecodeIdx+1:])
|
||
if bytesRead > 0 && timecodeIdx+1+bytesRead+int(tcSize) <= len(clusterData) {
|
||
tcData := clusterData[timecodeIdx+1+bytesRead : timecodeIdx+1+bytesRead+int(tcSize)]
|
||
for _, b := range tcData {
|
||
timestamp = (timestamp << 8) | uint32(b)
|
||
}
|
||
}
|
||
}
|
||
|
||
// 解析 SimpleBlock
|
||
simpleBlockMarker := []byte{0xA3} // SimpleBlock element ID
|
||
blockData := clusterData
|
||
|
||
for {
|
||
blockIdx := bytes.Index(blockData, simpleBlockMarker)
|
||
if blockIdx == -1 || blockIdx+1 >= len(blockData) {
|
||
break
|
||
}
|
||
|
||
// 解析 block 大小
|
||
sizeStart := blockIdx + 1
|
||
blockSize, bytesRead := readVarInt(blockData[sizeStart:])
|
||
if bytesRead == 0 {
|
||
break
|
||
}
|
||
|
||
dataStart := sizeStart + bytesRead
|
||
dataEnd := dataStart + int(blockSize)
|
||
if dataEnd > len(blockData) {
|
||
break
|
||
}
|
||
|
||
// 提取 block 内容
|
||
block := blockData[dataStart:dataEnd]
|
||
if len(block) < 4 {
|
||
blockData = blockData[dataEnd:]
|
||
continue
|
||
}
|
||
|
||
// 解析 track number（变长）
|
||
trackNum, trackBytes := readVarInt(block)
|
||
if trackBytes == 0 {
|
||
blockData = blockData[dataEnd:]
|
||
continue
|
||
}
|
||
|
||
// 解析相对时间戳（2 bytes, big-endian）
|
||
if trackBytes+2 > len(block) {
|
||
blockData = blockData[dataEnd:]
|
||
continue
|
||
}
|
||
relativeTimestamp := binary.BigEndian.Uint16(block[trackBytes : trackBytes+2])
|
||
|
||
// 解析 flags (1 byte)
|
||
flagsIdx := trackBytes + 2
|
||
if flagsIdx >= len(block) {
|
||
blockData = blockData[dataEnd:]
|
||
continue
|
||
}
|
||
flags := block[flagsIdx]
|
||
|
||
// 帧数据
|
||
frameData := block[flagsIdx+1:]
|
||
if len(frameData) == 0 {
|
||
blockData = blockData[dataEnd:]
|
||
continue
|
||
}
|
||
|
||
// 计算绝对时间戳
|
||
absoluteTimestamp := timestamp + uint32(relativeTimestamp)
|
||
session.lastTimestamp = absoluteTimestamp
|
||
|
||
// 判断是视频还是音频（简化：假设 track 1 是视频，track 2 是音频）
|
||
isVideo := trackNum == 1
|
||
isKeyframe := (flags & 0x80) != 0
|
||
|
||
// 创建 FLV tag 并广播
|
||
var flvTag []byte
|
||
if isVideo {
|
||
flvTag = p.createFLVVideoTag(absoluteTimestamp, frameData, isKeyframe)
|
||
} else {
|
||
flvTag = p.createFLVAudioTag(absoluteTimestamp, frameData)
|
||
}
|
||
|
||
if flvTag != nil {
|
||
p.broadcastFLVTag(session, flvTag)
|
||
}
|
||
|
||
blockData = blockData[dataEnd:]
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// createFLVVideoTag 创建 FLV 视频 tag
|
||
// 根据会话的编解码器信息选择正确的封装方式
|
||
func (p *WebMToRTMPProxy) createFLVVideoTag(timestamp uint32, data []byte, isKeyframe bool) []byte {
|
||
// FLV Video Tag Header:
|
||
// FrameType (4 bits): 1=keyframe, 2=inter frame
|
||
// CodecID (4 bits): 7=AVC (H.264), 12=VP8 (非标准)
|
||
|
||
frameType := byte(2) // inter frame
|
||
if isKeyframe {
|
||
frameType = 1 // keyframe
|
||
}
|
||
|
||
// 默认使用 VP8（非标准），如果是 H.264 则使用标准格式
|
||
// 注意：VP8 不能被标准 FLV 播放器播放，需要转码
|
||
// 这里保留 VP8 封装作为降级方案
|
||
codecID := byte(12) // 自定义：VP8
|
||
|
||
header := (frameType << 4) | codecID
|
||
|
||
// 构建完整数据
|
||
videoData := make([]byte, 1+len(data))
|
||
videoData[0] = header
|
||
copy(videoData[1:], data)
|
||
|
||
return p.createFLVTag(9, timestamp, videoData) // 9 = video
|
||
}
|
||
|
||
// createFLVVideoTagH264 创建 H.264 FLV 视频 tag
|
||
// H.264 可以直接封装到标准 FLV
|
||
func (p *WebMToRTMPProxy) createFLVVideoTagH264(timestamp uint32, data []byte, isKeyframe bool) []byte {
|
||
// FLV Video Tag Header:
|
||
// FrameType (4 bits): 1=keyframe, 2=inter frame
|
||
// CodecID (4 bits): 7=AVC (H.264)
|
||
|
||
frameType := byte(2) // inter frame
|
||
if isKeyframe {
|
||
frameType = 1 // keyframe
|
||
}
|
||
|
||
codecID := byte(7) // AVC (H.264)
|
||
header := (frameType << 4) | codecID
|
||
|
||
// AVC 数据需要额外的封装
|
||
// AVCPacketType: 0=AVC sequence header, 1=AVC NALU
|
||
// CompositionTime: 3 bytes (通常为 0)
|
||
avcPacketType := byte(1) // AVC NALU
|
||
if isKeyframe {
|
||
// 关键帧可能需要先发送 sequence header
|
||
// 这里简化处理，假设数据已经是正确格式
|
||
}
|
||
|
||
// 构建 AVC 数据
|
||
// 1 byte header + 1 byte AVCPacketType + 3 bytes CompositionTime + data
|
||
videoData := make([]byte, 5+len(data))
|
||
videoData[0] = header
|
||
videoData[1] = avcPacketType
|
||
videoData[2] = 0 // CompositionTime
|
||
videoData[3] = 0
|
||
videoData[4] = 0
|
||
copy(videoData[5:], data)
|
||
|
||
return p.createFLVTag(9, timestamp, videoData) // 9 = video
|
||
}
|
||
|
||
// createFLVAudioTag 创建 FLV 音频 tag
|
||
// Opus -> FLV (需要转换为 AAC，这里简化处理)
|
||
func (p *WebMToRTMPProxy) createFLVAudioTag(timestamp uint32, data []byte) []byte {
|
||
// 注意：Opus 不能直接封装到 FLV，需要转码为 AAC
|
||
// 这里使用简化方案
|
||
|
||
// FLV Audio Tag Header:
|
||
// SoundFormat (4 bits): 10=AAC, 13=Opus (非标准)
|
||
// SoundRate (2 bits): 3=44kHz
|
||
// SoundSize (1 bit): 1=16-bit
|
||
// SoundType (1 bit): 1=stereo
|
||
|
||
// 使用自定义格式标记 Opus
|
||
soundFormat := byte(13) // 自定义：Opus
|
||
soundRate := byte(3) // 44kHz
|
||
soundSize := byte(1) // 16-bit
|
||
soundType := byte(1) // stereo
|
||
|
||
header := (soundFormat << 4) | (soundRate << 2) | (soundSize << 1) | soundType
|
||
|
||
// 构建完整数据
|
||
audioData := make([]byte, 1+len(data))
|
||
audioData[0] = header
|
||
copy(audioData[1:], data)
|
||
|
||
return p.createFLVTag(8, timestamp, audioData) // 8 = audio
|
||
}
|
||
|
||
// createFLVTag 创建 FLV tag
|
||
func (p *WebMToRTMPProxy) createFLVTag(tagType byte, timestamp uint32, data []byte) []byte {
|
||
dataSize := len(data)
|
||
tagSize := 11 + dataSize + 4
|
||
|
||
tag := make([]byte, tagSize)
|
||
|
||
// Tag type
|
||
tag[0] = tagType
|
||
|
||
// Data size (24-bit big-endian)
|
||
tag[1] = byte((dataSize >> 16) & 0xff)
|
||
tag[2] = byte((dataSize >> 8) & 0xff)
|
||
tag[3] = byte(dataSize & 0xff)
|
||
|
||
// Timestamp (24-bit big-endian)
|
||
tag[4] = byte((timestamp >> 16) & 0xff)
|
||
tag[5] = byte((timestamp >> 8) & 0xff)
|
||
tag[6] = byte(timestamp & 0xff)
|
||
|
||
// Timestamp extended
|
||
tag[7] = byte((timestamp >> 24) & 0xff)
|
||
|
||
// Stream ID (always 0)
|
||
tag[8] = 0
|
||
tag[9] = 0
|
||
tag[10] = 0
|
||
|
||
// Data
|
||
copy(tag[11:], data)
|
||
|
||
// Previous tag size
|
||
prevTagSize := 11 + dataSize
|
||
tag[11+dataSize] = byte((prevTagSize >> 24) & 0xff)
|
||
tag[11+dataSize+1] = byte((prevTagSize >> 16) & 0xff)
|
||
tag[11+dataSize+2] = byte((prevTagSize >> 8) & 0xff)
|
||
tag[11+dataSize+3] = byte(prevTagSize & 0xff)
|
||
|
||
return tag
|
||
}
|
||
|
||
// broadcastFLVTag 广播 FLV tag 到订阅者
|
||
func (p *WebMToRTMPProxy) broadcastFLVTag(session *ProxySession, tag []byte) {
|
||
if session.Stream == nil {
|
||
return
|
||
}
|
||
|
||
// 缓存到 GOP
|
||
session.Stream.mu.Lock()
|
||
session.Stream.gopCache = append(session.Stream.gopCache, tag)
|
||
if len(session.Stream.gopCache) > 300 {
|
||
session.Stream.gopCache = session.Stream.gopCache[len(session.Stream.gopCache)-300:]
|
||
}
|
||
session.Stream.mu.Unlock()
|
||
|
||
// 广播到 HTTP-FLV 订阅者
|
||
p.rtmpServer.subMu.RLock()
|
||
subs, exists := p.rtmpServer.subscribers[session.StreamID]
|
||
if exists {
|
||
for _, sub := range subs {
|
||
select {
|
||
case sub.DataChan <- tag:
|
||
default:
|
||
// 缓冲区满，跳过
|
||
}
|
||
}
|
||
}
|
||
p.rtmpServer.subMu.RUnlock()
|
||
}
|
||
|
||
// GetSession 获取会话
|
||
func (p *WebMToRTMPProxy) GetSession(sessionID string) *ProxySession {
|
||
p.sessionsMu.RLock()
|
||
defer p.sessionsMu.RUnlock()
|
||
return p.sessions[sessionID]
|
||
}
|
||
|
||
// GetSessionsByStream 获取流的所有会话
|
||
func (p *WebMToRTMPProxy) GetSessionsByStream(streamID string) []*ProxySession {
|
||
p.sessionsMu.RLock()
|
||
defer p.sessionsMu.RUnlock()
|
||
|
||
sessions := make([]*ProxySession, 0)
|
||
for _, s := range p.sessions {
|
||
if s.StreamID == streamID {
|
||
sessions = append(sessions, s)
|
||
}
|
||
}
|
||
return sessions
|
||
}
|
||
|
||
// CloseSession 关闭会话
|
||
func (p *WebMToRTMPProxy) CloseSession(sessionID string) {
|
||
p.sessionsMu.Lock()
|
||
session, exists := p.sessions[sessionID]
|
||
if exists {
|
||
delete(p.sessions, sessionID)
|
||
}
|
||
p.sessionsMu.Unlock()
|
||
|
||
if session != nil {
|
||
select {
|
||
case <-session.StopChan:
|
||
default:
|
||
close(session.StopChan)
|
||
}
|
||
session.Conn.Close()
|
||
}
|
||
}
|
||
|
||
// CloseAllSessions 关闭所有会话
|
||
func (p *WebMToRTMPProxy) CloseAllSessions() {
|
||
p.sessionsMu.Lock()
|
||
sessions := make([]*ProxySession, 0, len(p.sessions))
|
||
for _, s := range p.sessions {
|
||
sessions = append(sessions, s)
|
||
}
|
||
p.sessions = make(map[string]*ProxySession)
|
||
p.sessionsMu.Unlock()
|
||
|
||
for _, s := range sessions {
|
||
select {
|
||
case <-s.StopChan:
|
||
default:
|
||
close(s.StopChan)
|
||
}
|
||
s.Conn.Close()
|
||
}
|
||
}
|
||
|
||
// 确保导入被使用
|
||
var _ io.Reader = (*bytes.Reader)(nil)
|