1223 lines
29 KiB
Go
1223 lines
29 KiB
Go
/**
|
||
* package mediaserver
|
||
*
|
||
* RTMP 服务（用于微信小程序 live-pusher/live-player）
|
||
* 自实现 RTMP 协议，解决第三方库的 SetChunkSize panic 问题
|
||
*
|
||
* 功能：
|
||
* 1. 接收小程序推流（publish）
|
||
* 2. 提供拉流服务（play）
|
||
* 3. 支持 HTTP-FLV 协议
|
||
*/
|
||
package mediaserver
|
||
|
||
import (
|
||
"crypto/hmac"
|
||
"crypto/sha1"
|
||
"encoding/base64"
|
||
"encoding/binary"
|
||
"fmt"
|
||
"log"
|
||
"net"
|
||
"net/http"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/spf13/viper"
|
||
)
|
||
|
||
// RTMPServer RTMP 服务器
|
||
type RTMPServer struct {
|
||
config *MediaServerConfig
|
||
mu sync.RWMutex
|
||
running bool
|
||
streams map[string]*RTMPStream // streamID -> stream
|
||
streamsMu sync.RWMutex
|
||
|
||
// 服务器监听器
|
||
rtmpListener net.Listener
|
||
httpServer *http.Server
|
||
|
||
// 订阅者管理
|
||
subscribers map[string]map[string]*Subscriber // streamID -> subscriberID -> subscriber
|
||
subMu sync.RWMutex
|
||
}
|
||
|
||
// RTMPStream RTMP 流信息
|
||
type RTMPStream struct {
|
||
ID string `json:"id"`
|
||
RoomID string `json:"room_id"`
|
||
UserID string `json:"user_id"`
|
||
PushURL string `json:"push_url"`
|
||
PullURL string `json:"pull_url"`
|
||
FLVURL string `json:"flv_url"`
|
||
CreatedAt time.Time `json:"created_at"`
|
||
IsActive bool `json:"is_active"`
|
||
|
||
// 流数据
|
||
mu sync.RWMutex
|
||
flvHeader []byte // FLV header
|
||
metaData []byte // FLV metadata tag
|
||
videoHeader []byte // Video sequence header (AVC)
|
||
audioHeader []byte // Audio sequence header (AAC)
|
||
gopCache [][]byte // GOP cache for new subscribers
|
||
|
||
// 控制
|
||
stopChan chan struct{} `json:"-"`
|
||
}
|
||
|
||
// Subscriber 订阅者
|
||
type Subscriber struct {
|
||
ID string
|
||
StreamID string
|
||
DataChan chan []byte
|
||
Done chan struct{}
|
||
}
|
||
|
||
// NewRTMPServer 创建 RTMP 服务器
|
||
func NewRTMPServer(config *MediaServerConfig) *RTMPServer {
|
||
return &RTMPServer{
|
||
config: config,
|
||
streams: make(map[string]*RTMPStream),
|
||
subscribers: make(map[string]map[string]*Subscriber),
|
||
}
|
||
}
|
||
|
||
// Start 启动 RTMP 服务器
|
||
func (r *RTMPServer) Start() {
|
||
r.mu.Lock()
|
||
if r.running {
|
||
r.mu.Unlock()
|
||
return
|
||
}
|
||
r.running = true
|
||
r.mu.Unlock()
|
||
|
||
// 启动 RTMP 服务器
|
||
go r.startRTMPServer()
|
||
|
||
// 启动 HTTP-FLV 服务器
|
||
go r.startHTTPFLVServer()
|
||
|
||
// 启动流清理任务
|
||
go r.cleanupTask()
|
||
|
||
log.Printf("🚀 [RTMP] RTMP 服务已启动 | RTMP端口: %d | HTTP-FLV端口: %d",
|
||
r.config.RTMPPort, r.config.HTTPFLVPort)
|
||
}
|
||
|
||
// startRTMPServer 启动 RTMP 服务
|
||
func (r *RTMPServer) startRTMPServer() {
|
||
addr := fmt.Sprintf(":%d", r.config.RTMPPort)
|
||
|
||
var err error
|
||
r.rtmpListener, err = net.Listen("tcp", addr)
|
||
if err != nil {
|
||
log.Printf("❌ [RTMP] RTMP 监听失败: %v", err)
|
||
return
|
||
}
|
||
|
||
log.Printf("🎬 [RTMP] RTMP 服务监听: %s", addr)
|
||
|
||
for {
|
||
conn, err := r.rtmpListener.Accept()
|
||
if err != nil {
|
||
r.mu.RLock()
|
||
running := r.running
|
||
r.mu.RUnlock()
|
||
if !running {
|
||
return
|
||
}
|
||
log.Printf("⚠️ [RTMP] Accept 失败: %v", err)
|
||
continue
|
||
}
|
||
|
||
go r.handleConnection(conn)
|
||
}
|
||
}
|
||
|
||
// handleConnection 处理 RTMP 连接
|
||
func (r *RTMPServer) handleConnection(conn net.Conn) {
|
||
remoteAddr := conn.RemoteAddr().String()
|
||
log.Printf("🎬 [RTMP] 新连接: %s", remoteAddr)
|
||
|
||
defer func() {
|
||
if err := recover(); err != nil {
|
||
log.Printf("❌ [RTMP] 连接处理 panic: %v", err)
|
||
}
|
||
conn.Close()
|
||
}()
|
||
|
||
// 1. 执行握手
|
||
if err := DoHandshake(conn, 30*time.Second); err != nil {
|
||
log.Printf("❌ [RTMP] 握手失败: %v", err)
|
||
return
|
||
}
|
||
|
||
// 2. 创建读写器
|
||
reader := NewChunkReader(conn)
|
||
writer := NewChunkWriter(conn)
|
||
|
||
// 3. 创建连接处理器
|
||
handler := &ConnectionHandler{
|
||
server: r,
|
||
conn: conn,
|
||
reader: reader,
|
||
writer: writer,
|
||
remoteAddr: remoteAddr,
|
||
}
|
||
|
||
// 4. 消息处理循环
|
||
handler.serve()
|
||
}
|
||
|
||
// ConnectionHandler RTMP 连接处理器
|
||
type ConnectionHandler struct {
|
||
server *RTMPServer
|
||
conn net.Conn
|
||
reader *ChunkReader
|
||
writer *ChunkWriter
|
||
remoteAddr string
|
||
|
||
// 连接状态
|
||
streamID string
|
||
isPublish bool
|
||
appName string
|
||
msgStreamID uint32
|
||
|
||
// 统计信息
|
||
audioCount int64
|
||
videoCount int64
|
||
firstAudio bool
|
||
firstVideo bool
|
||
}
|
||
|
||
// serve 消息处理循环
|
||
func (h *ConnectionHandler) serve() {
|
||
log.Printf("🎬 [RTMP] OnServe: %s", h.remoteAddr)
|
||
|
||
msgCount := int64(0)
|
||
lastLogTime := time.Now()
|
||
|
||
for {
|
||
msg, err := h.reader.ReadMessage()
|
||
if err != nil {
|
||
log.Printf("🔌 [RTMP] 连接关闭: %s (err: %v)", h.remoteAddr, err)
|
||
break
|
||
}
|
||
|
||
msgCount++
|
||
|
||
// 每 100 条消息或每 5 秒记录一次统计信息
|
||
if msgCount%100 == 0 || time.Since(lastLogTime) > 5*time.Second {
|
||
log.Printf("📊 [RTMP] 消息统计: stream=%s 总消息=%d audio=%d video=%d remote=%s",
|
||
h.streamID, msgCount, h.audioCount, h.videoCount, h.remoteAddr)
|
||
lastLogTime = time.Now()
|
||
}
|
||
|
||
if err := h.handleMessage(msg); err != nil {
|
||
log.Printf("❌ [RTMP] 处理消息失败: %v", err)
|
||
break
|
||
}
|
||
}
|
||
|
||
// 连接关闭时的清理
|
||
h.onClose()
|
||
}
|
||
|
||
// handleMessage 处理单个消息
|
||
func (h *ConnectionHandler) handleMessage(msg *RTMPMessage) error {
|
||
switch msg.TypeID {
|
||
case RTMP_MSG_CHUNK_SIZE:
|
||
return h.handleSetChunkSize(msg)
|
||
case RTMP_MSG_ACK:
|
||
// 忽略 ACK
|
||
return nil
|
||
case RTMP_MSG_USER_CONTROL:
|
||
// 忽略 User Control
|
||
return nil
|
||
case RTMP_MSG_WIN_ACK_SIZE:
|
||
// 忽略 Window Ack Size
|
||
return nil
|
||
case RTMP_MSG_SET_PEER_BW:
|
||
// 忽略 Set Peer Bandwidth
|
||
return nil
|
||
case RTMP_MSG_AUDIO:
|
||
return h.handleAudio(msg)
|
||
case RTMP_MSG_VIDEO:
|
||
return h.handleVideo(msg)
|
||
case RTMP_MSG_AMF0_DATA:
|
||
return h.handleDataMessage(msg)
|
||
case RTMP_MSG_AMF0_CMD:
|
||
return h.handleCommand(msg)
|
||
default:
|
||
// 忽略未知消息
|
||
return nil
|
||
}
|
||
}
|
||
|
||
// handleSetChunkSize 处理 SetChunkSize 消息
|
||
func (h *ConnectionHandler) handleSetChunkSize(msg *RTMPMessage) error {
|
||
if len(msg.Data) < 4 {
|
||
return fmt.Errorf("SetChunkSize 数据不足")
|
||
}
|
||
|
||
newSize := binary.BigEndian.Uint32(msg.Data)
|
||
log.Printf("📝 [RTMP] SetChunkSize: %d -> %d (remote: %s)", h.reader.GetChunkSize(), newSize, h.remoteAddr)
|
||
|
||
// 关键：立即更新 chunk size
|
||
h.reader.SetChunkSize(newSize)
|
||
return nil
|
||
}
|
||
|
||
// handleCommand 处理 AMF0 命令
|
||
func (h *ConnectionHandler) handleCommand(msg *RTMPMessage) error {
|
||
values, err := DecodeAMF0(msg.Data)
|
||
if err != nil {
|
||
log.Printf("⚠️ [RTMP] 解析命令失败: %v", err)
|
||
return nil
|
||
}
|
||
|
||
if len(values) == 0 {
|
||
return nil
|
||
}
|
||
|
||
command, ok := values[0].(string)
|
||
if !ok {
|
||
return nil
|
||
}
|
||
|
||
transactionID := float64(0)
|
||
if len(values) > 1 {
|
||
if tid, ok := values[1].(float64); ok {
|
||
transactionID = tid
|
||
}
|
||
}
|
||
|
||
log.Printf("🎬 [RTMP] 命令: %s (tid: %.0f) from %s", command, transactionID, h.remoteAddr)
|
||
|
||
switch command {
|
||
case "connect":
|
||
return h.handleConnect(values, transactionID)
|
||
case "releaseStream":
|
||
return h.handleReleaseStream(values)
|
||
case "FCPublish":
|
||
return h.handleFCPublish(values)
|
||
case "createStream":
|
||
return h.handleCreateStream(transactionID)
|
||
case "publish":
|
||
return h.handlePublish(values, msg.StreamID)
|
||
case "play":
|
||
return h.handlePlay(values, msg.StreamID)
|
||
case "deleteStream":
|
||
return h.handleDeleteStream(values)
|
||
case "FCUnpublish":
|
||
return h.handleFCUnpublish(values)
|
||
default:
|
||
// 忽略未知命令
|
||
return nil
|
||
}
|
||
}
|
||
|
||
// handleConnect 处理 connect 命令
|
||
func (h *ConnectionHandler) handleConnect(values []interface{}, tid float64) error {
|
||
// 解析 app 名称
|
||
if len(values) > 2 {
|
||
if obj, ok := values[2].(AMF0Object); ok {
|
||
if app, ok := obj["app"].(string); ok {
|
||
h.appName = app
|
||
}
|
||
}
|
||
}
|
||
log.Printf("🎬 [RTMP] OnConnect: app=%s", h.appName)
|
||
|
||
// 发送 Window Ack Size
|
||
if err := h.writer.WriteWindowAckSize(DEFAULT_WINDOW_SIZE); err != nil {
|
||
return err
|
||
}
|
||
|
||
// 发送 Set Peer Bandwidth
|
||
if err := h.writer.WriteSetPeerBandwidth(DEFAULT_WINDOW_SIZE, 2); err != nil {
|
||
return err
|
||
}
|
||
|
||
// 发送 Set Chunk Size
|
||
if err := h.writer.WriteSetChunkSize(4096); err != nil {
|
||
return err
|
||
}
|
||
|
||
// 发送 _result
|
||
result := EncodeConnectResult(tid)
|
||
if err := h.writer.WriteCommand(3, 0, result); err != nil {
|
||
return err
|
||
}
|
||
|
||
// 发送 onBWDone
|
||
bwDone := EncodeOnBWDone()
|
||
return h.writer.WriteCommand(3, 0, bwDone)
|
||
}
|
||
|
||
// handleReleaseStream 处理 releaseStream 命令
|
||
func (h *ConnectionHandler) handleReleaseStream(values []interface{}) error {
|
||
streamName := ""
|
||
if len(values) > 3 {
|
||
if name, ok := values[3].(string); ok {
|
||
streamName = name
|
||
}
|
||
}
|
||
if idx := strings.Index(streamName, "?"); idx != -1 {
|
||
streamName = streamName[:idx]
|
||
}
|
||
log.Printf("🎬 [RTMP] OnReleaseStream: %s", streamName)
|
||
return nil
|
||
}
|
||
|
||
// handleFCPublish 处理 FCPublish 命令
|
||
func (h *ConnectionHandler) handleFCPublish(values []interface{}) error {
|
||
streamName := ""
|
||
if len(values) > 3 {
|
||
if name, ok := values[3].(string); ok {
|
||
streamName = name
|
||
}
|
||
}
|
||
log.Printf("🎬 [RTMP] OnFCPublish: %s", streamName)
|
||
return nil
|
||
}
|
||
|
||
// handleCreateStream 处理 createStream 命令
|
||
func (h *ConnectionHandler) handleCreateStream(tid float64) error {
|
||
log.Printf("🎬 [RTMP] OnCreateStream")
|
||
|
||
h.msgStreamID = 1
|
||
|
||
// 发送 _result
|
||
result := EncodeCreateStreamResult(tid, float64(h.msgStreamID))
|
||
return h.writer.WriteCommand(3, 0, result)
|
||
}
|
||
|
||
// handlePublish 处理 publish 命令
|
||
func (h *ConnectionHandler) handlePublish(values []interface{}, streamID uint32) error {
|
||
publishingName := ""
|
||
publishingType := ""
|
||
|
||
if len(values) > 3 {
|
||
if name, ok := values[3].(string); ok {
|
||
publishingName = name
|
||
}
|
||
}
|
||
if len(values) > 4 {
|
||
if ptype, ok := values[4].(string); ok {
|
||
publishingType = ptype
|
||
}
|
||
}
|
||
|
||
log.Printf("🎬 [RTMP] OnPublish: name=%s, type=%s, remote=%s", publishingName, publishingType, h.remoteAddr)
|
||
|
||
// 解析 stream name，去掉 token 参数
|
||
streamName := publishingName
|
||
if idx := strings.Index(streamName, "?"); idx != -1 {
|
||
streamName = streamName[:idx]
|
||
}
|
||
|
||
h.streamID = streamName
|
||
h.isPublish = true
|
||
h.audioCount = 0
|
||
h.videoCount = 0
|
||
h.firstAudio = false
|
||
h.firstVideo = false
|
||
|
||
log.Printf("🎬 [RTMP] 解析后的流ID: %s", h.streamID)
|
||
|
||
// 检查流是否已注册
|
||
h.server.streamsMu.RLock()
|
||
stream, exists := h.server.streams[h.streamID]
|
||
h.server.streamsMu.RUnlock()
|
||
|
||
if !exists {
|
||
log.Printf("⚠️ [RTMP] 流未注册，自动创建: %s", h.streamID)
|
||
stream = &RTMPStream{
|
||
ID: h.streamID,
|
||
CreatedAt: time.Now(),
|
||
IsActive: true,
|
||
stopChan: make(chan struct{}),
|
||
gopCache: make([][]byte, 0),
|
||
}
|
||
h.server.streamsMu.Lock()
|
||
h.server.streams[h.streamID] = stream
|
||
h.server.streamsMu.Unlock()
|
||
} else {
|
||
log.Printf("✅ [RTMP] 找到已注册的流: %s", h.streamID)
|
||
}
|
||
|
||
stream.mu.Lock()
|
||
stream.IsActive = true
|
||
// 重置流数据
|
||
stream.audioHeader = nil
|
||
stream.videoHeader = nil
|
||
stream.metaData = nil
|
||
stream.gopCache = make([][]byte, 0)
|
||
stream.mu.Unlock()
|
||
|
||
// 发送 Stream Begin
|
||
if err := h.writer.WriteStreamBegin(streamID); err != nil {
|
||
return err
|
||
}
|
||
|
||
// 发送 onStatus (NetStream.Publish.Start)
|
||
status := EncodeOnStatus("NetStream.Publish.Start", "status", "Publishing started.")
|
||
if err := h.writer.WriteCommand(5, streamID, status); err != nil {
|
||
return err
|
||
}
|
||
|
||
log.Printf("✅ [RTMP] 开始推流: %s (等待音视频数据...)", h.streamID)
|
||
return nil
|
||
}
|
||
|
||
// handlePlay 处理 play 命令
|
||
func (h *ConnectionHandler) handlePlay(values []interface{}, streamID uint32) error {
|
||
streamName := ""
|
||
if len(values) > 3 {
|
||
if name, ok := values[3].(string); ok {
|
||
streamName = name
|
||
}
|
||
}
|
||
|
||
log.Printf("🎬 [RTMP] OnPlay: name=%s", streamName)
|
||
|
||
h.streamID = streamName
|
||
h.isPublish = false
|
||
|
||
// 检查流是否存在
|
||
h.server.streamsMu.RLock()
|
||
stream, exists := h.server.streams[h.streamID]
|
||
h.server.streamsMu.RUnlock()
|
||
|
||
if !exists || !stream.IsActive {
|
||
log.Printf("⚠️ [RTMP] 流不存在或未激活: %s", h.streamID)
|
||
status := EncodeOnStatus("NetStream.Play.StreamNotFound", "error", "Stream not found.")
|
||
return h.writer.WriteCommand(5, streamID, status)
|
||
}
|
||
|
||
// 发送 Stream Begin
|
||
if err := h.writer.WriteStreamBegin(streamID); err != nil {
|
||
return err
|
||
}
|
||
|
||
// 发送 onStatus (NetStream.Play.Start)
|
||
status := EncodeOnStatus("NetStream.Play.Start", "status", "Playing started.")
|
||
if err := h.writer.WriteCommand(5, streamID, status); err != nil {
|
||
return err
|
||
}
|
||
|
||
log.Printf("✅ [RTMP] 开始播放: %s", h.streamID)
|
||
return nil
|
||
}
|
||
|
||
// handleDeleteStream 处理 deleteStream 命令
|
||
func (h *ConnectionHandler) handleDeleteStream(values []interface{}) error {
|
||
log.Printf("🎬 [RTMP] OnDeleteStream")
|
||
return nil
|
||
}
|
||
|
||
// handleFCUnpublish 处理 FCUnpublish 命令
|
||
func (h *ConnectionHandler) handleFCUnpublish(values []interface{}) error {
|
||
streamName := ""
|
||
if len(values) > 3 {
|
||
if name, ok := values[3].(string); ok {
|
||
streamName = name
|
||
}
|
||
}
|
||
log.Printf("🎬 [RTMP] OnFCUnpublish: %s", streamName)
|
||
return nil
|
||
}
|
||
|
||
// handleDataMessage 处理数据消息 (@setDataFrame)
|
||
func (h *ConnectionHandler) handleDataMessage(msg *RTMPMessage) error {
|
||
log.Printf("🎬 [RTMP] OnSetDataFrame: stream=%s payloadLen=%d audioCount=%d videoCount=%d",
|
||
h.streamID, len(msg.Data), h.audioCount, h.videoCount)
|
||
|
||
if h.streamID == "" {
|
||
log.Printf("⚠️ [RTMP] OnSetDataFrame: streamID 为空")
|
||
return nil
|
||
}
|
||
|
||
h.server.streamsMu.RLock()
|
||
stream, exists := h.server.streams[h.streamID]
|
||
h.server.streamsMu.RUnlock()
|
||
|
||
if exists && len(msg.Data) > 0 {
|
||
stream.mu.Lock()
|
||
stream.metaData = make([]byte, len(msg.Data))
|
||
copy(stream.metaData, msg.Data)
|
||
stream.mu.Unlock()
|
||
log.Printf("✅ [RTMP] OnSetDataFrame: 已保存 metadata")
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// handleAudio 处理音频数据
|
||
func (h *ConnectionHandler) handleAudio(msg *RTMPMessage) error {
|
||
h.audioCount++
|
||
|
||
if !h.firstAudio {
|
||
h.firstAudio = true
|
||
log.Printf("🎵 [RTMP] OnAudio 首帧: stream=%s timestamp=%d isPublish=%v", h.streamID, msg.Timestamp, h.isPublish)
|
||
}
|
||
|
||
if h.streamID == "" || !h.isPublish {
|
||
return nil
|
||
}
|
||
|
||
data := msg.Data
|
||
if len(data) == 0 {
|
||
return nil
|
||
}
|
||
|
||
if h.audioCount%100 == 0 {
|
||
log.Printf("🎵 [RTMP] OnAudio 统计: stream=%s audioCount=%d videoCount=%d", h.streamID, h.audioCount, h.videoCount)
|
||
}
|
||
|
||
h.server.streamsMu.RLock()
|
||
stream, exists := h.server.streams[h.streamID]
|
||
h.server.streamsMu.RUnlock()
|
||
|
||
if !exists {
|
||
return nil
|
||
}
|
||
|
||
// 保存 AAC sequence header
|
||
if len(data) > 1 {
|
||
soundFormat := (data[0] >> 4) & 0x0f
|
||
if soundFormat == 10 { // AAC
|
||
aacPacketType := data[1]
|
||
if aacPacketType == 0 { // Sequence header
|
||
stream.mu.Lock()
|
||
stream.audioHeader = make([]byte, len(data))
|
||
copy(stream.audioHeader, data)
|
||
stream.mu.Unlock()
|
||
}
|
||
}
|
||
}
|
||
|
||
// 创建 FLV audio tag
|
||
flvData := createFLVTag(8, msg.Timestamp, data)
|
||
h.broadcastToSubscribers(h.streamID, flvData)
|
||
|
||
return nil
|
||
}
|
||
|
||
// handleVideo 处理视频数据
|
||
func (h *ConnectionHandler) handleVideo(msg *RTMPMessage) error {
|
||
h.videoCount++
|
||
|
||
if !h.firstVideo {
|
||
h.firstVideo = true
|
||
log.Printf("🎥 [RTMP] OnVideo 首帧: stream=%s timestamp=%d isPublish=%v", h.streamID, msg.Timestamp, h.isPublish)
|
||
}
|
||
|
||
if h.streamID == "" || !h.isPublish {
|
||
return nil
|
||
}
|
||
|
||
data := msg.Data
|
||
if len(data) == 0 {
|
||
return nil
|
||
}
|
||
|
||
if h.videoCount == 1 {
|
||
log.Printf("🎥 [RTMP] OnVideo 首帧数据: stream=%s len=%d", h.streamID, len(data))
|
||
}
|
||
|
||
if h.videoCount%30 == 0 {
|
||
log.Printf("🎥 [RTMP] OnVideo 统计: stream=%s audioCount=%d videoCount=%d", h.streamID, h.audioCount, h.videoCount)
|
||
}
|
||
|
||
h.server.streamsMu.RLock()
|
||
stream, exists := h.server.streams[h.streamID]
|
||
h.server.streamsMu.RUnlock()
|
||
|
||
if !exists {
|
||
return nil
|
||
}
|
||
|
||
// 解析视频帧信息
|
||
frameType := (data[0] >> 4) & 0x0f
|
||
codecID := data[0] & 0x0f
|
||
|
||
// 保存 AVC sequence header
|
||
if codecID == 7 && len(data) > 1 { // AVC
|
||
avcPacketType := data[1]
|
||
if avcPacketType == 0 { // Sequence header
|
||
stream.mu.Lock()
|
||
stream.videoHeader = make([]byte, len(data))
|
||
copy(stream.videoHeader, data)
|
||
stream.mu.Unlock()
|
||
}
|
||
|
||
// 关键帧时清空 GOP 缓存
|
||
if frameType == 1 { // Keyframe
|
||
stream.mu.Lock()
|
||
stream.gopCache = make([][]byte, 0)
|
||
stream.mu.Unlock()
|
||
}
|
||
}
|
||
|
||
// 创建 FLV video tag
|
||
flvData := createFLVTag(9, msg.Timestamp, data)
|
||
|
||
// 缓存到 GOP
|
||
stream.mu.Lock()
|
||
stream.gopCache = append(stream.gopCache, flvData)
|
||
if len(stream.gopCache) > 300 {
|
||
stream.gopCache = stream.gopCache[len(stream.gopCache)-300:]
|
||
}
|
||
stream.mu.Unlock()
|
||
|
||
h.broadcastToSubscribers(h.streamID, flvData)
|
||
|
||
return nil
|
||
}
|
||
|
||
// broadcastToSubscribers 向所有订阅者广播数据
|
||
func (h *ConnectionHandler) broadcastToSubscribers(streamID string, data []byte) {
|
||
h.server.subMu.RLock()
|
||
subs, exists := h.server.subscribers[streamID]
|
||
if !exists {
|
||
h.server.subMu.RUnlock()
|
||
return
|
||
}
|
||
|
||
for _, sub := range subs {
|
||
select {
|
||
case sub.DataChan <- data:
|
||
default:
|
||
// 缓冲区满，跳过
|
||
}
|
||
}
|
||
h.server.subMu.RUnlock()
|
||
}
|
||
|
||
// onClose 连接关闭
|
||
func (h *ConnectionHandler) onClose() {
|
||
log.Printf("🔌 [RTMP] OnClose: stream=%s, publish=%v, remote=%s, audioCount=%d, videoCount=%d",
|
||
h.streamID, h.isPublish, h.remoteAddr, h.audioCount, h.videoCount)
|
||
|
||
if h.isPublish && h.audioCount > 0 && h.videoCount == 0 {
|
||
log.Printf("⚠️ [RTMP] 警告: 只收到音频(%d帧)，没有收到视频！stream=%s", h.audioCount, h.streamID)
|
||
log.Printf("⚠️ [RTMP] 可能原因: 1.推流端未启用视频 2.视频编码格式不支持 3.连接过早断开")
|
||
}
|
||
|
||
if h.isPublish && h.streamID != "" {
|
||
h.server.streamsMu.Lock()
|
||
if stream, exists := h.server.streams[h.streamID]; exists {
|
||
stream.mu.Lock()
|
||
stream.IsActive = false
|
||
stream.mu.Unlock()
|
||
}
|
||
h.server.streamsMu.Unlock()
|
||
log.Printf("⏹️ [RTMP] 停止推流: %s (总计: 音频%d帧, 视频%d帧)", h.streamID, h.audioCount, h.videoCount)
|
||
}
|
||
}
|
||
|
||
// createFLVTag 创建 FLV tag
|
||
func createFLVTag(tagType byte, timestamp uint32, data []byte) []byte {
|
||
dataSize := len(data)
|
||
tagSize := 11 + dataSize + 4
|
||
|
||
tag := make([]byte, tagSize)
|
||
|
||
tag[0] = tagType
|
||
tag[1] = byte((dataSize >> 16) & 0xff)
|
||
tag[2] = byte((dataSize >> 8) & 0xff)
|
||
tag[3] = byte(dataSize & 0xff)
|
||
tag[4] = byte((timestamp >> 16) & 0xff)
|
||
tag[5] = byte((timestamp >> 8) & 0xff)
|
||
tag[6] = byte(timestamp & 0xff)
|
||
tag[7] = byte((timestamp >> 24) & 0xff)
|
||
tag[8] = 0
|
||
tag[9] = 0
|
||
tag[10] = 0
|
||
|
||
copy(tag[11:], data)
|
||
|
||
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
|
||
}
|
||
|
||
// startHTTPFLVServer 启动 HTTP-FLV 服务
|
||
func (r *RTMPServer) startHTTPFLVServer() {
|
||
addr := fmt.Sprintf(":%d", r.config.HTTPFLVPort)
|
||
|
||
mux := http.NewServeMux()
|
||
mux.HandleFunc("/live/", r.handleFLVRequest)
|
||
mux.HandleFunc("/api/streams", r.handleStreamsAPI)
|
||
|
||
r.httpServer = &http.Server{
|
||
Addr: addr,
|
||
Handler: mux,
|
||
}
|
||
|
||
log.Printf("🎬 [RTMP] HTTP-FLV 服务监听: %s", addr)
|
||
|
||
if err := r.httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||
log.Printf("❌ [RTMP] HTTP-FLV 启动失败: %v", err)
|
||
}
|
||
}
|
||
|
||
// handleFLVRequest 处理 FLV 请求
|
||
func (r *RTMPServer) handleFLVRequest(w http.ResponseWriter, req *http.Request) {
|
||
// 添加 CORS 头，允许跨域访问
|
||
w.Header().Set("Access-Control-Allow-Origin", "*")
|
||
w.Header().Set("Access-Control-Allow-Methods", "GET, OPTIONS")
|
||
w.Header().Set("Access-Control-Allow-Headers", "Origin, Content-Type, Accept, Range")
|
||
w.Header().Set("Access-Control-Expose-Headers", "Content-Length, Content-Range")
|
||
|
||
// 处理 OPTIONS 预检请求
|
||
if req.Method == "OPTIONS" {
|
||
w.WriteHeader(http.StatusOK)
|
||
return
|
||
}
|
||
|
||
path := req.URL.Path
|
||
if len(path) < 7 {
|
||
http.Error(w, "Invalid path", http.StatusBadRequest)
|
||
return
|
||
}
|
||
|
||
streamPath := path[6:] // 去掉 "/live/"
|
||
if len(streamPath) > 4 && streamPath[len(streamPath)-4:] == ".flv" {
|
||
streamPath = streamPath[:len(streamPath)-4]
|
||
}
|
||
|
||
log.Printf("🎬 [RTMP] FLV 请求: %s (from: %s)", streamPath, req.RemoteAddr)
|
||
|
||
// 等待流就绪
|
||
var stream *RTMPStream
|
||
var exists bool
|
||
maxWait := 10 * time.Second
|
||
waitInterval := 500 * time.Millisecond
|
||
waited := time.Duration(0)
|
||
|
||
for waited < maxWait {
|
||
r.streamsMu.RLock()
|
||
stream, exists = r.streams[streamPath]
|
||
r.streamsMu.RUnlock()
|
||
|
||
if exists && stream.IsActive {
|
||
break
|
||
}
|
||
|
||
select {
|
||
case <-req.Context().Done():
|
||
log.Printf("⚠️ [RTMP] FLV 请求已取消: %s", streamPath)
|
||
return
|
||
default:
|
||
}
|
||
|
||
if waited == 0 {
|
||
log.Printf("⏳ [RTMP] 等待流就绪: %s", streamPath)
|
||
}
|
||
|
||
time.Sleep(waitInterval)
|
||
waited += waitInterval
|
||
}
|
||
|
||
if !exists || !stream.IsActive {
|
||
log.Printf("❌ [RTMP] 流不存在或未激活: %s (waited: %v)", streamPath, waited)
|
||
http.Error(w, "Stream not found or not active", http.StatusNotFound)
|
||
return
|
||
}
|
||
|
||
log.Printf("✅ [RTMP] 流已就绪，开始 FLV 传输: %s (waited: %v)", streamPath, waited)
|
||
|
||
// 设置响应头
|
||
w.Header().Set("Content-Type", "video/x-flv")
|
||
w.Header().Set("Access-Control-Allow-Origin", "*")
|
||
w.Header().Set("Transfer-Encoding", "chunked")
|
||
w.Header().Set("Connection", "keep-alive")
|
||
w.Header().Set("Cache-Control", "no-cache")
|
||
|
||
// 发送 FLV header
|
||
flvHeader := []byte{0x46, 0x4C, 0x56, 0x01, 0x05, 0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00, 0x00}
|
||
w.Write(flvHeader)
|
||
|
||
if f, ok := w.(http.Flusher); ok {
|
||
f.Flush()
|
||
}
|
||
|
||
// 创建订阅者
|
||
subID := fmt.Sprintf("%d", time.Now().UnixNano())
|
||
sub := &Subscriber{
|
||
ID: subID,
|
||
StreamID: streamPath,
|
||
DataChan: make(chan []byte, 100),
|
||
Done: make(chan struct{}),
|
||
}
|
||
|
||
// 注册订阅者
|
||
r.subMu.Lock()
|
||
if r.subscribers[streamPath] == nil {
|
||
r.subscribers[streamPath] = make(map[string]*Subscriber)
|
||
}
|
||
r.subscribers[streamPath][subID] = sub
|
||
r.subMu.Unlock()
|
||
|
||
defer func() {
|
||
r.subMu.Lock()
|
||
delete(r.subscribers[streamPath], subID)
|
||
r.subMu.Unlock()
|
||
close(sub.Done)
|
||
}()
|
||
|
||
// 发送缓存的头信息
|
||
stream.mu.RLock()
|
||
hasMetadata := len(stream.metaData) > 0
|
||
hasVideoHeader := len(stream.videoHeader) > 0
|
||
hasAudioHeader := len(stream.audioHeader) > 0
|
||
gopLen := len(stream.gopCache)
|
||
|
||
log.Printf("📦 [RTMP] FLV 缓存状态: stream=%s metadata=%v videoHeader=%v audioHeader=%v gopLen=%d",
|
||
streamPath, hasMetadata, hasVideoHeader, hasAudioHeader, gopLen)
|
||
|
||
if hasMetadata {
|
||
w.Write(stream.metaData)
|
||
}
|
||
|
||
if hasVideoHeader {
|
||
flvTag := createFLVTag(9, 0, stream.videoHeader)
|
||
w.Write(flvTag)
|
||
}
|
||
|
||
if hasAudioHeader {
|
||
flvTag := createFLVTag(8, 0, stream.audioHeader)
|
||
w.Write(flvTag)
|
||
}
|
||
|
||
for _, data := range stream.gopCache {
|
||
w.Write(data)
|
||
}
|
||
stream.mu.RUnlock()
|
||
|
||
if f, ok := w.(http.Flusher); ok {
|
||
f.Flush()
|
||
}
|
||
|
||
log.Printf("▶️ [RTMP] FLV 开始实时传输: stream=%s subscriber=%s", streamPath, subID)
|
||
|
||
// 持续发送数据
|
||
dataCount := 0
|
||
for {
|
||
select {
|
||
case <-req.Context().Done():
|
||
log.Printf("⏹️ [RTMP] FLV 传输结束 (客户端断开): stream=%s dataCount=%d", streamPath, dataCount)
|
||
return
|
||
case data := <-sub.DataChan:
|
||
if _, err := w.Write(data); err != nil {
|
||
log.Printf("⏹️ [RTMP] FLV 传输结束 (写入失败): stream=%s dataCount=%d err=%v", streamPath, dataCount, err)
|
||
return
|
||
}
|
||
dataCount++
|
||
if f, ok := w.(http.Flusher); ok {
|
||
f.Flush()
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// handleStreamsAPI 处理流列表 API
|
||
func (r *RTMPServer) handleStreamsAPI(w http.ResponseWriter, req *http.Request) {
|
||
w.Header().Set("Content-Type", "application/json")
|
||
w.Header().Set("Access-Control-Allow-Origin", "*")
|
||
|
||
r.streamsMu.RLock()
|
||
defer r.streamsMu.RUnlock()
|
||
|
||
streams := make([]map[string]interface{}, 0)
|
||
for _, s := range r.streams {
|
||
if s.IsActive {
|
||
streams = append(streams, map[string]interface{}{
|
||
"id": s.ID,
|
||
"room_id": s.RoomID,
|
||
"user_id": s.UserID,
|
||
"pull_url": s.PullURL,
|
||
"flv_url": s.FLVURL,
|
||
})
|
||
}
|
||
}
|
||
|
||
fmt.Fprintf(w, `{"streams":%d,"data":%v}`, len(streams), streams)
|
||
}
|
||
|
||
// Stop 停止 RTMP 服务器
|
||
func (r *RTMPServer) Stop() {
|
||
r.mu.Lock()
|
||
defer r.mu.Unlock()
|
||
|
||
r.running = false
|
||
|
||
// 关闭 RTMP 监听器
|
||
if r.rtmpListener != nil {
|
||
r.rtmpListener.Close()
|
||
}
|
||
|
||
// 关闭 HTTP 服务器
|
||
if r.httpServer != nil {
|
||
r.httpServer.Close()
|
||
}
|
||
|
||
// 关闭所有流
|
||
r.streamsMu.Lock()
|
||
for _, stream := range r.streams {
|
||
stream.IsActive = false
|
||
if stream.stopChan != nil {
|
||
select {
|
||
case <-stream.stopChan:
|
||
default:
|
||
close(stream.stopChan)
|
||
}
|
||
}
|
||
}
|
||
r.streams = make(map[string]*RTMPStream)
|
||
r.streamsMu.Unlock()
|
||
|
||
// 关闭所有订阅者
|
||
r.subMu.Lock()
|
||
for _, subs := range r.subscribers {
|
||
for _, sub := range subs {
|
||
close(sub.DataChan)
|
||
}
|
||
}
|
||
r.subscribers = make(map[string]map[string]*Subscriber)
|
||
r.subMu.Unlock()
|
||
|
||
log.Println("🛑 [RTMP] 已停止")
|
||
}
|
||
|
||
// GenerateStreamURLs 为用户生成推拉流地址
|
||
func (r *RTMPServer) GenerateStreamURLs(roomID, userID string) (*RTMPStream, error) {
|
||
r.mu.RLock()
|
||
if !r.running {
|
||
r.mu.RUnlock()
|
||
return nil, ErrRTMPNotReady
|
||
}
|
||
r.mu.RUnlock()
|
||
|
||
streamID := fmt.Sprintf("%s_%s_%d", roomID, userID, time.Now().UnixNano())
|
||
|
||
publicIP := r.config.PublicIP
|
||
if publicIP == "" {
|
||
publicIP = viper.GetString("turn.public_ip")
|
||
}
|
||
|
||
token := generateStreamToken(streamID)
|
||
|
||
stream := &RTMPStream{
|
||
ID: streamID,
|
||
RoomID: roomID,
|
||
UserID: userID,
|
||
PushURL: fmt.Sprintf("rtmp://%s:%d/live/%s?token=%s", publicIP, r.config.RTMPPort, streamID, token),
|
||
PullURL: fmt.Sprintf("rtmp://%s:%d/live/%s", publicIP, r.config.RTMPPort, streamID),
|
||
// 使用 HTTPS 域名，通过 Nginx 反向代理到 HTTP-FLV 服务
|
||
// 避免浏览器混合内容安全策略阻止 HTTP 请求
|
||
FLVURL: fmt.Sprintf("https://g-ws.nailaoyun.cn/live/%s.flv", streamID),
|
||
CreatedAt: time.Now(),
|
||
IsActive: true,
|
||
stopChan: make(chan struct{}),
|
||
gopCache: make([][]byte, 0),
|
||
}
|
||
|
||
r.streamsMu.Lock()
|
||
r.streams[streamID] = stream
|
||
r.streamsMu.Unlock()
|
||
|
||
log.Printf("🎬 [RTMP] 创建流 | Room:%s User:%s Stream:%s", roomID, userID, streamID)
|
||
|
||
return stream, nil
|
||
}
|
||
|
||
// GetStream 获取流信息
|
||
func (r *RTMPServer) GetStream(streamID string) *RTMPStream {
|
||
r.streamsMu.RLock()
|
||
defer r.streamsMu.RUnlock()
|
||
return r.streams[streamID]
|
||
}
|
||
|
||
// GetStreamsByRoom 获取房间内所有流
|
||
func (r *RTMPServer) GetStreamsByRoom(roomID string) []*RTMPStream {
|
||
r.streamsMu.RLock()
|
||
defer r.streamsMu.RUnlock()
|
||
|
||
streams := make([]*RTMPStream, 0)
|
||
for _, s := range r.streams {
|
||
if s.RoomID == roomID && s.IsActive {
|
||
streams = append(streams, s)
|
||
}
|
||
}
|
||
return streams
|
||
}
|
||
|
||
// GetStreamsByUser 获取用户的所有流
|
||
func (r *RTMPServer) GetStreamsByUser(userID string) []*RTMPStream {
|
||
r.streamsMu.RLock()
|
||
defer r.streamsMu.RUnlock()
|
||
|
||
streams := make([]*RTMPStream, 0)
|
||
for _, s := range r.streams {
|
||
if s.UserID == userID && s.IsActive {
|
||
streams = append(streams, s)
|
||
}
|
||
}
|
||
return streams
|
||
}
|
||
|
||
// RemoveStream 移除流
|
||
func (r *RTMPServer) RemoveStream(streamID string) {
|
||
r.streamsMu.Lock()
|
||
defer r.streamsMu.Unlock()
|
||
|
||
if stream, exists := r.streams[streamID]; exists {
|
||
stream.IsActive = false
|
||
if stream.stopChan != nil {
|
||
select {
|
||
case <-stream.stopChan:
|
||
default:
|
||
close(stream.stopChan)
|
||
}
|
||
}
|
||
delete(r.streams, streamID)
|
||
log.Printf("🎬 [RTMP] 移除流 | Stream:%s", streamID)
|
||
}
|
||
|
||
r.subMu.Lock()
|
||
delete(r.subscribers, streamID)
|
||
r.subMu.Unlock()
|
||
}
|
||
|
||
// RemoveStreamsByUser 移除用户的所有流
|
||
func (r *RTMPServer) RemoveStreamsByUser(userID string) {
|
||
r.streamsMu.Lock()
|
||
defer r.streamsMu.Unlock()
|
||
|
||
for streamID, stream := range r.streams {
|
||
if stream.UserID == userID {
|
||
stream.IsActive = false
|
||
if stream.stopChan != nil {
|
||
select {
|
||
case <-stream.stopChan:
|
||
default:
|
||
close(stream.stopChan)
|
||
}
|
||
}
|
||
delete(r.streams, streamID)
|
||
log.Printf("🎬 [RTMP] 移除用户流 | User:%s Stream:%s", userID, streamID)
|
||
|
||
r.subMu.Lock()
|
||
delete(r.subscribers, streamID)
|
||
r.subMu.Unlock()
|
||
}
|
||
}
|
||
}
|
||
|
||
// RemoveStreamsByRoom 移除房间的所有流
|
||
func (r *RTMPServer) RemoveStreamsByRoom(roomID string) {
|
||
r.streamsMu.Lock()
|
||
defer r.streamsMu.Unlock()
|
||
|
||
for streamID, stream := range r.streams {
|
||
if stream.RoomID == roomID {
|
||
stream.IsActive = false
|
||
if stream.stopChan != nil {
|
||
select {
|
||
case <-stream.stopChan:
|
||
default:
|
||
close(stream.stopChan)
|
||
}
|
||
}
|
||
delete(r.streams, streamID)
|
||
log.Printf("🎬 [RTMP] 移除房间流 | Room:%s Stream:%s", roomID, streamID)
|
||
|
||
r.subMu.Lock()
|
||
delete(r.subscribers, streamID)
|
||
r.subMu.Unlock()
|
||
}
|
||
}
|
||
}
|
||
|
||
// GetStreamCount 获取流数量
|
||
func (r *RTMPServer) GetStreamCount() int {
|
||
r.streamsMu.RLock()
|
||
defer r.streamsMu.RUnlock()
|
||
return len(r.streams)
|
||
}
|
||
|
||
// cleanupTask 清理过期流
|
||
func (r *RTMPServer) cleanupTask() {
|
||
ticker := time.NewTicker(5 * time.Minute)
|
||
defer ticker.Stop()
|
||
|
||
for {
|
||
select {
|
||
case <-ticker.C:
|
||
r.cleanupExpiredStreams()
|
||
}
|
||
|
||
r.mu.RLock()
|
||
if !r.running {
|
||
r.mu.RUnlock()
|
||
return
|
||
}
|
||
r.mu.RUnlock()
|
||
}
|
||
}
|
||
|
||
// cleanupExpiredStreams 清理过期的流
|
||
func (r *RTMPServer) cleanupExpiredStreams() {
|
||
r.streamsMu.Lock()
|
||
defer r.streamsMu.Unlock()
|
||
|
||
expireTime := time.Now().Add(-1 * time.Hour)
|
||
|
||
for streamID, stream := range r.streams {
|
||
if stream.CreatedAt.Before(expireTime) && !stream.IsActive {
|
||
if stream.stopChan != nil {
|
||
select {
|
||
case <-stream.stopChan:
|
||
default:
|
||
close(stream.stopChan)
|
||
}
|
||
}
|
||
delete(r.streams, streamID)
|
||
log.Printf("🗑️ [RTMP] 清理过期流 | Stream:%s", streamID)
|
||
}
|
||
}
|
||
}
|
||
|
||
// generateStreamToken 生成流鉴权 Token
|
||
func generateStreamToken(streamID string) string {
|
||
secret := viper.GetString("turn.shared_secret")
|
||
timestamp := time.Now().Add(24 * time.Hour).Unix()
|
||
data := fmt.Sprintf("%s:%d", streamID, timestamp)
|
||
|
||
mac := hmac.New(sha1.New, []byte(secret))
|
||
mac.Write([]byte(data))
|
||
return base64.URLEncoding.EncodeToString(mac.Sum(nil))
|
||
}
|
||
|
||
// GenerateStreamToken 生成流鉴权 Token (公开接口)
|
||
func GenerateStreamToken(streamID string) string {
|
||
return generateStreamToken(streamID)
|
||
}
|
||
|
||
// ValidateStreamToken 验证流鉴权 Token
|
||
func ValidateStreamToken(streamID, token string) bool {
|
||
expectedToken := generateStreamToken(streamID)
|
||
return token == expectedToken
|
||
}
|