zishuo/rtsp-simple-server
1
0
Fork 0
mirror of https://github.com/aler9/rtsp-simple-server synced 2026-04-22 23:17:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

rtmp: prevent legacy clients from reading multiple video/audio tracks (#5478)

Browse Source
This commit is contained in:
Alessandro Ros
2026-02-18 11:02:58 +01:00
committed by GitHub
parent 2c0defbb07
commit ddaafd4ef3
2 changed files with 294 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/protocols/rtmp/from_stream.go
+133 -116
View File
@@ -48,6 +48,10 @@ func FromStream(
var tracks []*gortmplib.Track
var w *gortmplib.Writer
isEnhanced := len(conn.FourCcList) != 0
legacyVideoTrackCount := 0
legacyAudioTrackCount := 0
for _, media := range desc.Medias {
for _, forma := range media.Formats {
switch forma := forma.(type) {
@@ -142,63 +146,66 @@ func FromStream(
}
case *format.H264:
sps, pps := forma.SafeParams()
track := &gortmplib.Track{
Codec: &codecs.H264{
SPS: sps,
PPS: pps,
},
}
tracks = append(tracks, track)
if isEnhanced || legacyVideoTrackCount == 0 {
legacyVideoTrackCount++
sps, pps := forma.SafeParams()
track := &gortmplib.Track{
Codec: &codecs.H264{
SPS: sps,
PPS: pps,
},
}
tracks = append(tracks, track)
var videoDTSExtractor *h264.DTSExtractor
var videoDTSExtractor *h264.DTSExtractor
r.OnData(
media,
forma,
func(u *unit.Unit) error {
if u.NilPayload() {
return nil
}
idrPresent := false
nonIDRPresent := false
for _, nalu := range u.Payload.(unit.PayloadH264) {
typ := h264.NALUType(nalu[0] & 0x1F)
switch typ {
case h264.NALUTypeIDR:
idrPresent = true
case h264.NALUTypeNonIDR:
nonIDRPresent = true
}
}
// wait until we receive an IDR
if videoDTSExtractor == nil {
if !idrPresent {
r.OnData(
media,
forma,
func(u *unit.Unit) error {
if u.NilPayload() {
return nil
}
videoDTSExtractor = &h264.DTSExtractor{}
videoDTSExtractor.Initialize()
} else if !idrPresent && !nonIDRPresent {
return nil
}
idrPresent := false
nonIDRPresent := false
dts, err := videoDTSExtractor.Extract(u.Payload.(unit.PayloadH264), u.PTS)
if err != nil {
return err
}
for _, nalu := range u.Payload.(unit.PayloadH264) {
typ := h264.NALUType(nalu[0] & 0x1F)
switch typ {
case h264.NALUTypeIDR:
idrPresent = true
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
return (*w).WriteH264(
track,
timestampToDuration(u.PTS, forma.ClockRate()),
timestampToDuration(dts, forma.ClockRate()),
u.Payload.(unit.PayloadH264))
})
case h264.NALUTypeNonIDR:
nonIDRPresent = true
}
}
// wait until we receive an IDR
if videoDTSExtractor == nil {
if !idrPresent {
return nil
}
videoDTSExtractor = &h264.DTSExtractor{}
videoDTSExtractor.Initialize()
} else if !idrPresent && !nonIDRPresent {
return nil
}
dts, err := videoDTSExtractor.Extract(u.Payload.(unit.PayloadH264), u.PTS)
if err != nil {
return err
}
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
return (*w).WriteH264(
track,
timestampToDuration(u.PTS, forma.ClockRate()),
timestampToDuration(dts, forma.ClockRate()),
u.Payload.(unit.PayloadH264))
})
}
case *format.Opus:
if slices.Contains(conn.FourCcList, any(fourCCToString(message.FourCCOpus))) {
@@ -238,40 +245,44 @@ func FromStream(
}
case *format.MPEG4Audio:
track := &gortmplib.Track{
Codec: &codecs.MPEG4Audio{
Config: forma.Config,
},
}
tracks = append(tracks, track)
if isEnhanced || legacyAudioTrackCount == 0 {
legacyAudioTrackCount++
track := &gortmplib.Track{
Codec: &codecs.MPEG4Audio{
Config: forma.Config,
},
}
tracks = append(tracks, track)
r.OnData(
media,
forma,
func(u *unit.Unit) error {
if u.NilPayload() {
return nil
}
for i, au := range u.Payload.(unit.PayloadMPEG4Audio) {
pts := u.PTS + int64(i)*mpeg4audio.SamplesPerAccessUnit
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
err := (*w).WriteMPEG4Audio(
track,
timestampToDuration(pts, forma.ClockRate()),
au,
)
if err != nil {
return err
r.OnData(
media,
forma,
func(u *unit.Unit) error {
if u.NilPayload() {
return nil
}
}
return nil
})
for i, au := range u.Payload.(unit.PayloadMPEG4Audio) {
pts := u.PTS + int64(i)*mpeg4audio.SamplesPerAccessUnit
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
err := (*w).WriteMPEG4Audio(
track,
timestampToDuration(pts, forma.ClockRate()),
au,
)
if err != nil {
return err
}
}
return nil
})
}
case *format.MPEG4AudioLATM:
if !forma.CPresent {
if !forma.CPresent && (isEnhanced || legacyAudioTrackCount == 0) {
legacyAudioTrackCount++
track := &gortmplib.Track{
Codec: &codecs.MPEG4Audio{
Config: forma.StreamMuxConfig.Programs[0].Layers[0].AudioSpecificConfig,
@@ -304,43 +315,46 @@ func FromStream(
}
case *format.MPEG1Audio:
track := &gortmplib.Track{
Codec: &codecs.MPEG1Audio{},
}
tracks = append(tracks, track)
if isEnhanced || legacyAudioTrackCount == 0 {
legacyAudioTrackCount++
track := &gortmplib.Track{
Codec: &codecs.MPEG1Audio{},
}
tracks = append(tracks, track)
r.OnData(
media,
forma,
func(u *unit.Unit) error {
if u.NilPayload() {
return nil
}
pts := u.PTS
for _, frame := range u.Payload.(unit.PayloadMPEG1Audio) {
var h mpeg1audio.FrameHeader
err := h.Unmarshal(frame)
if err != nil {
return err
}
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
err = (*w).WriteMPEG1Audio(
track,
timestampToDuration(pts, forma.ClockRate()),
frame)
if err != nil {
return err
}
pts += int64(h.SampleCount()) *
int64(forma.ClockRate()) / int64(h.SampleRate)
}
r.OnData(
media,
forma,
func(u *unit.Unit) error {
if u.NilPayload() {
return nil
}
pts := u.PTS
for _, frame := range u.Payload.(unit.PayloadMPEG1Audio) {
var h mpeg1audio.FrameHeader
err := h.Unmarshal(frame)
if err != nil {
return err
}
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
err = (*w).WriteMPEG1Audio(
track,
timestampToDuration(pts, forma.ClockRate()),
frame)
if err != nil {
return err
}
pts += int64(h.SampleCount()) *
int64(forma.ClockRate()) / int64(h.SampleRate)
}
return nil
})
})
}
case *format.AC3:
if slices.Contains(conn.FourCcList, any(fourCCToString(message.FourCCAC3))) {
@@ -378,7 +392,8 @@ func FromStream(
}
case *format.G711:
if forma.SampleRate == 8000 {
if forma.SampleRate == 8000 && (isEnhanced || legacyAudioTrackCount == 0) {
legacyAudioTrackCount++
track := &gortmplib.Track{
Codec: &codecs.G711{
MULaw: forma.MULaw,
@@ -409,7 +424,9 @@ func FromStream(
(forma.SampleRate == 5512 ||
forma.SampleRate == 11025 ||
forma.SampleRate == 22050 ||
forma.SampleRate == 44100) {
forma.SampleRate == 44100) &&
(isEnhanced || legacyAudioTrackCount == 0) {
legacyAudioTrackCount++
track := &gortmplib.Track{
Codec: &codecs.LPCM{
BitDepth: forma.BitDepth,
internal/protocols/rtmp/from_stream_test.go
+161 -1
View File
@@ -12,6 +12,7 @@ import (
"github.com/bluenviron/gortmplib/pkg/codecs"
"github.com/bluenviron/gortsplib/v5/pkg/description"
"github.com/bluenviron/gortsplib/v5/pkg/format"
"github.com/bluenviron/mediacommon/v2/pkg/codecs/mpeg4audio"
"github.com/bluenviron/mediamtx/internal/logger"
"github.com/bluenviron/mediamtx/internal/stream"
"github.com/bluenviron/mediamtx/internal/test"
@@ -630,6 +631,163 @@ func TestFromStream(t *testing.T) {
}
}
func TestFromStreamLegacyClientMultipleTracks(t *testing.T) {
// Test that legacy RTMP clients (without enhanced RTMP)
// only receive one H264 track and one MPEG4-audio track
// when multiple tracks of each type are available
h264SPS1 := []byte{
0x67, 0x42, 0xc0, 0x28, 0xd9, 0x00, 0x78, 0x02,
0x27, 0xe5, 0x84, 0x00, 0x00, 0x03, 0x00, 0x04,
0x00, 0x00, 0x03, 0x00, 0xf0, 0x3c, 0x60, 0xc9, 0x20,
}
h264PPS1 := []byte{0x08, 0x06, 0x07, 0x08}
h264SPS2 := []byte{
0x67, 0x42, 0xc0, 0x28, 0xd9, 0x00, 0x78, 0x02,
0x27, 0xe5, 0x84, 0x00, 0x00, 0x03, 0x00, 0x04,
0x00, 0x00, 0x03, 0x00, 0xf0, 0x3c, 0x60, 0xc9, 0x21,
}
h264PPS2 := []byte{0x08, 0x06, 0x07, 0x09}
aacConfig1 := test.FormatMPEG4Audio.Config
aacConfig2 := &mpeg4audio.AudioSpecificConfig{
Type: 2, // MPEG4-AAC LC
SampleRate: 48000,
ChannelCount: 2,
ChannelConfig: 2,
}
medias := []*description.Media{
{
Formats: []format.Format{&format.H264{
PayloadTyp: 96,
PacketizationMode: 1,
SPS: h264SPS1,
PPS: h264PPS1,
}},
},
{
Formats: []format.Format{&format.H264{
PayloadTyp: 97,
PacketizationMode: 1,
SPS: h264SPS2,
PPS: h264PPS2,
}},
},
{
Formats: []format.Format{&format.MPEG4Audio{
PayloadTyp: 98,
Config: aacConfig1,
SizeLength: 13,
IndexLength: 3,
IndexDeltaLength: 3,
}},
},
{
Formats: []format.Format{&format.MPEG4Audio{
PayloadTyp: 99,
Config: aacConfig2,
SizeLength: 13,
IndexLength: 3,
IndexDeltaLength: 3,
}},
},
}
strm := &stream.Stream{
Desc: &description.Session{Medias: medias},
WriteQueueSize: 512,
RTPMaxPayloadSize: 1450,
Parent: test.NilLogger,
}
err := strm.Initialize()
require.NoError(t, err)
subStream := &stream.SubStream{
Stream: strm,
UseRTPPackets: false,
}
err = subStream.Initialize()
require.NoError(t, err)
ln, err := net.Listen("tcp", "127.0.0.1:9121")
require.NoError(t, err)
defer ln.Close()
done := make(chan struct{})
go func() {
u, err2 := url.Parse("rtmp://127.0.0.1:9121/stream")
require.NoError(t, err2)
c := &gortmplib.Client{
URL: u,
}
err2 = c.Initialize(context.Background())
require.NoError(t, err2)
r := &gortmplib.Reader{
Conn: c,
}
err2 = r.Initialize()
require.NoError(t, err2)
require.Equal(t, []*gortmplib.Track{
{Codec: &codecs.H264{
SPS: h264SPS1,
PPS: h264PPS1,
}},
{Codec: &codecs.MPEG4Audio{
Config: aacConfig1,
}},
}, r.Tracks())
close(done)
}()
nconn, err := ln.Accept()
require.NoError(t, err)
defer nconn.Close()
conn := &gortmplib.ServerConn{
RW: nconn,
}
err = conn.Initialize()
require.NoError(t, err)
err = conn.Accept()
require.NoError(t, err)
// Simulate a legacy client by clearing the FourCcList
conn.FourCcList = []any{}
r := &stream.Reader{Parent: test.NilLogger}
err = FromStream(strm.Desc, r, conn, nconn, 10*time.Second)
require.NoError(t, err)
strm.AddReader(r)
defer strm.RemoveReader(r)
// Write units to trigger track setup
subStream.WriteUnit(medias[0], medias[0].Formats[0], &unit.Unit{
PTS: 0,
Payload: unit.PayloadH264{
{5, 1}, // IDR
},
})
subStream.WriteUnit(medias[2], medias[2].Formats[0], &unit.Unit{
PTS: 90000,
Payload: unit.PayloadMPEG4Audio{
{3, 4},
},
})
<-done
}
func TestFromStreamNoSupportedCodecs(t *testing.T) {
desc := &description.Session{Medias: []*description.Media{{
Type: description.MediaTypeVideo,
@@ -642,7 +800,9 @@ func TestFromStreamNoSupportedCodecs(t *testing.T) {
}),
}
err := FromStream(desc, r, nil, nil, 0)
conn := &gortmplib.ServerConn{}
err := FromStream(desc, r, conn, nil, 0)
require.Equal(t, errNoSupportedCodecsFrom, err)
}