Files
rtsp-simple-server/internal/protocols/rtmp/from_stream.go
2025-09-16 13:10:34 +02:00

434 lines
9.7 KiB
Go

// Package rtmp provides RTMP utilities.
package rtmp
import (
"errors"
"net"
"slices"
"time"
"github.com/bluenviron/gortmplib"
"github.com/bluenviron/gortmplib/pkg/message"
"github.com/bluenviron/gortsplib/v5/pkg/format"
"github.com/bluenviron/mediacommon/v2/pkg/codecs/ac3"
"github.com/bluenviron/mediacommon/v2/pkg/codecs/h264"
"github.com/bluenviron/mediacommon/v2/pkg/codecs/h265"
"github.com/bluenviron/mediacommon/v2/pkg/codecs/mpeg1audio"
"github.com/bluenviron/mediacommon/v2/pkg/codecs/mpeg4audio"
"github.com/bluenviron/mediacommon/v2/pkg/codecs/opus"
"github.com/bluenviron/mediamtx/internal/logger"
"github.com/bluenviron/mediamtx/internal/stream"
"github.com/bluenviron/mediamtx/internal/unit"
)
var errNoSupportedCodecsFrom = errors.New(
"the stream doesn't contain any supported codec, which are currently " +
"AV1, VP9, H265, H264, Opus, MPEG-4 Audio, MPEG-1/2 Audio, AC-3, G711, LPCM")
func multiplyAndDivide2(v, m, d time.Duration) time.Duration {
secs := v / d
dec := v % d
return (secs*m + dec*m/d)
}
func timestampToDuration(t int64, clockRate int) time.Duration {
return multiplyAndDivide2(time.Duration(t), time.Second, time.Duration(clockRate))
}
// FromStream maps a MediaMTX stream to a RTMP stream.
func FromStream(
str *stream.Stream,
reader stream.Reader,
conn *gortmplib.ServerConn,
nconn net.Conn,
writeTimeout time.Duration,
) error {
var tracks []format.Format
var w *gortmplib.Writer
for _, media := range str.Desc.Medias {
for _, forma := range media.Formats {
switch forma := forma.(type) {
case *format.AV1:
if slices.Contains(conn.FourCcList, interface{}(fourCCToString(message.FourCCAV1))) {
str.AddReader(
reader,
media,
forma,
func(u unit.Unit) error {
tunit := u.(*unit.AV1)
if tunit.TU == nil {
return nil
}
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
return (*w).WriteAV1(
forma,
timestampToDuration(tunit.PTS, forma.ClockRate()),
tunit.TU)
})
tracks = append(tracks, forma)
}
case *format.VP9:
if slices.Contains(conn.FourCcList, interface{}(fourCCToString(message.FourCCVP9))) {
str.AddReader(
reader,
media,
forma,
func(u unit.Unit) error {
tunit := u.(*unit.VP9)
if tunit.Frame == nil {
return nil
}
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
return (*w).WriteVP9(
forma,
timestampToDuration(tunit.PTS, forma.ClockRate()),
tunit.Frame)
})
tracks = append(tracks, forma)
}
case *format.H265:
if slices.Contains(conn.FourCcList, interface{}(fourCCToString(message.FourCCHEVC))) {
var videoDTSExtractor *h265.DTSExtractor
str.AddReader(
reader,
media,
forma,
func(u unit.Unit) error {
tunit := u.(*unit.H265)
if tunit.AU == nil {
return nil
}
if videoDTSExtractor == nil {
if !h265.IsRandomAccess(tunit.AU) {
return nil
}
videoDTSExtractor = &h265.DTSExtractor{}
videoDTSExtractor.Initialize()
}
dts, err := videoDTSExtractor.Extract(tunit.AU, tunit.PTS)
if err != nil {
return err
}
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
return (*w).WriteH265(
forma,
timestampToDuration(tunit.PTS, forma.ClockRate()),
timestampToDuration(dts, forma.ClockRate()),
tunit.AU)
})
tracks = append(tracks, forma)
}
case *format.H264:
var videoDTSExtractor *h264.DTSExtractor
str.AddReader(
reader,
media,
forma,
func(u unit.Unit) error {
tunit := u.(*unit.H264)
if tunit.AU == nil {
return nil
}
idrPresent := false
nonIDRPresent := false
for _, nalu := range tunit.AU {
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 {
return nil
}
videoDTSExtractor = &h264.DTSExtractor{}
videoDTSExtractor.Initialize()
} else if !idrPresent && !nonIDRPresent {
return nil
}
dts, err := videoDTSExtractor.Extract(tunit.AU, tunit.PTS)
if err != nil {
return err
}
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
return (*w).WriteH264(
forma,
timestampToDuration(tunit.PTS, forma.ClockRate()),
timestampToDuration(dts, forma.ClockRate()),
tunit.AU)
})
tracks = append(tracks, forma)
case *format.Opus:
if slices.Contains(conn.FourCcList, interface{}(fourCCToString(message.FourCCOpus))) {
str.AddReader(
reader,
media,
forma,
func(u unit.Unit) error {
tunit := u.(*unit.Opus)
if tunit.Packets == nil {
return nil
}
pts := tunit.PTS
for _, pkt := range tunit.Packets {
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
err := (*w).WriteOpus(
forma,
timestampToDuration(pts, forma.ClockRate()),
pkt,
)
if err != nil {
return err
}
pts += opus.PacketDuration2(pkt)
}
return nil
})
tracks = append(tracks, forma)
}
case *format.MPEG4Audio:
str.AddReader(
reader,
media,
forma,
func(u unit.Unit) error {
tunit := u.(*unit.MPEG4Audio)
if tunit.AUs == nil {
return nil
}
for i, au := range tunit.AUs {
pts := tunit.PTS + int64(i)*mpeg4audio.SamplesPerAccessUnit
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
err := (*w).WriteMPEG4Audio(
forma,
timestampToDuration(pts, forma.ClockRate()),
au,
)
if err != nil {
return err
}
}
return nil
})
tracks = append(tracks, forma)
case *format.MPEG4AudioLATM:
if !forma.CPresent {
str.AddReader(
reader,
media,
forma,
func(u unit.Unit) error {
tunit := u.(*unit.MPEG4AudioLATM)
if tunit.Element == nil {
return nil
}
var ame mpeg4audio.AudioMuxElement
ame.StreamMuxConfig = forma.StreamMuxConfig
err := ame.Unmarshal(tunit.Element)
if err != nil {
return err
}
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
return (*w).WriteMPEG4Audio(
forma,
timestampToDuration(tunit.PTS, forma.ClockRate()),
ame.Payloads[0][0][0],
)
})
tracks = append(tracks, forma)
}
case *format.MPEG1Audio:
// TODO: check sample rate and layer,
// unfortunately they are not available at this stage.
str.AddReader(
reader,
media,
forma,
func(u unit.Unit) error {
tunit := u.(*unit.MPEG1Audio)
pts := tunit.PTS
for _, frame := range tunit.Frames {
var h mpeg1audio.FrameHeader
err := h.Unmarshal(frame)
if err != nil {
return err
}
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
err = (*w).WriteMPEG1Audio(
forma,
timestampToDuration(pts, forma.ClockRate()),
frame)
if err != nil {
return err
}
pts += int64(h.SampleCount()) *
int64(forma.ClockRate()) / int64(h.SampleRate)
}
return nil
})
tracks = append(tracks, forma)
case *format.AC3:
if slices.Contains(conn.FourCcList, interface{}(fourCCToString(message.FourCCAC3))) {
str.AddReader(
reader,
media,
forma,
func(u unit.Unit) error {
tunit := u.(*unit.AC3)
for i, frame := range tunit.Frames {
pts := tunit.PTS + int64(i)*ac3.SamplesPerFrame
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
err := (*w).WriteAC3(
forma,
timestampToDuration(pts, forma.ClockRate()),
frame)
if err != nil {
return err
}
}
return nil
})
tracks = append(tracks, forma)
}
case *format.G711:
if forma.SampleRate == 8000 {
str.AddReader(
reader,
media,
forma,
func(u unit.Unit) error {
tunit := u.(*unit.G711)
if tunit.Samples == nil {
return nil
}
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
return (*w).WriteG711(
forma,
timestampToDuration(tunit.PTS, forma.ClockRate()),
tunit.Samples,
)
})
tracks = append(tracks, forma)
}
case *format.LPCM:
if (forma.ChannelCount == 1 || forma.ChannelCount == 2) &&
(forma.SampleRate == 5512 ||
forma.SampleRate == 11025 ||
forma.SampleRate == 22050 ||
forma.SampleRate == 44100) {
str.AddReader(
reader,
media,
forma,
func(u unit.Unit) error {
tunit := u.(*unit.LPCM)
if tunit.Samples == nil {
return nil
}
nconn.SetWriteDeadline(time.Now().Add(writeTimeout))
return (*w).WriteLPCM(
forma,
timestampToDuration(tunit.PTS, forma.ClockRate()),
tunit.Samples,
)
})
tracks = append(tracks, forma)
}
}
}
}
if len(tracks) == 0 {
return errNoSupportedCodecsFrom
}
w = &gortmplib.Writer{
Conn: conn,
Tracks: tracks,
}
err := w.Initialize()
if err != nil {
return err
}
n := 1
for _, media := range str.Desc.Medias {
for _, forma := range media.Formats {
if !slices.Contains(tracks, forma) {
reader.Log(logger.Warn, "skipping track %d (%s)", n, forma.Codec())
}
n++
}
}
return nil
}