2020-10-19 20:17:48 +00:00
|
|
|
package sourcertmp
|
2020-10-03 19:10:41 +00:00
|
|
|
|
|
|
|
import (
|
|
|
|
"fmt"
|
|
|
|
"net"
|
2020-11-01 15:51:12 +00:00
|
|
|
"sync"
|
2020-10-03 19:10:41 +00:00
|
|
|
"sync/atomic"
|
|
|
|
"time"
|
|
|
|
|
|
|
|
"github.com/aler9/gortsplib"
|
2020-11-22 14:46:16 +00:00
|
|
|
"github.com/aler9/gortsplib/pkg/rtcpsender"
|
2020-11-15 16:56:54 +00:00
|
|
|
"github.com/aler9/gortsplib/pkg/rtpaac"
|
|
|
|
"github.com/aler9/gortsplib/pkg/rtph264"
|
2020-10-03 19:10:41 +00:00
|
|
|
"github.com/notedit/rtmp/av"
|
|
|
|
"github.com/notedit/rtmp/codec/h264"
|
|
|
|
"github.com/notedit/rtmp/format/rtmp"
|
2020-10-25 10:41:41 +00:00
|
|
|
|
2020-12-08 11:21:06 +00:00
|
|
|
"github.com/aler9/rtsp-simple-server/internal/logger"
|
2020-11-01 21:56:56 +00:00
|
|
|
"github.com/aler9/rtsp-simple-server/internal/stats"
|
2020-10-03 19:10:41 +00:00
|
|
|
)
|
|
|
|
|
|
|
|
const (
|
2020-11-11 10:07:40 +00:00
|
|
|
retryPause = 5 * time.Second
|
|
|
|
analyzeTimeout = 8 * time.Second
|
2020-10-03 19:10:41 +00:00
|
|
|
)
|
|
|
|
|
2020-11-05 11:30:25 +00:00
|
|
|
// Parent is implemeneted by path.Path.
|
2020-10-19 20:17:48 +00:00
|
|
|
type Parent interface {
|
2020-12-08 11:21:06 +00:00
|
|
|
Log(logger.Level, string, ...interface{})
|
2020-11-01 15:51:12 +00:00
|
|
|
OnSourceSetReady(gortsplib.Tracks)
|
|
|
|
OnSourceSetNotReady()
|
2020-10-19 20:17:48 +00:00
|
|
|
OnFrame(int, gortsplib.StreamType, []byte)
|
|
|
|
}
|
|
|
|
|
2020-11-05 11:30:25 +00:00
|
|
|
// Source is a RTMP source.
|
2020-10-19 20:17:48 +00:00
|
|
|
type Source struct {
|
|
|
|
ur string
|
2020-11-01 15:51:12 +00:00
|
|
|
wg *sync.WaitGroup
|
2020-10-25 10:41:41 +00:00
|
|
|
stats *stats.Stats
|
2020-10-19 20:17:48 +00:00
|
|
|
parent Parent
|
|
|
|
|
|
|
|
// in
|
2020-11-01 15:51:12 +00:00
|
|
|
terminate chan struct{}
|
2020-10-03 19:10:41 +00:00
|
|
|
}
|
|
|
|
|
2020-11-05 11:30:25 +00:00
|
|
|
// New allocates a Source.
|
2020-10-19 20:17:48 +00:00
|
|
|
func New(ur string,
|
2020-11-01 15:51:12 +00:00
|
|
|
wg *sync.WaitGroup,
|
2020-10-25 10:41:41 +00:00
|
|
|
stats *stats.Stats,
|
2020-10-19 20:17:48 +00:00
|
|
|
parent Parent) *Source {
|
|
|
|
s := &Source{
|
2020-11-01 15:51:12 +00:00
|
|
|
ur: ur,
|
|
|
|
wg: wg,
|
|
|
|
stats: stats,
|
|
|
|
parent: parent,
|
|
|
|
terminate: make(chan struct{}),
|
2020-10-03 19:10:41 +00:00
|
|
|
}
|
|
|
|
|
2020-10-25 10:41:41 +00:00
|
|
|
atomic.AddInt64(s.stats.CountSourcesRtmp, +1)
|
2020-12-08 11:21:06 +00:00
|
|
|
s.log(logger.Info, "started")
|
2020-10-25 10:41:41 +00:00
|
|
|
|
2020-11-01 15:51:12 +00:00
|
|
|
s.wg.Add(1)
|
2020-10-25 10:41:41 +00:00
|
|
|
go s.run()
|
2020-10-19 20:17:48 +00:00
|
|
|
return s
|
|
|
|
}
|
2020-10-03 19:10:41 +00:00
|
|
|
|
2020-11-05 11:30:25 +00:00
|
|
|
// Close closes a Source.
|
2020-10-19 20:17:48 +00:00
|
|
|
func (s *Source) Close() {
|
2020-11-01 15:51:12 +00:00
|
|
|
atomic.AddInt64(s.stats.CountSourcesRtmpRunning, -1)
|
2020-12-08 11:21:06 +00:00
|
|
|
s.log(logger.Info, "stopped")
|
2020-10-19 20:17:48 +00:00
|
|
|
close(s.terminate)
|
|
|
|
}
|
2020-10-03 19:10:41 +00:00
|
|
|
|
2020-11-05 11:30:25 +00:00
|
|
|
// IsSource implements path.source.
|
2020-10-19 20:17:48 +00:00
|
|
|
func (s *Source) IsSource() {}
|
|
|
|
|
2020-11-05 11:30:25 +00:00
|
|
|
// IsSourceExternal implements path.sourceExternal.
|
2020-11-01 15:51:12 +00:00
|
|
|
func (s *Source) IsSourceExternal() {}
|
2020-10-03 19:10:41 +00:00
|
|
|
|
2020-12-08 11:21:06 +00:00
|
|
|
func (s *Source) log(level logger.Level, format string, args ...interface{}) {
|
|
|
|
s.parent.Log(level, "[rtmp source] "+format, args...)
|
|
|
|
}
|
|
|
|
|
2020-10-25 10:41:41 +00:00
|
|
|
func (s *Source) run() {
|
2020-11-01 15:51:12 +00:00
|
|
|
defer s.wg.Done()
|
2020-10-03 19:10:41 +00:00
|
|
|
|
|
|
|
for {
|
2020-10-31 16:03:03 +00:00
|
|
|
ok := func() bool {
|
2020-11-01 15:51:12 +00:00
|
|
|
ok := s.runInner()
|
2020-10-31 16:03:03 +00:00
|
|
|
if !ok {
|
|
|
|
return false
|
|
|
|
}
|
2020-10-03 19:10:41 +00:00
|
|
|
|
2020-11-01 15:51:12 +00:00
|
|
|
t := time.NewTimer(retryPause)
|
2020-10-31 16:03:03 +00:00
|
|
|
defer t.Stop()
|
2020-10-03 19:10:41 +00:00
|
|
|
|
2020-10-31 16:03:03 +00:00
|
|
|
select {
|
|
|
|
case <-t.C:
|
|
|
|
return true
|
2020-11-01 15:51:12 +00:00
|
|
|
case <-s.terminate:
|
2020-10-31 16:03:03 +00:00
|
|
|
return false
|
|
|
|
}
|
|
|
|
}()
|
|
|
|
if !ok {
|
|
|
|
break
|
2020-10-03 19:10:41 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2020-11-01 15:51:12 +00:00
|
|
|
func (s *Source) runInner() bool {
|
2020-12-08 11:21:06 +00:00
|
|
|
s.log(logger.Info, "connecting")
|
2020-10-03 19:10:41 +00:00
|
|
|
|
|
|
|
var conn *rtmp.Conn
|
|
|
|
var nconn net.Conn
|
|
|
|
var err error
|
|
|
|
dialDone := make(chan struct{}, 1)
|
|
|
|
go func() {
|
|
|
|
defer close(dialDone)
|
2020-10-19 20:17:48 +00:00
|
|
|
conn, nconn, err = rtmp.NewClient().Dial(s.ur, rtmp.PrepareReading)
|
2020-10-03 19:10:41 +00:00
|
|
|
}()
|
|
|
|
|
|
|
|
select {
|
2020-11-01 15:51:12 +00:00
|
|
|
case <-s.terminate:
|
2020-10-03 19:10:41 +00:00
|
|
|
return false
|
|
|
|
case <-dialDone:
|
|
|
|
}
|
|
|
|
|
|
|
|
if err != nil {
|
2020-12-08 11:21:06 +00:00
|
|
|
s.log(logger.Info, "ERR: %s", err)
|
2020-10-03 19:10:41 +00:00
|
|
|
return true
|
|
|
|
}
|
|
|
|
|
2020-10-05 19:07:34 +00:00
|
|
|
// gather video and audio features
|
|
|
|
var h264Sps []byte
|
|
|
|
var h264Pps []byte
|
|
|
|
var aacConfig []byte
|
|
|
|
confDone := make(chan struct{})
|
|
|
|
confClose := uint32(0)
|
|
|
|
go func() {
|
|
|
|
defer close(confDone)
|
|
|
|
|
2020-10-03 19:10:41 +00:00
|
|
|
for {
|
2020-10-05 19:07:34 +00:00
|
|
|
var pkt av.Packet
|
|
|
|
pkt, err = conn.ReadPacket()
|
2020-10-03 19:10:41 +00:00
|
|
|
if err != nil {
|
2020-10-05 19:07:34 +00:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
if atomic.LoadUint32(&confClose) > 0 {
|
|
|
|
return
|
2020-10-03 19:10:41 +00:00
|
|
|
}
|
|
|
|
|
2020-10-05 19:07:34 +00:00
|
|
|
switch pkt.Type {
|
|
|
|
case av.H264DecoderConfig:
|
2020-10-03 19:10:41 +00:00
|
|
|
codec, err := h264.FromDecoderConfig(pkt.Data)
|
|
|
|
if err != nil {
|
|
|
|
panic(err)
|
|
|
|
}
|
|
|
|
|
2020-11-11 10:07:40 +00:00
|
|
|
h264Sps, h264Pps = codec.SPS[0], codec.PPS[0]
|
2020-10-05 19:07:34 +00:00
|
|
|
|
|
|
|
if aacConfig != nil {
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
case av.AACDecoderConfig:
|
|
|
|
aacConfig = pkt.Data
|
|
|
|
|
|
|
|
if h264Sps != nil {
|
|
|
|
return
|
|
|
|
}
|
2020-10-03 19:10:41 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}()
|
2020-10-05 19:07:34 +00:00
|
|
|
|
2020-11-11 10:07:40 +00:00
|
|
|
timer := time.NewTimer(analyzeTimeout)
|
2020-10-31 16:03:03 +00:00
|
|
|
defer timer.Stop()
|
2020-10-05 19:07:34 +00:00
|
|
|
|
|
|
|
select {
|
|
|
|
case <-confDone:
|
|
|
|
case <-timer.C:
|
|
|
|
atomic.StoreUint32(&confClose, 1)
|
|
|
|
<-confDone
|
|
|
|
}
|
|
|
|
|
2020-10-03 19:10:41 +00:00
|
|
|
if err != nil {
|
2020-12-08 11:21:06 +00:00
|
|
|
s.log(logger.Info, "ERR: %s", err)
|
2020-10-03 19:10:41 +00:00
|
|
|
return true
|
|
|
|
}
|
|
|
|
|
2020-10-05 19:07:34 +00:00
|
|
|
var tracks gortsplib.Tracks
|
2020-11-22 14:46:16 +00:00
|
|
|
|
2020-10-05 19:07:34 +00:00
|
|
|
var videoTrack *gortsplib.Track
|
2020-11-22 14:46:16 +00:00
|
|
|
var videoRtcpSender *rtcpsender.RtcpSender
|
2020-10-05 19:07:34 +00:00
|
|
|
var h264Encoder *rtph264.Encoder
|
2020-11-22 14:46:16 +00:00
|
|
|
|
|
|
|
var audioTrack *gortsplib.Track
|
|
|
|
var audioRtcpSender *rtcpsender.RtcpSender
|
2020-10-05 19:07:34 +00:00
|
|
|
var aacEncoder *rtpaac.Encoder
|
|
|
|
|
|
|
|
if h264Sps != nil {
|
2020-11-30 14:21:01 +00:00
|
|
|
videoTrack, err = gortsplib.NewTrackH264(96, h264Sps, h264Pps)
|
2020-10-05 19:07:34 +00:00
|
|
|
if err != nil {
|
2020-12-08 11:21:06 +00:00
|
|
|
s.log(logger.Info, "ERR: %s", err)
|
2020-10-05 19:07:34 +00:00
|
|
|
return true
|
|
|
|
}
|
|
|
|
|
2020-11-22 14:46:16 +00:00
|
|
|
clockRate, _ := videoTrack.ClockRate()
|
|
|
|
videoRtcpSender = rtcpsender.New(clockRate)
|
|
|
|
|
2020-11-30 14:21:01 +00:00
|
|
|
h264Encoder, err = rtph264.NewEncoder(96)
|
2020-10-05 19:07:34 +00:00
|
|
|
if err != nil {
|
2020-12-08 11:21:06 +00:00
|
|
|
s.log(logger.Info, "ERR: %s", err)
|
2020-10-05 19:07:34 +00:00
|
|
|
return true
|
|
|
|
}
|
|
|
|
|
|
|
|
tracks = append(tracks, videoTrack)
|
|
|
|
}
|
|
|
|
|
|
|
|
if aacConfig != nil {
|
2020-11-30 14:21:01 +00:00
|
|
|
audioTrack, err = gortsplib.NewTrackAAC(96, aacConfig)
|
2020-10-05 19:07:34 +00:00
|
|
|
if err != nil {
|
2020-12-08 11:21:06 +00:00
|
|
|
s.log(logger.Info, "ERR: %s", err)
|
2020-10-05 19:07:34 +00:00
|
|
|
return true
|
|
|
|
}
|
|
|
|
|
2020-11-22 14:46:16 +00:00
|
|
|
clockRate, _ := audioTrack.ClockRate()
|
|
|
|
audioRtcpSender = rtcpsender.New(clockRate)
|
|
|
|
|
2020-11-30 14:21:01 +00:00
|
|
|
aacEncoder, err = rtpaac.NewEncoder(96, clockRate)
|
2020-10-05 19:07:34 +00:00
|
|
|
if err != nil {
|
2020-12-08 11:21:06 +00:00
|
|
|
s.log(logger.Info, "ERR: %s", err)
|
2020-10-05 19:07:34 +00:00
|
|
|
return true
|
|
|
|
}
|
|
|
|
|
|
|
|
tracks = append(tracks, audioTrack)
|
|
|
|
}
|
|
|
|
|
|
|
|
if len(tracks) == 0 {
|
2020-12-08 11:21:06 +00:00
|
|
|
s.log(logger.Info, "ERR: no tracks found")
|
2020-10-05 19:07:34 +00:00
|
|
|
return true
|
|
|
|
}
|
|
|
|
|
2020-11-30 14:21:01 +00:00
|
|
|
for i, t := range tracks {
|
2020-12-05 19:42:59 +00:00
|
|
|
t.ID = i
|
2020-11-30 14:21:01 +00:00
|
|
|
}
|
|
|
|
|
2020-12-08 11:21:06 +00:00
|
|
|
s.log(logger.Info, "ready")
|
2020-11-01 15:51:12 +00:00
|
|
|
s.parent.OnSourceSetReady(tracks)
|
2020-11-19 11:25:06 +00:00
|
|
|
defer s.parent.OnSourceSetNotReady()
|
2020-10-03 19:10:41 +00:00
|
|
|
|
2020-11-22 14:46:16 +00:00
|
|
|
rtcpTerminate := make(chan struct{})
|
|
|
|
rtcpDone := make(chan struct{})
|
|
|
|
go func() {
|
|
|
|
close(rtcpDone)
|
|
|
|
|
|
|
|
t := time.NewTicker(10 * time.Second)
|
|
|
|
defer t.Stop()
|
|
|
|
|
|
|
|
for {
|
|
|
|
select {
|
|
|
|
case <-t.C:
|
|
|
|
now := time.Now()
|
|
|
|
|
|
|
|
if videoRtcpSender != nil {
|
|
|
|
r := videoRtcpSender.Report(now)
|
|
|
|
if r != nil {
|
2020-12-05 19:42:59 +00:00
|
|
|
s.parent.OnFrame(videoTrack.ID, gortsplib.StreamTypeRtcp, r)
|
2020-11-22 14:46:16 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
if audioRtcpSender != nil {
|
|
|
|
r := audioRtcpSender.Report(now)
|
|
|
|
if r != nil {
|
2020-12-05 19:42:59 +00:00
|
|
|
s.parent.OnFrame(audioTrack.ID, gortsplib.StreamTypeRtcp, r)
|
2020-11-22 14:46:16 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
case <-rtcpTerminate:
|
|
|
|
return
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}()
|
|
|
|
|
2020-11-15 19:15:06 +00:00
|
|
|
readerDone := make(chan error)
|
2020-10-03 19:10:41 +00:00
|
|
|
go func() {
|
|
|
|
for {
|
|
|
|
pkt, err := conn.ReadPacket()
|
|
|
|
if err != nil {
|
2020-11-15 19:15:06 +00:00
|
|
|
readerDone <- err
|
2020-10-03 19:10:41 +00:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2020-10-05 19:07:34 +00:00
|
|
|
switch pkt.Type {
|
|
|
|
case av.H264:
|
|
|
|
if h264Sps == nil {
|
2020-11-15 19:15:06 +00:00
|
|
|
readerDone <- fmt.Errorf("rtmp source ERR: received an H264 frame, but track is not setup up")
|
2020-10-05 19:07:34 +00:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2020-10-03 19:10:41 +00:00
|
|
|
// decode from AVCC format
|
|
|
|
nalus, typ := h264.SplitNALUs(pkt.Data)
|
|
|
|
if typ != h264.NALU_AVCC {
|
2020-11-15 19:15:06 +00:00
|
|
|
readerDone <- fmt.Errorf("invalid NALU format (%d)", typ)
|
2020-10-03 19:10:41 +00:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
// encode into RTP/H264 format
|
2020-11-28 23:32:30 +00:00
|
|
|
frames, err := h264Encoder.Write(pkt.Time+pkt.CTime, nalus)
|
2020-10-03 19:10:41 +00:00
|
|
|
if err != nil {
|
2020-11-15 19:15:06 +00:00
|
|
|
readerDone <- err
|
2020-10-03 19:10:41 +00:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
for _, f := range frames {
|
2020-11-30 14:21:01 +00:00
|
|
|
videoRtcpSender.ProcessFrame(time.Now(), gortsplib.StreamTypeRtp, f)
|
2020-12-05 19:42:59 +00:00
|
|
|
s.parent.OnFrame(videoTrack.ID, gortsplib.StreamTypeRtp, f)
|
2020-10-03 19:10:41 +00:00
|
|
|
}
|
2020-10-05 19:07:34 +00:00
|
|
|
|
|
|
|
case av.AAC:
|
|
|
|
if aacConfig == nil {
|
2020-11-15 19:15:06 +00:00
|
|
|
readerDone <- fmt.Errorf("rtmp source ERR: received an AAC frame, but track is not setup up")
|
2020-10-05 19:07:34 +00:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2020-11-28 23:32:30 +00:00
|
|
|
frames, err := aacEncoder.Write(pkt.Time+pkt.CTime, pkt.Data)
|
2020-10-05 19:07:34 +00:00
|
|
|
if err != nil {
|
2020-11-15 19:15:06 +00:00
|
|
|
readerDone <- err
|
2020-10-05 19:07:34 +00:00
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
for _, f := range frames {
|
2020-11-30 14:21:01 +00:00
|
|
|
audioRtcpSender.ProcessFrame(time.Now(), gortsplib.StreamTypeRtp, f)
|
2020-12-05 19:42:59 +00:00
|
|
|
s.parent.OnFrame(audioTrack.ID, gortsplib.StreamTypeRtp, f)
|
2020-10-05 19:07:34 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
default:
|
2020-11-15 19:15:06 +00:00
|
|
|
readerDone <- fmt.Errorf("rtmp source ERR: unexpected packet: %v", pkt.Type)
|
2020-10-05 19:07:34 +00:00
|
|
|
return
|
2020-10-03 19:10:41 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}()
|
|
|
|
|
|
|
|
for {
|
|
|
|
select {
|
2020-11-01 15:51:12 +00:00
|
|
|
case <-s.terminate:
|
2020-10-03 19:10:41 +00:00
|
|
|
nconn.Close()
|
2020-11-15 19:15:06 +00:00
|
|
|
<-readerDone
|
2020-11-22 14:46:16 +00:00
|
|
|
|
|
|
|
close(rtcpTerminate)
|
|
|
|
<-rtcpDone
|
2020-11-19 11:25:06 +00:00
|
|
|
return false
|
2020-10-03 19:10:41 +00:00
|
|
|
|
2020-11-15 19:15:06 +00:00
|
|
|
case err := <-readerDone:
|
2020-10-03 19:10:41 +00:00
|
|
|
nconn.Close()
|
2020-12-08 11:21:06 +00:00
|
|
|
s.log(logger.Info, "ERR: %s", err)
|
2020-11-22 14:46:16 +00:00
|
|
|
|
|
|
|
close(rtcpTerminate)
|
|
|
|
<-rtcpDone
|
2020-11-19 11:25:06 +00:00
|
|
|
return true
|
2020-10-03 19:10:41 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|