mirror of
https://github.com/aler9/rtsp-simple-server
synced 2025-09-26 19:51:26 +08:00
526 lines
11 KiB
Go
526 lines
11 KiB
Go
package recorder
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"fmt"
|
|
"io"
|
|
"time"
|
|
|
|
rtspformat "github.com/bluenviron/gortsplib/v4/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/mpeg4audio"
|
|
"github.com/bluenviron/mediacommon/v2/pkg/codecs/mpeg4video"
|
|
"github.com/bluenviron/mediacommon/v2/pkg/formats/mpegts"
|
|
|
|
"github.com/bluenviron/mediamtx/internal/defs"
|
|
"github.com/bluenviron/mediamtx/internal/logger"
|
|
"github.com/bluenviron/mediamtx/internal/unit"
|
|
)
|
|
|
|
const (
|
|
mpegtsBufferSize = 64 * 1024
|
|
)
|
|
|
|
func multiplyAndDivide(v, m, d int64) int64 {
|
|
secs := v / d
|
|
dec := v % d
|
|
return (secs*m + dec*m/d)
|
|
}
|
|
|
|
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))
|
|
}
|
|
|
|
type dynamicWriter struct {
|
|
w io.Writer
|
|
}
|
|
|
|
func (d *dynamicWriter) Write(p []byte) (int, error) {
|
|
return d.w.Write(p)
|
|
}
|
|
|
|
func (d *dynamicWriter) setTarget(w io.Writer) {
|
|
d.w = w
|
|
}
|
|
|
|
type formatMPEGTS struct {
|
|
ri *recorderInstance
|
|
|
|
dw *dynamicWriter
|
|
bw *bufio.Writer
|
|
mw *mpegts.Writer
|
|
hasVideo bool
|
|
currentSegment *formatMPEGTSSegment
|
|
}
|
|
|
|
func (f *formatMPEGTS) initialize() bool {
|
|
var tracks []*mpegts.Track
|
|
var setuppedFormats []rtspformat.Format
|
|
setuppedFormatsMap := make(map[rtspformat.Format]struct{})
|
|
|
|
addTrack := func(format rtspformat.Format, codec mpegts.Codec) *mpegts.Track {
|
|
track := &mpegts.Track{
|
|
Codec: codec,
|
|
}
|
|
|
|
tracks = append(tracks, track)
|
|
setuppedFormats = append(setuppedFormats, format)
|
|
setuppedFormatsMap[format] = struct{}{}
|
|
return track
|
|
}
|
|
|
|
for _, media := range f.ri.stream.Desc.Medias {
|
|
for _, forma := range media.Formats {
|
|
clockRate := forma.ClockRate()
|
|
|
|
switch forma := forma.(type) {
|
|
case *rtspformat.H265: //nolint:dupl
|
|
track := addTrack(forma, &mpegts.CodecH265{})
|
|
|
|
var dtsExtractor *h265.DTSExtractor
|
|
|
|
f.ri.stream.AddReader(
|
|
f.ri,
|
|
media,
|
|
forma,
|
|
func(u unit.Unit) error {
|
|
tunit := u.(*unit.H265)
|
|
if tunit.AU == nil {
|
|
return nil
|
|
}
|
|
|
|
randomAccess := h265.IsRandomAccess(tunit.AU)
|
|
|
|
if dtsExtractor == nil {
|
|
if !randomAccess {
|
|
return nil
|
|
}
|
|
dtsExtractor = &h265.DTSExtractor{}
|
|
dtsExtractor.Initialize()
|
|
}
|
|
|
|
dts, err := dtsExtractor.Extract(tunit.AU, tunit.PTS)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return f.write(
|
|
timestampToDuration(dts, clockRate),
|
|
tunit.NTP,
|
|
true,
|
|
randomAccess,
|
|
func() error {
|
|
return f.mw.WriteH265(
|
|
track,
|
|
tunit.PTS, // no conversion is needed since clock rate is 90khz in both MPEG-TS and RTSP
|
|
dts,
|
|
tunit.AU)
|
|
},
|
|
)
|
|
})
|
|
|
|
case *rtspformat.H264: //nolint:dupl
|
|
track := addTrack(forma, &mpegts.CodecH264{})
|
|
|
|
var dtsExtractor *h264.DTSExtractor
|
|
|
|
f.ri.stream.AddReader(
|
|
f.ri,
|
|
media,
|
|
forma,
|
|
func(u unit.Unit) error {
|
|
tunit := u.(*unit.H264)
|
|
if tunit.AU == nil {
|
|
return nil
|
|
}
|
|
|
|
randomAccess := h264.IsRandomAccess(tunit.AU)
|
|
|
|
if dtsExtractor == nil {
|
|
if !randomAccess {
|
|
return nil
|
|
}
|
|
dtsExtractor = &h264.DTSExtractor{}
|
|
dtsExtractor.Initialize()
|
|
}
|
|
|
|
dts, err := dtsExtractor.Extract(tunit.AU, tunit.PTS)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return f.write(
|
|
timestampToDuration(dts, clockRate),
|
|
tunit.NTP,
|
|
true,
|
|
randomAccess,
|
|
func() error {
|
|
return f.mw.WriteH264(
|
|
track,
|
|
tunit.PTS, // no conversion is needed since clock rate is 90khz in both MPEG-TS and RTSP
|
|
dts,
|
|
tunit.AU)
|
|
},
|
|
)
|
|
})
|
|
|
|
case *rtspformat.MPEG4Video:
|
|
track := addTrack(forma, &mpegts.CodecMPEG4Video{})
|
|
|
|
firstReceived := false
|
|
var lastPTS int64
|
|
|
|
f.ri.stream.AddReader(
|
|
f.ri,
|
|
media,
|
|
forma,
|
|
func(u unit.Unit) error {
|
|
tunit := u.(*unit.MPEG4Video)
|
|
if tunit.Frame == nil {
|
|
return nil
|
|
}
|
|
|
|
if !firstReceived {
|
|
firstReceived = true
|
|
} else if tunit.PTS < lastPTS {
|
|
return fmt.Errorf("MPEG-4 Video streams with B-frames are not supported (yet)")
|
|
}
|
|
lastPTS = tunit.PTS
|
|
|
|
randomAccess := bytes.Contains(tunit.Frame, []byte{0, 0, 1, byte(mpeg4video.GroupOfVOPStartCode)})
|
|
|
|
return f.write(
|
|
timestampToDuration(tunit.PTS, clockRate),
|
|
tunit.NTP,
|
|
true,
|
|
randomAccess,
|
|
func() error {
|
|
return f.mw.WriteMPEG4Video(
|
|
track,
|
|
tunit.PTS, // no conversion is needed since clock rate is 90khz in both MPEG-TS and RTSP
|
|
tunit.Frame)
|
|
},
|
|
)
|
|
})
|
|
|
|
case *rtspformat.MPEG1Video:
|
|
track := addTrack(forma, &mpegts.CodecMPEG1Video{})
|
|
|
|
firstReceived := false
|
|
var lastPTS int64
|
|
|
|
f.ri.stream.AddReader(
|
|
f.ri,
|
|
media,
|
|
forma,
|
|
func(u unit.Unit) error {
|
|
tunit := u.(*unit.MPEG1Video)
|
|
if tunit.Frame == nil {
|
|
return nil
|
|
}
|
|
|
|
if !firstReceived {
|
|
firstReceived = true
|
|
} else if tunit.PTS < lastPTS {
|
|
return fmt.Errorf("MPEG-1 Video streams with B-frames are not supported (yet)")
|
|
}
|
|
lastPTS = tunit.PTS
|
|
|
|
randomAccess := bytes.Contains(tunit.Frame, []byte{0, 0, 1, 0xB8})
|
|
|
|
return f.write(
|
|
timestampToDuration(tunit.PTS, clockRate),
|
|
tunit.NTP,
|
|
true,
|
|
randomAccess,
|
|
func() error {
|
|
return f.mw.WriteMPEG1Video(
|
|
track,
|
|
tunit.PTS, // no conversion is needed since clock rate is 90khz in both MPEG-TS and RTSP
|
|
tunit.Frame)
|
|
},
|
|
)
|
|
})
|
|
|
|
case *rtspformat.Opus:
|
|
track := addTrack(forma, &mpegts.CodecOpus{
|
|
ChannelCount: forma.ChannelCount,
|
|
})
|
|
|
|
f.ri.stream.AddReader(
|
|
f.ri,
|
|
media,
|
|
forma,
|
|
func(u unit.Unit) error {
|
|
tunit := u.(*unit.Opus)
|
|
if tunit.Packets == nil {
|
|
return nil
|
|
}
|
|
|
|
return f.write(
|
|
timestampToDuration(tunit.PTS, clockRate),
|
|
tunit.NTP,
|
|
false,
|
|
true,
|
|
func() error {
|
|
return f.mw.WriteOpus(
|
|
track,
|
|
multiplyAndDivide(tunit.PTS, 90000, int64(clockRate)),
|
|
tunit.Packets)
|
|
},
|
|
)
|
|
})
|
|
|
|
case *rtspformat.KLV:
|
|
track := addTrack(forma, &mpegts.CodecKLV{
|
|
Synchronous: true,
|
|
})
|
|
|
|
f.ri.stream.AddReader(
|
|
f.ri,
|
|
media,
|
|
forma,
|
|
func(u unit.Unit) error {
|
|
tunit := u.(*unit.KLV)
|
|
if tunit.Unit == nil {
|
|
return nil
|
|
}
|
|
|
|
return f.write(
|
|
timestampToDuration(tunit.PTS, 90000),
|
|
tunit.NTP,
|
|
false,
|
|
true,
|
|
func() error {
|
|
return f.mw.WriteKLV(track, multiplyAndDivide(tunit.PTS, 90000, 90000), tunit.Unit)
|
|
},
|
|
)
|
|
})
|
|
|
|
case *rtspformat.MPEG4Audio:
|
|
track := addTrack(forma, &mpegts.CodecMPEG4Audio{
|
|
Config: *forma.Config,
|
|
})
|
|
|
|
f.ri.stream.AddReader(
|
|
f.ri,
|
|
media,
|
|
forma,
|
|
func(u unit.Unit) error {
|
|
tunit := u.(*unit.MPEG4Audio)
|
|
if tunit.AUs == nil {
|
|
return nil
|
|
}
|
|
|
|
return f.write(
|
|
timestampToDuration(tunit.PTS, clockRate),
|
|
tunit.NTP,
|
|
false,
|
|
true,
|
|
func() error {
|
|
return f.mw.WriteMPEG4Audio(
|
|
track,
|
|
multiplyAndDivide(tunit.PTS, 90000, int64(clockRate)),
|
|
tunit.AUs)
|
|
},
|
|
)
|
|
})
|
|
|
|
case *rtspformat.MPEG4AudioLATM:
|
|
if !forma.CPresent {
|
|
track := addTrack(forma, &mpegts.CodecMPEG4Audio{
|
|
Config: *forma.StreamMuxConfig.Programs[0].Layers[0].AudioSpecificConfig,
|
|
})
|
|
|
|
f.ri.stream.AddReader(
|
|
f.ri,
|
|
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
|
|
}
|
|
|
|
return f.write(
|
|
timestampToDuration(tunit.PTS, clockRate),
|
|
tunit.NTP,
|
|
false,
|
|
true,
|
|
func() error {
|
|
return f.mw.WriteMPEG4Audio(
|
|
track,
|
|
multiplyAndDivide(tunit.PTS, 90000, int64(clockRate)),
|
|
[][]byte{ame.Payloads[0][0][0]})
|
|
},
|
|
)
|
|
})
|
|
}
|
|
|
|
case *rtspformat.MPEG1Audio:
|
|
track := addTrack(forma, &mpegts.CodecMPEG1Audio{})
|
|
|
|
f.ri.stream.AddReader(
|
|
f.ri,
|
|
media,
|
|
forma,
|
|
func(u unit.Unit) error {
|
|
tunit := u.(*unit.MPEG1Audio)
|
|
if tunit.Frames == nil {
|
|
return nil
|
|
}
|
|
|
|
return f.write(
|
|
timestampToDuration(tunit.PTS, clockRate),
|
|
tunit.NTP,
|
|
false,
|
|
true,
|
|
func() error {
|
|
return f.mw.WriteMPEG1Audio(
|
|
track,
|
|
tunit.PTS, // no conversion is needed since clock rate is 90khz in both MPEG-TS and RTSP
|
|
tunit.Frames)
|
|
},
|
|
)
|
|
})
|
|
|
|
case *rtspformat.AC3:
|
|
track := addTrack(forma, &mpegts.CodecAC3{})
|
|
|
|
f.ri.stream.AddReader(
|
|
f.ri,
|
|
media,
|
|
forma,
|
|
func(u unit.Unit) error {
|
|
tunit := u.(*unit.AC3)
|
|
if tunit.Frames == nil {
|
|
return nil
|
|
}
|
|
|
|
return f.write(
|
|
timestampToDuration(tunit.PTS, clockRate),
|
|
tunit.NTP,
|
|
false,
|
|
true,
|
|
func() error {
|
|
for i, frame := range tunit.Frames {
|
|
framePTS := tunit.PTS + int64(i)*ac3.SamplesPerFrame
|
|
|
|
err := f.mw.WriteAC3(
|
|
track,
|
|
multiplyAndDivide(framePTS, 90000, int64(clockRate)),
|
|
frame)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
},
|
|
)
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
if len(setuppedFormats) == 0 {
|
|
f.ri.Log(logger.Warn, "no supported tracks found, skipping recording")
|
|
return false
|
|
}
|
|
|
|
n := 1
|
|
for _, medi := range f.ri.stream.Desc.Medias {
|
|
for _, forma := range medi.Formats {
|
|
if _, ok := setuppedFormatsMap[forma]; !ok {
|
|
f.ri.Log(logger.Warn, "skipping track %d (%s)", n, forma.Codec())
|
|
}
|
|
n++
|
|
}
|
|
}
|
|
|
|
f.dw = &dynamicWriter{}
|
|
f.bw = bufio.NewWriterSize(f.dw, mpegtsBufferSize)
|
|
|
|
f.mw = &mpegts.Writer{W: f.bw, Tracks: tracks}
|
|
err := f.mw.Initialize()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
f.ri.Log(logger.Info, "recording %s",
|
|
defs.FormatsInfo(setuppedFormats))
|
|
|
|
return true
|
|
}
|
|
|
|
func (f *formatMPEGTS) close() {
|
|
if f.currentSegment != nil {
|
|
f.currentSegment.close() //nolint:errcheck
|
|
}
|
|
}
|
|
|
|
func (f *formatMPEGTS) write(
|
|
dts time.Duration,
|
|
ntp time.Time,
|
|
isVideo bool,
|
|
randomAccess bool,
|
|
writeCB func() error,
|
|
) error {
|
|
if isVideo {
|
|
f.hasVideo = true
|
|
}
|
|
|
|
switch {
|
|
case f.currentSegment == nil:
|
|
f.currentSegment = &formatMPEGTSSegment{
|
|
f: f,
|
|
startDTS: dts,
|
|
startNTP: ntp,
|
|
}
|
|
f.currentSegment.initialize()
|
|
case (!f.hasVideo || isVideo) &&
|
|
randomAccess &&
|
|
(dts-f.currentSegment.startDTS) >= f.ri.segmentDuration:
|
|
f.currentSegment.lastDTS = dts
|
|
err := f.currentSegment.close()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
f.currentSegment = &formatMPEGTSSegment{
|
|
f: f,
|
|
startDTS: dts,
|
|
startNTP: ntp,
|
|
}
|
|
f.currentSegment.initialize()
|
|
|
|
case (dts - f.currentSegment.lastFlush) >= f.ri.partDuration:
|
|
err := f.bw.Flush()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
f.currentSegment.lastFlush = dts
|
|
}
|
|
|
|
f.currentSegment.lastDTS = dts
|
|
|
|
return writeCB()
|
|
}
|