微信小程序\推拉流
This commit is contained in:
509
internal/mediaserver/rtmp_protocol.go
Normal file
509
internal/mediaserver/rtmp_protocol.go
Normal file
@@ -0,0 +1,509 @@
|
||||
/**
|
||||
* package mediaserver
|
||||
*
|
||||
* 自实现的 RTMP 协议处理
|
||||
* 用于替代 github.com/yutopp/go-rtmp，解决 SetChunkSize 导致的 panic 问题
|
||||
*/
|
||||
package mediaserver
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// RTMP 消息类型常量
|
||||
const (
|
||||
RTMP_MSG_CHUNK_SIZE = 1 // SetChunkSize
|
||||
RTMP_MSG_ABORT = 2 // Abort Message
|
||||
RTMP_MSG_ACK = 3 // Acknowledgement
|
||||
RTMP_MSG_USER_CONTROL = 4 // User Control Message
|
||||
RTMP_MSG_WIN_ACK_SIZE = 5 // Window Acknowledgement Size
|
||||
RTMP_MSG_SET_PEER_BW = 6 // Set Peer Bandwidth
|
||||
RTMP_MSG_AUDIO = 8 // Audio Message
|
||||
RTMP_MSG_VIDEO = 9 // Video Message
|
||||
RTMP_MSG_AMF3_DATA = 15 // AMF3 Data Message
|
||||
RTMP_MSG_AMF3_SHARED_OBJ = 16 // AMF3 Shared Object Message
|
||||
RTMP_MSG_AMF3_CMD = 17 // AMF3 Command Message
|
||||
RTMP_MSG_AMF0_DATA = 18 // AMF0 Data Message (@setDataFrame)
|
||||
RTMP_MSG_AMF0_SHARED_OBJ = 19 // AMF0 Shared Object Message
|
||||
RTMP_MSG_AMF0_CMD = 20 // AMF0 Command Message (connect, publish, etc)
|
||||
RTMP_MSG_AGGREGATE = 22 // Aggregate Message
|
||||
)
|
||||
|
||||
// Chunk 格式类型
|
||||
const (
|
||||
CHUNK_FMT_0 = 0 // 11 bytes header
|
||||
CHUNK_FMT_1 = 1 // 7 bytes header
|
||||
CHUNK_FMT_2 = 2 // 3 bytes header
|
||||
CHUNK_FMT_3 = 3 // 0 bytes header
|
||||
)
|
||||
|
||||
// 默认值
|
||||
const (
|
||||
DEFAULT_CHUNK_SIZE = 128
|
||||
MAX_CHUNK_SIZE = 65536
|
||||
DEFAULT_WINDOW_SIZE = 2500000
|
||||
RTMP_PROTOCOL_VERSION = 3
|
||||
)
|
||||
|
||||
// RTMPMessage 表示一个完整的 RTMP 消息
|
||||
type RTMPMessage struct {
|
||||
ChunkStreamID uint32
|
||||
Timestamp uint32
|
||||
TypeID uint8
|
||||
StreamID uint32
|
||||
Data []byte
|
||||
}
|
||||
|
||||
// ChunkHeader 表示 chunk 的头部信息
|
||||
type ChunkHeader struct {
|
||||
Format uint8 // 0-3
|
||||
ChunkStreamID uint32 // 2-65599
|
||||
Timestamp uint32 // 24-bit or 32-bit (extended)
|
||||
MessageLength uint32 // 24-bit
|
||||
MessageTypeID uint8
|
||||
MessageSID uint32 // 32-bit little-endian
|
||||
ExtendedTS bool // Timestamp >= 0xFFFFFF
|
||||
}
|
||||
|
||||
// ChunkReader RTMP Chunk 读取器
|
||||
type ChunkReader struct {
|
||||
conn net.Conn
|
||||
reader *bufio.Reader
|
||||
chunkSize uint32
|
||||
mu sync.RWMutex
|
||||
|
||||
// 缓存每个 chunk stream 的头部信息（用于 fmt 1/2/3）
|
||||
prevHeaders map[uint32]*ChunkHeader
|
||||
// 缓存不完整的消息数据
|
||||
messageBuffer map[uint32]*messageState
|
||||
}
|
||||
|
||||
// messageState 追踪消息的读取状态
|
||||
type messageState struct {
|
||||
header *ChunkHeader
|
||||
data []byte
|
||||
bytesRead uint32
|
||||
}
|
||||
|
||||
// ChunkWriter RTMP Chunk 写入器
|
||||
type ChunkWriter struct {
|
||||
conn net.Conn
|
||||
writer *bufio.Writer
|
||||
chunkSize uint32
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// NewChunkReader 创建新的 Chunk 读取器
|
||||
func NewChunkReader(conn net.Conn) *ChunkReader {
|
||||
return &ChunkReader{
|
||||
conn: conn,
|
||||
reader: bufio.NewReaderSize(conn, 4096),
|
||||
chunkSize: DEFAULT_CHUNK_SIZE,
|
||||
prevHeaders: make(map[uint32]*ChunkHeader),
|
||||
messageBuffer: make(map[uint32]*messageState),
|
||||
}
|
||||
}
|
||||
|
||||
// NewChunkWriter 创建新的 Chunk 写入器
|
||||
func NewChunkWriter(conn net.Conn) *ChunkWriter {
|
||||
return &ChunkWriter{
|
||||
conn: conn,
|
||||
writer: bufio.NewWriterSize(conn, 4096),
|
||||
chunkSize: DEFAULT_CHUNK_SIZE,
|
||||
}
|
||||
}
|
||||
|
||||
// SetChunkSize 设置读取的 chunk 大小
|
||||
func (r *ChunkReader) SetChunkSize(size uint32) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if size > 0 && size <= MAX_CHUNK_SIZE {
|
||||
log.Printf("📝 [RTMP Protocol] ChunkReader: 更新 chunkSize %d -> %d", r.chunkSize, size)
|
||||
r.chunkSize = size
|
||||
}
|
||||
}
|
||||
|
||||
// GetChunkSize 获取当前 chunk 大小
|
||||
func (r *ChunkReader) GetChunkSize() uint32 {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.chunkSize
|
||||
}
|
||||
|
||||
// SetChunkSize 设置写入的 chunk 大小
|
||||
func (w *ChunkWriter) SetChunkSize(size uint32) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
if size > 0 && size <= MAX_CHUNK_SIZE {
|
||||
log.Printf("📝 [RTMP Protocol] ChunkWriter: 更新 chunkSize %d -> %d", w.chunkSize, size)
|
||||
w.chunkSize = size
|
||||
}
|
||||
}
|
||||
|
||||
// ReadMessage 读取一个完整的 RTMP 消息
|
||||
func (r *ChunkReader) ReadMessage() (*RTMPMessage, error) {
|
||||
for {
|
||||
// 读取 chunk header
|
||||
header, err := r.readChunkHeader()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 获取或创建消息状态
|
||||
// 关键修复：检查消息是否已完成，如果是则为新消息创建新状态
|
||||
state, exists := r.messageBuffer[header.ChunkStreamID]
|
||||
if !exists || state.bytesRead >= state.header.MessageLength {
|
||||
// 新消息或上一条消息已完成，创建新状态
|
||||
state = &messageState{
|
||||
header: header,
|
||||
data: make([]byte, 0, header.MessageLength),
|
||||
bytesRead: 0,
|
||||
}
|
||||
r.messageBuffer[header.ChunkStreamID] = state
|
||||
}
|
||||
|
||||
// 计算本次要读取的字节数
|
||||
// 关键修复：使用 state.header.MessageLength 而非 header.MessageLength
|
||||
// 因为后续 chunk (fmt 1/2/3) 的 header 可能从前一个 header 继承值
|
||||
remaining := state.header.MessageLength - state.bytesRead
|
||||
r.mu.RLock()
|
||||
toRead := r.chunkSize
|
||||
r.mu.RUnlock()
|
||||
if remaining < toRead {
|
||||
toRead = remaining
|
||||
}
|
||||
|
||||
// toRead=0 说明 MessageLength=0，这是无效的 RTMP 消息
|
||||
// 直接返回错误而不是继续，避免失去同步
|
||||
if toRead == 0 {
|
||||
return nil, fmt.Errorf("invalid RTMP: toRead=0 on csid=%d msgLen=%d bytesRead=%d",
|
||||
header.ChunkStreamID, state.header.MessageLength, state.bytesRead)
|
||||
}
|
||||
|
||||
// 读取 chunk 数据
|
||||
chunkData := make([]byte, toRead)
|
||||
if _, err := io.ReadFull(r.reader, chunkData); err != nil {
|
||||
return nil, fmt.Errorf("读取 chunk 数据失败: %w", err)
|
||||
}
|
||||
state.data = append(state.data, chunkData...)
|
||||
state.bytesRead += toRead
|
||||
|
||||
// 检查消息是否完整
|
||||
if state.bytesRead >= state.header.MessageLength {
|
||||
msg := &RTMPMessage{
|
||||
ChunkStreamID: header.ChunkStreamID,
|
||||
Timestamp: state.header.Timestamp,
|
||||
TypeID: state.header.MessageTypeID,
|
||||
StreamID: state.header.MessageSID,
|
||||
Data: state.data,
|
||||
}
|
||||
// 清除消息缓冲
|
||||
delete(r.messageBuffer, header.ChunkStreamID)
|
||||
return msg, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// readChunkHeader 读取 chunk 头部
|
||||
func (r *ChunkReader) readChunkHeader() (*ChunkHeader, error) {
|
||||
// 读取第一个字节：Basic Header (1-3 bytes)
|
||||
firstByte, err := r.reader.ReadByte()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
format := (firstByte >> 6) & 0x03
|
||||
csid := uint32(firstByte & 0x3F)
|
||||
|
||||
// 扩展 chunk stream ID
|
||||
if csid == 0 {
|
||||
// 2 byte header
|
||||
secondByte, err := r.reader.ReadByte()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
csid = uint32(secondByte) + 64
|
||||
} else if csid == 1 {
|
||||
// 3 byte header
|
||||
bytes := make([]byte, 2)
|
||||
if _, err := io.ReadFull(r.reader, bytes); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
csid = uint32(bytes[0]) + uint32(bytes[1])*256 + 64
|
||||
}
|
||||
|
||||
// 获取上一个头部（用于 fmt 1/2/3）
|
||||
prevHeader := r.prevHeaders[csid]
|
||||
if prevHeader == nil {
|
||||
prevHeader = &ChunkHeader{
|
||||
ChunkStreamID: csid,
|
||||
}
|
||||
}
|
||||
|
||||
header := &ChunkHeader{
|
||||
Format: format,
|
||||
ChunkStreamID: csid,
|
||||
Timestamp: prevHeader.Timestamp,
|
||||
MessageLength: prevHeader.MessageLength,
|
||||
MessageTypeID: prevHeader.MessageTypeID,
|
||||
MessageSID: prevHeader.MessageSID,
|
||||
}
|
||||
|
||||
// 根据 format 读取 Message Header
|
||||
switch format {
|
||||
case CHUNK_FMT_0:
|
||||
// 11 bytes: timestamp(3) + length(3) + typeID(1) + streamID(4)
|
||||
data := make([]byte, 11)
|
||||
if _, err := io.ReadFull(r.reader, data); err != nil {
|
||||
return nil, fmt.Errorf("读取 fmt0 header 失败: %w", err)
|
||||
}
|
||||
header.Timestamp = uint32(data[0])<<16 | uint32(data[1])<<8 | uint32(data[2])
|
||||
header.MessageLength = uint32(data[3])<<16 | uint32(data[4])<<8 | uint32(data[5])
|
||||
header.MessageTypeID = data[6]
|
||||
header.MessageSID = binary.LittleEndian.Uint32(data[7:11])
|
||||
|
||||
case CHUNK_FMT_1:
|
||||
// 7 bytes: timestamp delta(3) + length(3) + typeID(1)
|
||||
data := make([]byte, 7)
|
||||
if _, err := io.ReadFull(r.reader, data); err != nil {
|
||||
return nil, fmt.Errorf("读取 fmt1 header 失败: %w", err)
|
||||
}
|
||||
delta := uint32(data[0])<<16 | uint32(data[1])<<8 | uint32(data[2])
|
||||
header.Timestamp = prevHeader.Timestamp + delta
|
||||
header.MessageLength = uint32(data[3])<<16 | uint32(data[4])<<8 | uint32(data[5])
|
||||
header.MessageTypeID = data[6]
|
||||
|
||||
case CHUNK_FMT_2:
|
||||
// 验证：fmt 2 必须有有效的 prevHeader（MessageLength 从 prevHeader 继承）
|
||||
if prevHeader.MessageLength == 0 {
|
||||
return nil, fmt.Errorf("invalid RTMP: fmt 2 on csid %d without valid prevHeader (msgLen=0)", csid)
|
||||
}
|
||||
// 3 bytes: timestamp delta(3)
|
||||
data := make([]byte, 3)
|
||||
if _, err := io.ReadFull(r.reader, data); err != nil {
|
||||
return nil, fmt.Errorf("读取 fmt2 header 失败: %w", err)
|
||||
}
|
||||
delta := uint32(data[0])<<16 | uint32(data[1])<<8 | uint32(data[2])
|
||||
header.Timestamp = prevHeader.Timestamp + delta
|
||||
|
||||
case CHUNK_FMT_3:
|
||||
// 验证：fmt 3 必须有有效的 prevHeader（所有字段从 prevHeader 继承）
|
||||
if prevHeader.MessageLength == 0 {
|
||||
return nil, fmt.Errorf("invalid RTMP: fmt 3 on csid %d without valid prevHeader (msgLen=0)", csid)
|
||||
}
|
||||
// 0 bytes: 使用上一个头部的所有字段
|
||||
// 已经复制了 prevHeader 的值
|
||||
}
|
||||
|
||||
// 检查是否有扩展时间戳
|
||||
if header.Timestamp == 0xFFFFFF {
|
||||
header.ExtendedTS = true
|
||||
extTS := make([]byte, 4)
|
||||
if _, err := io.ReadFull(r.reader, extTS); err != nil {
|
||||
return nil, fmt.Errorf("读取扩展时间戳失败: %w", err)
|
||||
}
|
||||
header.Timestamp = binary.BigEndian.Uint32(extTS)
|
||||
}
|
||||
|
||||
// 保存当前头部供后续 chunk 使用
|
||||
r.prevHeaders[csid] = header
|
||||
|
||||
return header, nil
|
||||
}
|
||||
|
||||
// WriteMessage 写入一个完整的 RTMP 消息
|
||||
func (w *ChunkWriter) WriteMessage(msg *RTMPMessage) error {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
data := msg.Data
|
||||
dataLen := uint32(len(data))
|
||||
offset := uint32(0)
|
||||
firstChunk := true
|
||||
|
||||
for offset < dataLen {
|
||||
// 计算本次写入的字节数
|
||||
remaining := dataLen - offset
|
||||
toWrite := w.chunkSize
|
||||
if remaining < toWrite {
|
||||
toWrite = remaining
|
||||
}
|
||||
|
||||
// 写入 chunk header
|
||||
if firstChunk {
|
||||
// fmt 0: 完整头部
|
||||
if err := w.writeChunkHeader0(msg, dataLen); err != nil {
|
||||
return err
|
||||
}
|
||||
firstChunk = false
|
||||
} else {
|
||||
// fmt 3: 无头部
|
||||
if err := w.writeChunkHeader3(msg.ChunkStreamID); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// 写入数据
|
||||
if _, err := w.writer.Write(data[offset : offset+toWrite]); err != nil {
|
||||
return err
|
||||
}
|
||||
offset += toWrite
|
||||
}
|
||||
|
||||
return w.writer.Flush()
|
||||
}
|
||||
|
||||
// writeChunkHeader0 写入 fmt 0 头部
|
||||
func (w *ChunkWriter) writeChunkHeader0(msg *RTMPMessage, dataLen uint32) error {
|
||||
// Basic Header
|
||||
csid := msg.ChunkStreamID
|
||||
if csid < 64 {
|
||||
if err := w.writer.WriteByte(byte(csid)); err != nil {
|
||||
return err
|
||||
}
|
||||
} else if csid < 320 {
|
||||
if _, err := w.writer.Write([]byte{0, byte(csid - 64)}); err != nil {
|
||||
return err
|
||||
}
|
||||
} else {
|
||||
csid -= 64
|
||||
if _, err := w.writer.Write([]byte{1, byte(csid & 0xFF), byte(csid >> 8)}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Message Header (11 bytes)
|
||||
header := make([]byte, 11)
|
||||
ts := msg.Timestamp
|
||||
if ts >= 0xFFFFFF {
|
||||
ts = 0xFFFFFF
|
||||
}
|
||||
header[0] = byte(ts >> 16)
|
||||
header[1] = byte(ts >> 8)
|
||||
header[2] = byte(ts)
|
||||
header[3] = byte(dataLen >> 16)
|
||||
header[4] = byte(dataLen >> 8)
|
||||
header[5] = byte(dataLen)
|
||||
header[6] = msg.TypeID
|
||||
binary.LittleEndian.PutUint32(header[7:11], msg.StreamID)
|
||||
|
||||
if _, err := w.writer.Write(header); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Extended Timestamp
|
||||
if msg.Timestamp >= 0xFFFFFF {
|
||||
extTS := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(extTS, msg.Timestamp)
|
||||
if _, err := w.writer.Write(extTS); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// writeChunkHeader3 写入 fmt 3 头部
|
||||
func (w *ChunkWriter) writeChunkHeader3(csid uint32) error {
|
||||
// Basic Header with fmt = 3
|
||||
if csid < 64 {
|
||||
return w.writer.WriteByte(byte(0xC0 | csid))
|
||||
} else if csid < 320 {
|
||||
_, err := w.writer.Write([]byte{0xC0, byte(csid - 64)})
|
||||
return err
|
||||
} else {
|
||||
csid -= 64
|
||||
_, err := w.writer.Write([]byte{0xC1, byte(csid & 0xFF), byte(csid >> 8)})
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// WriteSetChunkSize 发送 SetChunkSize 消息
|
||||
func (w *ChunkWriter) WriteSetChunkSize(size uint32) error {
|
||||
data := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(data, size)
|
||||
|
||||
msg := &RTMPMessage{
|
||||
ChunkStreamID: 2,
|
||||
Timestamp: 0,
|
||||
TypeID: RTMP_MSG_CHUNK_SIZE,
|
||||
StreamID: 0,
|
||||
Data: data,
|
||||
}
|
||||
|
||||
if err := w.WriteMessage(msg); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
w.SetChunkSize(size)
|
||||
return nil
|
||||
}
|
||||
|
||||
// WriteWindowAckSize 发送 Window Acknowledgement Size 消息
|
||||
func (w *ChunkWriter) WriteWindowAckSize(size uint32) error {
|
||||
data := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(data, size)
|
||||
|
||||
return w.WriteMessage(&RTMPMessage{
|
||||
ChunkStreamID: 2,
|
||||
Timestamp: 0,
|
||||
TypeID: RTMP_MSG_WIN_ACK_SIZE,
|
||||
StreamID: 0,
|
||||
Data: data,
|
||||
})
|
||||
}
|
||||
|
||||
// WriteSetPeerBandwidth 发送 Set Peer Bandwidth 消息
|
||||
func (w *ChunkWriter) WriteSetPeerBandwidth(size uint32, limitType uint8) error {
|
||||
data := make([]byte, 5)
|
||||
binary.BigEndian.PutUint32(data, size)
|
||||
data[4] = limitType
|
||||
|
||||
return w.WriteMessage(&RTMPMessage{
|
||||
ChunkStreamID: 2,
|
||||
Timestamp: 0,
|
||||
TypeID: RTMP_MSG_SET_PEER_BW,
|
||||
StreamID: 0,
|
||||
Data: data,
|
||||
})
|
||||
}
|
||||
|
||||
// WriteUserControl 发送 User Control Message
|
||||
func (w *ChunkWriter) WriteUserControl(eventType uint16, data []byte) error {
|
||||
payload := make([]byte, 2+len(data))
|
||||
binary.BigEndian.PutUint16(payload, eventType)
|
||||
copy(payload[2:], data)
|
||||
|
||||
return w.WriteMessage(&RTMPMessage{
|
||||
ChunkStreamID: 2,
|
||||
Timestamp: 0,
|
||||
TypeID: RTMP_MSG_USER_CONTROL,
|
||||
StreamID: 0,
|
||||
Data: payload,
|
||||
})
|
||||
}
|
||||
|
||||
// WriteStreamBegin 发送 Stream Begin 事件
|
||||
func (w *ChunkWriter) WriteStreamBegin(streamID uint32) error {
|
||||
data := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(data, streamID)
|
||||
return w.WriteUserControl(0, data) // 0 = StreamBegin
|
||||
}
|
||||
|
||||
// WriteCommand 发送 AMF0 命令消息
|
||||
func (w *ChunkWriter) WriteCommand(csid uint32, streamID uint32, data []byte) error {
|
||||
return w.WriteMessage(&RTMPMessage{
|
||||
ChunkStreamID: csid,
|
||||
Timestamp: 0,
|
||||
TypeID: RTMP_MSG_AMF0_CMD,
|
||||
StreamID: streamID,
|
||||
Data: data,
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user