fix: rtsp pull and push

This commit is contained in:
langhuihui
2024-08-13 14:44:35 +08:00
parent 43c0fc7be4
commit 78e8d74fec
18 changed files with 213 additions and 236 deletions

View File

@@ -2,6 +2,7 @@ package rtmp
import (
"errors"
"m7s.live/m7s/v5"
"net"
"runtime"
"sync/atomic"
@@ -57,6 +58,7 @@ type NetConnection struct {
chunkHeaderBuf util.Buffer
mediaDataPool util.RecyclableMemory
writing atomic.Bool // false 可写true 不可写
Receivers map[uint32]*m7s.Publisher
}
func NewNetConnection(conn net.Conn) (ret *NetConnection) {
@@ -69,30 +71,31 @@ func NewNetConnection(conn net.Conn) (ret *NetConnection) {
bandwidth: RTMP_MAX_CHUNK_SIZE << 3,
tmpBuf: make(util.Buffer, 4),
chunkHeaderBuf: make(util.Buffer, 0, 20),
Receivers: make(map[uint32]*m7s.Publisher),
}
ret.mediaDataPool.SetAllocator(util.NewScalableMemoryAllocator(1 << util.MinPowerOf2))
return
}
func (conn *NetConnection) Dispose() {
conn.Conn.Close()
conn.BufReader.Recycle()
conn.mediaDataPool.Recycle()
func (nc *NetConnection) Dispose() {
nc.Conn.Close()
nc.BufReader.Recycle()
nc.mediaDataPool.Recycle()
}
func (conn *NetConnection) SendStreamID(eventType uint16, streamID uint32) (err error) {
return conn.SendMessage(RTMP_MSG_USER_CONTROL, &StreamIDMessage{UserControlMessage{EventType: eventType}, streamID})
func (nc *NetConnection) SendStreamID(eventType uint16, streamID uint32) (err error) {
return nc.SendMessage(RTMP_MSG_USER_CONTROL, &StreamIDMessage{UserControlMessage{EventType: eventType}, streamID})
}
func (conn *NetConnection) SendUserControl(eventType uint16) error {
return conn.SendMessage(RTMP_MSG_USER_CONTROL, &UserControlMessage{EventType: eventType})
func (nc *NetConnection) SendUserControl(eventType uint16) error {
return nc.SendMessage(RTMP_MSG_USER_CONTROL, &UserControlMessage{EventType: eventType})
}
func (conn *NetConnection) ResponseCreateStream(tid uint64, streamID uint32) error {
func (nc *NetConnection) ResponseCreateStream(tid uint64, streamID uint32) error {
m := &ResponseCreateStreamMessage{}
m.CommandName = Response_Result
m.TransactionId = tid
m.StreamId = streamID
return conn.SendMessage(RTMP_MSG_AMF0_COMMAND, m)
return nc.SendMessage(RTMP_MSG_AMF0_COMMAND, m)
}
// func (conn *NetConnection) SendCommand(message string, args any) error {
@@ -110,21 +113,21 @@ func (conn *NetConnection) ResponseCreateStream(tid uint64, streamID uint32) err
// return errors.New("send message no exist")
// }
func (conn *NetConnection) readChunk() (msg *Chunk, err error) {
head, err := conn.ReadByte()
func (nc *NetConnection) readChunk() (msg *Chunk, err error) {
head, err := nc.ReadByte()
if err != nil {
return nil, err
}
conn.readSeqNum++
nc.readSeqNum++
ChunkStreamID := uint32(head & 0x3f) // 0011 1111
ChunkType := head >> 6 // 1100 0000
// 如果块流ID为0,1的话,就需要计算.
ChunkStreamID, err = conn.readChunkStreamID(ChunkStreamID)
ChunkStreamID, err = nc.readChunkStreamID(ChunkStreamID)
if err != nil {
return nil, errors.New("get chunk stream id error :" + err.Error())
}
//println("ChunkStreamID:", ChunkStreamID, "ChunkType:", ChunkType)
chunk, ok := conn.incommingChunks[ChunkStreamID]
chunk, ok := nc.incommingChunks[ChunkStreamID]
if ChunkType != 3 && ok && chunk.bufLen > 0 {
// 如果块类型不为3,那么这个rtmp的body应该为空.
@@ -132,10 +135,10 @@ func (conn *NetConnection) readChunk() (msg *Chunk, err error) {
}
if !ok {
chunk = &Chunk{}
conn.incommingChunks[ChunkStreamID] = chunk
nc.incommingChunks[ChunkStreamID] = chunk
}
if err = conn.readChunkType(&chunk.ChunkHeader, ChunkType); err != nil {
if err = nc.readChunkType(&chunk.ChunkHeader, ChunkType); err != nil {
return nil, errors.New("get chunk type error :" + err.Error())
}
msgLen := int(chunk.MessageLength)
@@ -143,19 +146,19 @@ func (conn *NetConnection) readChunk() (msg *Chunk, err error) {
return nil, nil
}
var bufSize = 0
if unRead := msgLen - chunk.bufLen; unRead < conn.ReadChunkSize {
if unRead := msgLen - chunk.bufLen; unRead < nc.ReadChunkSize {
bufSize = unRead
} else {
bufSize = conn.ReadChunkSize
bufSize = nc.ReadChunkSize
}
conn.readSeqNum += uint32(bufSize)
nc.readSeqNum += uint32(bufSize)
if chunk.bufLen == 0 {
chunk.AVData.RecyclableMemory = util.RecyclableMemory{}
chunk.AVData.SetAllocator(conn.mediaDataPool.GetAllocator())
chunk.AVData.SetAllocator(nc.mediaDataPool.GetAllocator())
chunk.AVData.NextN(msgLen)
}
buffer := chunk.AVData.Buffers[0]
err = conn.ReadRange(bufSize, func(buf []byte) {
err = nc.ReadRange(bufSize, func(buf []byte) {
copy(buffer[chunk.bufLen:], buf)
chunk.bufLen += len(buf)
})
@@ -176,14 +179,14 @@ func (conn *NetConnection) readChunk() (msg *Chunk, err error) {
return
}
func (conn *NetConnection) readChunkStreamID(csid uint32) (chunkStreamID uint32, err error) {
func (nc *NetConnection) readChunkStreamID(csid uint32) (chunkStreamID uint32, err error) {
chunkStreamID = csid
switch csid {
case 0:
{
u8, err := conn.ReadByte()
conn.readSeqNum++
u8, err := nc.ReadByte()
nc.readSeqNum++
if err != nil {
return 0, err
}
@@ -192,15 +195,15 @@ func (conn *NetConnection) readChunkStreamID(csid uint32) (chunkStreamID uint32,
}
case 1:
{
u16_0, err1 := conn.ReadByte()
u16_0, err1 := nc.ReadByte()
if err1 != nil {
return 0, err1
}
u16_1, err1 := conn.ReadByte()
u16_1, err1 := nc.ReadByte()
if err1 != nil {
return 0, err1
}
conn.readSeqNum += 2
nc.readSeqNum += 2
chunkStreamID = 64 + uint32(u16_0) + (uint32(u16_1) << 8)
}
}
@@ -208,27 +211,27 @@ func (conn *NetConnection) readChunkStreamID(csid uint32) (chunkStreamID uint32,
return chunkStreamID, nil
}
func (conn *NetConnection) readChunkType(h *ChunkHeader, chunkType byte) (err error) {
func (nc *NetConnection) readChunkType(h *ChunkHeader, chunkType byte) (err error) {
if chunkType == 3 {
// 3个字节的时间戳
} else {
// Timestamp 3 bytes
if h.Timestamp, err = conn.ReadBE32(3); err != nil {
if h.Timestamp, err = nc.ReadBE32(3); err != nil {
return err
}
if chunkType != 2 {
if h.MessageLength, err = conn.ReadBE32(3); err != nil {
if h.MessageLength, err = nc.ReadBE32(3); err != nil {
return err
}
// Message Type ID 1 bytes
if h.MessageTypeID, err = conn.ReadByte(); err != nil {
if h.MessageTypeID, err = nc.ReadByte(); err != nil {
return err
}
conn.readSeqNum++
nc.readSeqNum++
if chunkType == 0 {
// Message Stream ID 4bytes
if h.MessageStreamID, err = conn.ReadBE32(4); err != nil { // 读取Message Stream ID
if h.MessageStreamID, err = nc.ReadBE32(4); err != nil { // 读取Message Stream ID
return err
}
}
@@ -237,7 +240,7 @@ func (conn *NetConnection) readChunkType(h *ChunkHeader, chunkType byte) (err er
// ExtendTimestamp 4 bytes
if h.Timestamp >= 0xffffff { // 对于type 0的chunk,绝对时间戳在这里表示,如果时间戳值大于等于0xffffff(16777215),该值必须是0xffffff,且时间戳扩展字段必须发送,其他情况没有要求
if h.Timestamp, err = conn.ReadBE32(4); err != nil {
if h.Timestamp, err = nc.ReadBE32(4); err != nil {
return err
}
switch chunkType {
@@ -258,75 +261,90 @@ func (conn *NetConnection) readChunkType(h *ChunkHeader, chunkType byte) (err er
return nil
}
func (conn *NetConnection) RecvMessage() (msg *Chunk, err error) {
if conn.readSeqNum >= conn.bandwidth {
conn.totalRead += conn.readSeqNum
conn.readSeqNum = 0
err = conn.SendMessage(RTMP_MSG_ACK, Uint32Message(conn.totalRead))
func (nc *NetConnection) RecvMessage() (msg *Chunk, err error) {
if nc.readSeqNum >= nc.bandwidth {
nc.totalRead += nc.readSeqNum
nc.readSeqNum = 0
err = nc.SendMessage(RTMP_MSG_ACK, Uint32Message(nc.totalRead))
}
for msg == nil && err == nil {
if msg, err = conn.readChunk(); msg != nil && err == nil {
if msg, err = nc.readChunk(); msg != nil && err == nil {
switch msg.MessageTypeID {
case RTMP_MSG_CHUNK_SIZE:
conn.ReadChunkSize = int(msg.MsgData.(Uint32Message))
nc.ReadChunkSize = int(msg.MsgData.(Uint32Message))
nc.Info("msg read chunk size", "readChunkSize", nc.ReadChunkSize)
case RTMP_MSG_ABORT:
delete(conn.incommingChunks, uint32(msg.MsgData.(Uint32Message)))
delete(nc.incommingChunks, uint32(msg.MsgData.(Uint32Message)))
case RTMP_MSG_ACK, RTMP_MSG_EDGE:
case RTMP_MSG_USER_CONTROL:
if _, ok := msg.MsgData.(*PingRequestMessage); ok {
conn.SendUserControl(RTMP_USER_PING_RESPONSE)
nc.SendUserControl(RTMP_USER_PING_RESPONSE)
}
case RTMP_MSG_ACK_SIZE:
conn.bandwidth = uint32(msg.MsgData.(Uint32Message))
nc.bandwidth = uint32(msg.MsgData.(Uint32Message))
case RTMP_MSG_BANDWIDTH:
conn.bandwidth = msg.MsgData.(*SetPeerBandwidthMessage).AcknowledgementWindowsize
case RTMP_MSG_AMF0_COMMAND, RTMP_MSG_AUDIO, RTMP_MSG_VIDEO:
nc.bandwidth = msg.MsgData.(*SetPeerBandwidthMessage).AcknowledgementWindowsize
case RTMP_MSG_AMF0_COMMAND:
return msg, err
case RTMP_MSG_AUDIO:
if r, ok := nc.Receivers[msg.MessageStreamID]; ok && r.PubAudio {
err = r.WriteAudio(msg.AVData.WrapAudio())
} else {
msg.AVData.Recycle()
nc.Warn("ReceiveAudio", "MessageStreamID", msg.MessageStreamID)
}
case RTMP_MSG_VIDEO:
if r, ok := nc.Receivers[msg.MessageStreamID]; ok && r.PubVideo {
err = r.WriteVideo(msg.AVData.WrapVideo())
} else {
msg.AVData.Recycle()
nc.Warn("ReceiveVideo", "MessageStreamID", msg.MessageStreamID)
}
}
}
}
return
}
func (conn *NetConnection) SendMessage(t byte, msg RtmpMessage) (err error) {
if conn == nil {
func (nc *NetConnection) SendMessage(t byte, msg RtmpMessage) (err error) {
if nc == nil {
return errors.New("connection is nil")
}
if conn.writeSeqNum > conn.bandwidth {
conn.totalWrite += conn.writeSeqNum
conn.writeSeqNum = 0
err = conn.SendMessage(RTMP_MSG_ACK, Uint32Message(conn.totalWrite))
err = conn.SendStreamID(RTMP_USER_PING_REQUEST, 0)
if nc.writeSeqNum > nc.bandwidth {
nc.totalWrite += nc.writeSeqNum
nc.writeSeqNum = 0
err = nc.SendMessage(RTMP_MSG_ACK, Uint32Message(nc.totalWrite))
err = nc.SendStreamID(RTMP_USER_PING_REQUEST, 0)
}
for !conn.writing.CompareAndSwap(false, true) {
for !nc.writing.CompareAndSwap(false, true) {
runtime.Gosched()
}
defer conn.writing.Store(false)
conn.tmpBuf.Reset()
amf := AMF{conn.tmpBuf}
if conn.ObjectEncoding == 0 {
defer nc.writing.Store(false)
nc.tmpBuf.Reset()
amf := AMF{nc.tmpBuf}
if nc.ObjectEncoding == 0 {
msg.Encode(&amf)
} else {
amf := AMF3{AMF: amf}
msg.Encode(&amf)
}
conn.tmpBuf = amf.Buffer
nc.tmpBuf = amf.Buffer
head := newChunkHeader(t)
head.MessageLength = uint32(conn.tmpBuf.Len())
head.MessageLength = uint32(nc.tmpBuf.Len())
if sid, ok := msg.(HaveStreamID); ok {
head.MessageStreamID = sid.GetStreamID()
}
return conn.sendChunk(net.Buffers{conn.tmpBuf}, head, RTMP_CHUNK_HEAD_12)
return nc.sendChunk(net.Buffers{nc.tmpBuf}, head, RTMP_CHUNK_HEAD_12)
}
func (conn *NetConnection) sendChunk(data net.Buffers, head *ChunkHeader, headType byte) (err error) {
conn.chunkHeaderBuf.Reset()
head.WriteTo(headType, &conn.chunkHeaderBuf)
chunks := net.Buffers{conn.chunkHeaderBuf}
var chunk3 util.Buffer = conn.chunkHeaderBuf[conn.chunkHeaderBuf.Len():20]
func (nc *NetConnection) sendChunk(data net.Buffers, head *ChunkHeader, headType byte) (err error) {
nc.chunkHeaderBuf.Reset()
head.WriteTo(headType, &nc.chunkHeaderBuf)
chunks := net.Buffers{nc.chunkHeaderBuf}
var chunk3 util.Buffer = nc.chunkHeaderBuf[nc.chunkHeaderBuf.Len():20]
head.WriteTo(RTMP_CHUNK_HEAD_1, &chunk3)
r := util.NewReadableBuffersFromBytes(data...)
for {
r.RangeN(conn.WriteChunkSize, func(buf []byte) {
r.RangeN(nc.WriteChunkSize, func(buf []byte) {
chunks = append(chunks, buf)
})
if r.Length <= 0 {
@@ -336,7 +354,7 @@ func (conn *NetConnection) sendChunk(data net.Buffers, head *ChunkHeader, headTy
chunks = append(chunks, chunk3)
}
var nw int64
nw, err = chunks.WriteTo(conn.Conn)
conn.writeSeqNum += uint32(nw)
nw, err = chunks.WriteTo(nc.Conn)
nc.writeSeqNum += uint32(nw)
return err
}