pimhwang
/
mediamtx
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
Browse Source

rtmp: add MessageWriter

pull/956/head
aler9 4 years ago
parent
046a3f248d
commit
9519bf670a
6 changed files with 207 additions and 132 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 52
      internal/rtmp/base/chunk0.go
  2. 16
      internal/rtmp/base/chunk1.go
  3. 6
      internal/rtmp/base/chunk3.go
  4. 10
      internal/rtmp/base/message.go
  5. 83
      internal/rtmp/base/messagewriter.go
  6. 172
      internal/rtmp/conn_test.go

52
internal/rtmp/base/chunk0.go
Unescape View File

@ -10,16 +10,16 @@ import (
// the start of a chunk stream, and whenever the stream timestamp goes // the start of a chunk stream, and whenever the stream timestamp goes
// backward (e.g., because of a backward seek). // backward (e.g., because of a backward seek).
type Chunk0 struct { type Chunk0 struct {
ChunkStreamID byte ChunkStreamID byte
Timestamp uint32 Timestamp uint32
Typ byte Typ byte
StreamID uint32 MessageStreamID uint32
BodyLen uint32 BodyLen uint32
Body []byte Body []byte
} }
// Read reads the chunk. // Read reads the chunk.
func (m *Chunk0) Read(r io.Reader, chunkMaxBodyLen int) error { func (c *Chunk0) Read(r io.Reader, chunkMaxBodyLen int) error {
header := make([]byte, 12) header := make([]byte, 12)
_, err := r.Read(header) _, err := r.Read(header)
if err != nil { if err != nil {
@ -30,39 +30,39 @@ func (m *Chunk0) Read(r io.Reader, chunkMaxBodyLen int) error {
return fmt.Errorf("wrong chunk header type") return fmt.Errorf("wrong chunk header type")
} }
m.ChunkStreamID = header[0] & 0x3F c.ChunkStreamID = header[0] & 0x3F
m.Timestamp = uint32(header[3])<<16 | uint32(header[2])<<8 | uint32(header[1]) c.Timestamp = uint32(header[3])<<16 | uint32(header[2])<<8 | uint32(header[1])
m.BodyLen = uint32(header[4])<<16 | uint32(header[5])<<8 | uint32(header[6]) c.BodyLen = uint32(header[4])<<16 | uint32(header[5])<<8 | uint32(header[6])
m.Typ = header[7] c.Typ = header[7]
m.StreamID = uint32(header[8])<<24 | uint32(header[9])<<16 | uint32(header[10])<<8 | uint32(header[11]) c.MessageStreamID = uint32(header[8])<<24 | uint32(header[9])<<16 | uint32(header[10])<<8 | uint32(header[11])
chunkBodyLen := int(m.BodyLen) chunkBodyLen := int(c.BodyLen)
if chunkBodyLen > chunkMaxBodyLen { if chunkBodyLen > chunkMaxBodyLen {
chunkBodyLen = chunkMaxBodyLen chunkBodyLen = chunkMaxBodyLen
} }
m.Body = make([]byte, chunkBodyLen) c.Body = make([]byte, chunkBodyLen)
_, err = r.Read(m.Body) _, err = r.Read(c.Body)
return err return err
} }
// Write writes the chunk. // Write writes the chunk.
func (m Chunk0) Write(w io.Writer) error { func (c Chunk0) Write(w io.Writer) error {
header := make([]byte, 12) header := make([]byte, 12)
header[0] = m.ChunkStreamID header[0] = c.ChunkStreamID
header[1] = byte(m.Timestamp >> 16) header[1] = byte(c.Timestamp >> 16)
header[2] = byte(m.Timestamp >> 8) header[2] = byte(c.Timestamp >> 8)
header[3] = byte(m.Timestamp) header[3] = byte(c.Timestamp)
header[4] = byte(m.BodyLen >> 16) header[4] = byte(c.BodyLen >> 16)
header[5] = byte(m.BodyLen >> 8) header[5] = byte(c.BodyLen >> 8)
header[6] = byte(m.BodyLen) header[6] = byte(c.BodyLen)
header[7] = m.Typ header[7] = c.Typ
header[8] = byte(m.StreamID) header[8] = byte(c.MessageStreamID)
_, err := w.Write(header) _, err := w.Write(header)
if err != nil { if err != nil {
return err return err
} }
_, err = w.Write(m.Body) _, err = w.Write(c.Body)
return err return err
} }

16
internal/rtmp/base/chunk1.go
Unescape View File

@ -13,23 +13,23 @@ import (
type Chunk1 struct { type Chunk1 struct {
ChunkStreamID byte ChunkStreamID byte
Typ byte Typ byte
BodyLen uint32
Body []byte Body []byte
} }
// Write writes the chunk. // Write writes the chunk.
func (m Chunk1) Write(w io.Writer) error { func (c Chunk1) Write(w io.Writer) error {
header := make([]byte, 8) header := make([]byte, 8)
header[0] = 1<<6 | m.ChunkStreamID header[0] = 1<<6 | c.ChunkStreamID
l := uint32(len(m.Body)) header[4] = byte(c.BodyLen >> 16)
header[4] = byte(l >> 16) header[5] = byte(c.BodyLen >> 8)
header[5] = byte(l >> 8) header[6] = byte(c.BodyLen)
header[6] = byte(l) header[7] = c.Typ
header[7] = m.Typ
_, err := w.Write(header) _, err := w.Write(header)
if err != nil { if err != nil {
return err return err
} }
_, err = w.Write(m.Body) _, err = w.Write(c.Body)
return err return err
} }

6
internal/rtmp/base/chunk3.go
Unescape View File

@ -16,14 +16,14 @@ type Chunk3 struct {
} }
// Write writes the chunk. // Write writes the chunk.
func (m Chunk3) Write(w io.Writer) error { func (c Chunk3) Write(w io.Writer) error {
header := make([]byte, 1) header := make([]byte, 1)
header[0] = 3<<6 | m.ChunkStreamID header[0] = 3<<6 | c.ChunkStreamID
_, err := w.Write(header) _, err := w.Write(header)
if err != nil { if err != nil {
return err return err
} }
_, err = w.Write(m.Body) _, err = w.Write(c.Body)
return err return err
} }

10
internal/rtmp/base/message.go
Unescape View File

@ -0,0 +1,10 @@
package base
// Message is a message.
type Message struct {
ChunkStreamID byte
Timestamp uint32
Typ byte
MessageStreamID uint32
Body []byte
}

83
internal/rtmp/base/messagewriter.go
Unescape View File

@ -0,0 +1,83 @@
package base
import (
"io"
)
// MessageWriter is a message writer.
type MessageWriter struct {
w io.Writer
chunkMaxBodyLen int
lastMessageStreamIDPerChunkStreamID map[byte]uint32
}
// NewMessageWriter instantiates a MessageWriter.
func NewMessageWriter(w io.Writer) *MessageWriter {
return &MessageWriter{
w: w,
chunkMaxBodyLen: 128,
lastMessageStreamIDPerChunkStreamID: make(map[byte]uint32),
}
}
// SetChunkSize sets the chunk size.
func (mw *MessageWriter) SetChunkSize(v int) {
mw.chunkMaxBodyLen = v
}
// Write writes a Message.
func (mw *MessageWriter) Write(msg *Message) error {
bodyLen := len(msg.Body)
pos := 0
first := true
for {
chunkBodyLen := bodyLen - pos
if chunkBodyLen > mw.chunkMaxBodyLen {
chunkBodyLen = mw.chunkMaxBodyLen
}
if first {
first = false
if v, ok := mw.lastMessageStreamIDPerChunkStreamID[msg.ChunkStreamID]; !ok || v != msg.MessageStreamID {
err := Chunk0{
ChunkStreamID: msg.ChunkStreamID,
Typ: msg.Typ,
MessageStreamID: msg.MessageStreamID,
BodyLen: uint32(bodyLen),
Body: msg.Body[pos : pos+chunkBodyLen],
}.Write(mw.w)
if err != nil {
return err
}
mw.lastMessageStreamIDPerChunkStreamID[msg.ChunkStreamID] = msg.MessageStreamID
} else {
err := Chunk1{
ChunkStreamID: msg.ChunkStreamID,
Typ: msg.Typ,
BodyLen: uint32(bodyLen),
Body: msg.Body[pos : pos+chunkBodyLen],
}.Write(mw.w)
if err != nil {
return err
}
}
} else {
err := Chunk3{
ChunkStreamID: msg.ChunkStreamID,
Body: msg.Body[pos : pos+chunkBodyLen],
}.Write(mw.w)
if err != nil {
return err
}
}
pos += chunkBodyLen
if (bodyLen - pos) == 0 {
return nil
}
}
}

172
internal/rtmp/conn_test.go
Unescape View File

@ -134,6 +134,8 @@ func TestReadTracks(t *testing.T) {
err = base.HandshakeC2{}.Write(conn, s1s2) err = base.HandshakeC2{}.Write(conn, s1s2)
require.NoError(t, err) require.NoError(t, err)
mw := base.NewMessageWriter(conn)
// C->S connect // C->S connect
byts := flvio.FillAMF0ValsMalloc([]interface{}{ byts := flvio.FillAMF0ValsMalloc([]interface{}{
"connect", "connect",
@ -149,17 +151,11 @@ func TestReadTracks(t *testing.T) {
{K: "videoFunction", V: 1}, {K: "videoFunction", V: 1},
}, },
}) })
err = base.Chunk0{ err = mw.Write(&base.Message{
ChunkStreamID: 3, ChunkStreamID: 3,
Typ: 0x14, Typ: 0x14,
BodyLen: uint32(len(byts)), Body: byts,
Body: byts[:128], })
}.Write(conn)
require.NoError(t, err)
err = base.Chunk3{
ChunkStreamID: 3,
Body: byts[128:],
}.Write(conn)
require.NoError(t, err) require.NoError(t, err)
// S->C window acknowledgement size // S->C window acknowledgement size
@ -216,16 +212,17 @@ func TestReadTracks(t *testing.T) {
}, arr) }, arr)
// C->S set chunk size // C->S set chunk size
err = base.Chunk0{ err = mw.Write(&base.Message{
ChunkStreamID: 2, ChunkStreamID: 2,
Typ: 1, Typ: 1,
BodyLen: 4,
Body: []byte{0x00, 0x01, 0x00, 0x00}, Body: []byte{0x00, 0x01, 0x00, 0x00},
}.Write(conn) })
require.NoError(t, err) require.NoError(t, err)
mw.SetChunkSize(65536)
// C->S releaseStream // C->S releaseStream
err = base.Chunk1{ err = mw.Write(&base.Message{
ChunkStreamID: 3, ChunkStreamID: 3,
Typ: 0x14, Typ: 0x14,
Body: flvio.FillAMF0ValsMalloc([]interface{}{ Body: flvio.FillAMF0ValsMalloc([]interface{}{
@ -234,11 +231,11 @@ func TestReadTracks(t *testing.T) {
nil, nil,
"", "",
}), }),
}.Write(conn) })
require.NoError(t, err) require.NoError(t, err)
// C->S FCPublish // C->S FCPublish
err = base.Chunk1{ err = mw.Write(&base.Message{
ChunkStreamID: 3, ChunkStreamID: 3,
Typ: 0x14, Typ: 0x14,
Body: flvio.FillAMF0ValsMalloc([]interface{}{ Body: flvio.FillAMF0ValsMalloc([]interface{}{
@ -247,18 +244,19 @@ func TestReadTracks(t *testing.T) {
nil, nil,
"", "",
}), }),
}.Write(conn) })
require.NoError(t, err) require.NoError(t, err)
// C->S createStream // C->S createStream
err = base.Chunk3{ err = mw.Write(&base.Message{
ChunkStreamID: 3, ChunkStreamID: 3,
Typ: 0x14,
Body: flvio.FillAMF0ValsMalloc([]interface{}{ Body: flvio.FillAMF0ValsMalloc([]interface{}{
"createStream", "createStream",
float64(4), float64(4),
nil, nil,
}), }),
}.Write(conn) })
require.NoError(t, err) require.NoError(t, err)
// S->C result // S->C result
@ -276,20 +274,18 @@ func TestReadTracks(t *testing.T) {
}, arr) }, arr)
// C->S publish // C->S publish
byts = flvio.FillAMF0ValsMalloc([]interface{}{ err = mw.Write(&base.Message{
"publish", ChunkStreamID: 8,
float64(5), Typ: 0x14,
nil, MessageStreamID: 1,
"", Body: flvio.FillAMF0ValsMalloc([]interface{}{
"live", "publish",
float64(5),
nil,
"",
"live",
}),
}) })
err = base.Chunk0{
ChunkStreamID: 8,
Typ: 0x14,
StreamID: 1,
BodyLen: uint32(len(byts)),
Body: byts,
}.Write(conn)
require.NoError(t, err) require.NoError(t, err)
// S->C onStatus // S->C onStatus
@ -335,13 +331,12 @@ func TestReadTracks(t *testing.T) {
}, },
}, },
}) })
err = base.Chunk0{ err = mw.Write(&base.Message{
ChunkStreamID: 4, ChunkStreamID: 4,
Typ: 0x12, Typ: 0x12,
StreamID: 1, MessageStreamID: 1,
BodyLen: uint32(len(byts)), Body: byts,
Body: byts, })
}.Write(conn)
require.NoError(t, err) require.NoError(t, err)
// C->S H264 decoder config // C->S H264 decoder config
@ -357,13 +352,12 @@ func TestReadTracks(t *testing.T) {
var n int var n int
codec.ToConfig(b, &n) codec.ToConfig(b, &n)
body := append([]byte{flvio.FRAME_KEY<<4 | flvio.VIDEO_H264, 0, 0, 0, 0}, b[:n]...) body := append([]byte{flvio.FRAME_KEY<<4 | flvio.VIDEO_H264, 0, 0, 0, 0}, b[:n]...)
err = base.Chunk0{ err = mw.Write(&base.Message{
ChunkStreamID: 6, ChunkStreamID: 6,
Typ: flvio.TAG_VIDEO, Typ: flvio.TAG_VIDEO,
StreamID: 1, MessageStreamID: 1,
BodyLen: uint32(len(body)), Body: body,
Body: body, })
}.Write(conn)
require.NoError(t, err) require.NoError(t, err)
// C->S AAC decoder config // C->S AAC decoder config
@ -373,16 +367,15 @@ func TestReadTracks(t *testing.T) {
ChannelCount: 2, ChannelCount: 2,
}.Encode() }.Encode()
require.NoError(t, err) require.NoError(t, err)
err = base.Chunk0{ err = mw.Write(&base.Message{
ChunkStreamID: 4, ChunkStreamID: 4,
Typ: flvio.TAG_AUDIO, Typ: flvio.TAG_AUDIO,
StreamID: 1, MessageStreamID: 1,
BodyLen: uint32(len(enc) + 2),
Body: append([]byte{ Body: append([]byte{
flvio.SOUND_AAC<<4 | flvio.SOUND_44Khz<<2 | flvio.SOUND_16BIT<<1 | flvio.SOUND_STEREO, flvio.SOUND_AAC<<4 | flvio.SOUND_44Khz<<2 | flvio.SOUND_16BIT<<1 | flvio.SOUND_STEREO,
flvio.AAC_SEQHDR, flvio.AAC_SEQHDR,
}, enc...), }, enc...),
}.Write(conn) })
require.NoError(t, err) require.NoError(t, err)
case "metadata without codec id": case "metadata without codec id":
@ -405,13 +398,12 @@ func TestReadTracks(t *testing.T) {
}, },
}, },
}) })
err = base.Chunk0{ err = mw.Write(&base.Message{
ChunkStreamID: 4, ChunkStreamID: 4,
Typ: 0x12, Typ: 0x12,
StreamID: 1, MessageStreamID: 1,
BodyLen: uint32(len(byts)), Body: byts,
Body: byts, })
}.Write(conn)
require.NoError(t, err) require.NoError(t, err)
// C->S H264 decoder config // C->S H264 decoder config
@ -427,13 +419,12 @@ func TestReadTracks(t *testing.T) {
var n int var n int
codec.ToConfig(b, &n) codec.ToConfig(b, &n)
body := append([]byte{flvio.FRAME_KEY<<4 | flvio.VIDEO_H264, 0, 0, 0, 0}, b[:n]...) body := append([]byte{flvio.FRAME_KEY<<4 | flvio.VIDEO_H264, 0, 0, 0, 0}, b[:n]...)
err = base.Chunk0{ err = mw.Write(&base.Message{
ChunkStreamID: 6, ChunkStreamID: 6,
Typ: flvio.TAG_VIDEO, Typ: flvio.TAG_VIDEO,
StreamID: 1, MessageStreamID: 1,
BodyLen: uint32(len(body)), Body: body,
Body: body, })
}.Write(conn)
require.NoError(t, err) require.NoError(t, err)
case "no metadata": case "no metadata":
@ -450,13 +441,12 @@ func TestReadTracks(t *testing.T) {
var n int var n int
codec.ToConfig(b, &n) codec.ToConfig(b, &n)
body := append([]byte{flvio.FRAME_KEY<<4 | flvio.VIDEO_H264, 0, 0, 0, 0}, b[:n]...) body := append([]byte{flvio.FRAME_KEY<<4 | flvio.VIDEO_H264, 0, 0, 0, 0}, b[:n]...)
err = base.Chunk0{ err = mw.Write(&base.Message{
ChunkStreamID: 6, ChunkStreamID: 6,
Typ: flvio.TAG_VIDEO, Typ: flvio.TAG_VIDEO,
StreamID: 1, MessageStreamID: 1,
BodyLen: uint32(len(body)), Body: body,
Body: body, })
}.Write(conn)
require.NoError(t, err) require.NoError(t, err)
} }
@ -523,6 +513,8 @@ func TestWriteTracks(t *testing.T) {
err = base.HandshakeC2{}.Write(conn, s1s2) err = base.HandshakeC2{}.Write(conn, s1s2)
require.NoError(t, err) require.NoError(t, err)
mw := base.NewMessageWriter(conn)
// C->S connect // C->S connect
byts := flvio.FillAMF0ValsMalloc([]interface{}{ byts := flvio.FillAMF0ValsMalloc([]interface{}{
"connect", "connect",
@ -538,17 +530,11 @@ func TestWriteTracks(t *testing.T) {
{K: "videoFunction", V: 1}, {K: "videoFunction", V: 1},
}, },
}) })
err = base.Chunk0{ err = mw.Write(&base.Message{
ChunkStreamID: 3, ChunkStreamID: 3,
Typ: 0x14, Typ: 0x14,
BodyLen: uint32(len(byts)), Body: byts,
Body: byts[:128], })
}.Write(conn)
require.NoError(t, err)
err = base.Chunk3{
ChunkStreamID: 3,
Body: byts[128:],
}.Write(conn)
require.NoError(t, err) require.NoError(t, err)
// S->C window acknowledgement size // S->C window acknowledgement size
@ -605,25 +591,23 @@ func TestWriteTracks(t *testing.T) {
}, arr) }, arr)
// C->S window acknowledgement size // C->S window acknowledgement size
err = base.Chunk0{ err = mw.Write(&base.Message{
ChunkStreamID: 2, ChunkStreamID: 2,
Typ: 0x05, Typ: 0x05,
BodyLen: 4,
Body: []byte{0x00, 0x26, 0x25, 0xa0}, Body: []byte{0x00, 0x26, 0x25, 0xa0},
}.Write(conn) })
require.NoError(t, err) require.NoError(t, err)
// C->S set chunk size // C->S set chunk size
err = base.Chunk0{ err = mw.Write(&base.Message{
ChunkStreamID: 2, ChunkStreamID: 2,
Typ: 1, Typ: 1,
BodyLen: 4,
Body: []byte{0x00, 0x01, 0x00, 0x00}, Body: []byte{0x00, 0x01, 0x00, 0x00},
}.Write(conn) })
require.NoError(t, err) require.NoError(t, err)
// C->S createStream // C->S createStream
err = base.Chunk1{ err = mw.Write(&base.Message{
ChunkStreamID: 3, ChunkStreamID: 3,
Typ: 0x14, Typ: 0x14,
Body: flvio.FillAMF0ValsMalloc([]interface{}{ Body: flvio.FillAMF0ValsMalloc([]interface{}{
@ -631,7 +615,7 @@ func TestWriteTracks(t *testing.T) {
float64(2), float64(2),
nil, nil,
}), }),
}.Write(conn) })
require.NoError(t, err) require.NoError(t, err)
// S->C result // S->C result
@ -655,11 +639,10 @@ func TestWriteTracks(t *testing.T) {
nil, nil,
"", "",
}) })
err = base.Chunk0{ err = mw.Write(&base.Message{
ChunkStreamID: 8, ChunkStreamID: 8,
BodyLen: uint32(len(byts)),
Body: byts, Body: byts,
}.Write(conn) })
require.NoError(t, err) require.NoError(t, err)
// C->S play // C->S play
@ -670,12 +653,11 @@ func TestWriteTracks(t *testing.T) {
"", "",
float64(-2000), float64(-2000),
}) })
err = base.Chunk0{ err = mw.Write(&base.Message{
ChunkStreamID: 8, ChunkStreamID: 8,
Typ: 0x14, Typ: 0x14,
BodyLen: uint32(len(byts)),
Body: byts, Body: byts,
}.Write(conn) })
require.NoError(t, err) require.NoError(t, err)
// S->C event "stream is recorded" // S->C event "stream is recorded"

Write Preview
Loading…
Cancel
Save