/** * 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) { 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), FLVURL: fmt.Sprintf("http://%s:%d/live/%s.flv", publicIP, r.config.HTTPFLVPort, 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 }