2020-10-19 20:17:48 +00:00
|
|
|
package sourcertmp
|
2020-10-03 19:10:41 +00:00
|
|
|
|
|
|
|
import (
|
2021-04-24 18:47:43 +00:00
|
|
|
"context"
|
2020-10-03 19:10:41 +00:00
|
|
|
"fmt"
|
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-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"
|
2020-10-25 10:41:41 +00:00
|
|
|
|
2021-04-03 09:02:06 +00:00
|
|
|
"github.com/aler9/rtsp-simple-server/internal/h264"
|
2020-12-08 11:21:06 +00:00
|
|
|
"github.com/aler9/rtsp-simple-server/internal/logger"
|
2021-03-22 20:40:07 +00:00
|
|
|
"github.com/aler9/rtsp-simple-server/internal/rtcpsenderset"
|
2021-04-03 09:39:19 +00:00
|
|
|
"github.com/aler9/rtsp-simple-server/internal/rtmp"
|
2021-03-22 20:40:07 +00:00
|
|
|
"github.com/aler9/rtsp-simple-server/internal/source"
|
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 (
|
2021-01-30 20:19:50 +00:00
|
|
|
retryPause = 5 * time.Second
|
2020-10-03 19:10:41 +00:00
|
|
|
)
|
|
|
|
|
2021-04-27 17:21:13 +00:00
|
|
|
// Parent is implemented 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{})
|
2021-03-22 20:40:07 +00:00
|
|
|
OnExtSourceSetReady(req source.ExtSetReadyReq)
|
|
|
|
OnExtSourceSetNotReady(req source.ExtSetNotReadyReq)
|
2020-10-19 20:17:48 +00:00
|
|
|
}
|
|
|
|
|
2021-03-22 20:40:07 +00:00
|
|
|
// Source is a RTMP external source.
|
2020-10-19 20:17:48 +00:00
|
|
|
type Source struct {
|
2021-04-24 18:47:43 +00:00
|
|
|
ur string
|
|
|
|
readTimeout time.Duration
|
|
|
|
writeTimeout time.Duration
|
|
|
|
wg *sync.WaitGroup
|
|
|
|
stats *stats.Stats
|
|
|
|
parent Parent
|
2020-10-19 20:17:48 +00:00
|
|
|
|
|
|
|
// 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,
|
2021-01-31 15:24:58 +00:00
|
|
|
readTimeout time.Duration,
|
2021-04-24 18:47:43 +00:00
|
|
|
writeTimeout time.Duration,
|
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{
|
2021-04-24 18:47:43 +00:00
|
|
|
ur: ur,
|
|
|
|
readTimeout: readTimeout,
|
|
|
|
writeTimeout: writeTimeout,
|
|
|
|
wg: wg,
|
|
|
|
stats: stats,
|
|
|
|
parent: parent,
|
|
|
|
terminate: make(chan struct{}),
|
2020-10-03 19:10:41 +00:00
|
|
|
}
|
|
|
|
|
2021-04-25 14:44:10 +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() {
|
2021-04-25 14:44:10 +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
|
|
|
|
2021-03-22 20:40:07 +00:00
|
|
|
// IsSource implements source.Source.
|
2020-10-19 20:17:48 +00:00
|
|
|
func (s *Source) IsSource() {}
|
|
|
|
|
2021-04-27 17:21:13 +00:00
|
|
|
// IsExtSource implements source.ExtSource.
|
2021-03-14 16:27:04 +00:00
|
|
|
func (s *Source) IsExtSource() {}
|
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-10-31 16:03:03 +00:00
|
|
|
select {
|
2021-03-14 16:27:04 +00:00
|
|
|
case <-time.After(retryPause):
|
2020-10-31 16:03:03 +00:00
|
|
|
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 {
|
2021-04-24 18:47:43 +00:00
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
2020-10-03 19:10:41 +00:00
|
|
|
|
2021-04-24 18:47:43 +00:00
|
|
|
done := make(chan error)
|
2020-10-03 19:10:41 +00:00
|
|
|
go func() {
|
2021-04-24 18:47:43 +00:00
|
|
|
done <- func() error {
|
|
|
|
s.log(logger.Debug, "connecting")
|
2020-10-03 19:10:41 +00:00
|
|
|
|
2021-04-24 18:47:43 +00:00
|
|
|
ctx2, cancel2 := context.WithTimeout(ctx, s.readTimeout)
|
|
|
|
defer cancel2()
|
2021-01-31 20:42:27 +00:00
|
|
|
|
2021-04-24 18:47:43 +00:00
|
|
|
conn, err := rtmp.DialContext(ctx2, s.ur)
|
|
|
|
if err != nil {
|
|
|
|
return err
|
2021-03-22 20:40:07 +00:00
|
|
|
}
|
|
|
|
|
2021-04-24 18:47:43 +00:00
|
|
|
readDone := make(chan error)
|
|
|
|
go func() {
|
|
|
|
readDone <- func() error {
|
|
|
|
conn.NetConn().SetReadDeadline(time.Now().Add(s.readTimeout))
|
|
|
|
conn.NetConn().SetWriteDeadline(time.Now().Add(s.writeTimeout))
|
|
|
|
err = conn.ClientHandshake()
|
|
|
|
if err != nil {
|
|
|
|
return err
|
2021-01-30 20:19:50 +00:00
|
|
|
}
|
2020-10-03 19:10:41 +00:00
|
|
|
|
2021-04-24 18:47:43 +00:00
|
|
|
conn.NetConn().SetWriteDeadline(time.Time{})
|
|
|
|
|
|
|
|
conn.NetConn().SetReadDeadline(time.Now().Add(s.readTimeout))
|
|
|
|
videoTrack, audioTrack, err := conn.ReadMetadata()
|
2021-04-03 09:02:06 +00:00
|
|
|
if err != nil {
|
|
|
|
return err
|
2021-01-30 20:19:50 +00:00
|
|
|
}
|
2020-10-03 19:10:41 +00:00
|
|
|
|
2021-04-24 18:47:43 +00:00
|
|
|
var tracks gortsplib.Tracks
|
2021-04-05 16:14:06 +00:00
|
|
|
|
2021-04-24 18:47:43 +00:00
|
|
|
var h264Encoder *rtph264.Encoder
|
|
|
|
if videoTrack != nil {
|
|
|
|
h264Encoder = rtph264.NewEncoder(96, nil, nil, nil)
|
|
|
|
tracks = append(tracks, videoTrack)
|
2021-03-24 18:40:58 +00:00
|
|
|
}
|
2021-03-10 14:06:45 +00:00
|
|
|
|
2021-04-24 18:47:43 +00:00
|
|
|
var aacEncoder *rtpaac.Encoder
|
|
|
|
if audioTrack != nil {
|
|
|
|
clockRate, _ := audioTrack.ClockRate()
|
|
|
|
aacEncoder = rtpaac.NewEncoder(96, clockRate, nil, nil, nil)
|
|
|
|
tracks = append(tracks, audioTrack)
|
2021-03-24 18:40:58 +00:00
|
|
|
}
|
|
|
|
|
2021-04-24 18:47:43 +00:00
|
|
|
for i, t := range tracks {
|
|
|
|
t.ID = i
|
2021-01-30 20:19:50 +00:00
|
|
|
}
|
2020-10-05 19:07:34 +00:00
|
|
|
|
2021-04-24 18:47:43 +00:00
|
|
|
s.log(logger.Info, "ready")
|
|
|
|
|
|
|
|
cres := make(chan source.ExtSetReadyRes)
|
|
|
|
s.parent.OnExtSourceSetReady(source.ExtSetReadyReq{
|
|
|
|
Tracks: tracks,
|
|
|
|
Res: cres,
|
|
|
|
})
|
|
|
|
res := <-cres
|
|
|
|
|
|
|
|
defer func() {
|
|
|
|
res := make(chan struct{})
|
|
|
|
s.parent.OnExtSourceSetNotReady(source.ExtSetNotReadyReq{
|
|
|
|
Res: res,
|
|
|
|
})
|
|
|
|
<-res
|
|
|
|
}()
|
|
|
|
|
|
|
|
rtcpSenders := rtcpsenderset.New(tracks, res.SP.OnFrame)
|
|
|
|
defer rtcpSenders.Close()
|
|
|
|
|
|
|
|
onFrame := func(trackID int, payload []byte) {
|
|
|
|
rtcpSenders.OnFrame(trackID, gortsplib.StreamTypeRTP, payload)
|
|
|
|
res.SP.OnFrame(trackID, gortsplib.StreamTypeRTP, payload)
|
2021-01-30 20:19:50 +00:00
|
|
|
}
|
2020-10-05 19:07:34 +00:00
|
|
|
|
2021-04-24 18:47:43 +00:00
|
|
|
for {
|
|
|
|
conn.NetConn().SetReadDeadline(time.Now().Add(s.readTimeout))
|
|
|
|
pkt, err := conn.ReadPacket()
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
2020-10-05 19:07:34 +00:00
|
|
|
|
2021-04-24 18:47:43 +00:00
|
|
|
switch pkt.Type {
|
|
|
|
case av.H264:
|
|
|
|
if videoTrack == nil {
|
|
|
|
return fmt.Errorf("ERR: received an H264 frame, but track is not set up")
|
|
|
|
}
|
|
|
|
|
|
|
|
nalus, err := h264.DecodeAVCC(pkt.Data)
|
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
|
|
|
|
var outNALUs [][]byte
|
|
|
|
for _, nalu := range nalus {
|
|
|
|
// remove SPS, PPS and AUD, not needed by RTSP / RTMP
|
|
|
|
typ := h264.NALUType(nalu[0] & 0x1F)
|
|
|
|
switch typ {
|
|
|
|
case h264.NALUTypeSPS, h264.NALUTypePPS, h264.NALUTypeAccessUnitDelimiter:
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
|
|
|
|
outNALUs = append(outNALUs, nalu)
|
|
|
|
}
|
|
|
|
|
|
|
|
pkts, err := h264Encoder.Encode(outNALUs, pkt.Time+pkt.CTime)
|
|
|
|
if err != nil {
|
|
|
|
return fmt.Errorf("ERR while encoding H264: %v", err)
|
|
|
|
}
|
|
|
|
|
|
|
|
for _, pkt := range pkts {
|
|
|
|
onFrame(videoTrack.ID, pkt)
|
|
|
|
}
|
|
|
|
|
|
|
|
case av.AAC:
|
|
|
|
if audioTrack == nil {
|
|
|
|
return fmt.Errorf("ERR: received an AAC frame, but track is not set up")
|
|
|
|
}
|
|
|
|
|
|
|
|
pkts, err := aacEncoder.Encode([][]byte{pkt.Data}, pkt.Time+pkt.CTime)
|
|
|
|
if err != nil {
|
|
|
|
return fmt.Errorf("ERR while encoding AAC: %v", err)
|
|
|
|
}
|
|
|
|
|
|
|
|
for _, pkt := range pkts {
|
|
|
|
onFrame(audioTrack.ID, pkt)
|
|
|
|
}
|
|
|
|
|
|
|
|
default:
|
|
|
|
return fmt.Errorf("ERR: unexpected packet: %v", pkt.Type)
|
|
|
|
}
|
2021-03-24 20:35:23 +00:00
|
|
|
}
|
2021-04-24 18:47:43 +00:00
|
|
|
}()
|
|
|
|
}()
|
2021-01-30 20:19:50 +00:00
|
|
|
|
2021-04-24 18:47:43 +00:00
|
|
|
select {
|
|
|
|
case err := <-readDone:
|
|
|
|
conn.NetConn().Close()
|
|
|
|
return err
|
|
|
|
|
|
|
|
case <-ctx.Done():
|
|
|
|
conn.NetConn().Close()
|
|
|
|
<-readDone
|
|
|
|
return nil
|
2020-10-03 19:10:41 +00:00
|
|
|
}
|
2021-01-30 20:19:50 +00:00
|
|
|
}()
|
2020-10-03 19:10:41 +00:00
|
|
|
}()
|
|
|
|
|
2021-04-04 17:45:41 +00:00
|
|
|
select {
|
2021-04-24 18:47:43 +00:00
|
|
|
case err := <-done:
|
|
|
|
cancel()
|
2021-04-04 17:45:41 +00:00
|
|
|
s.log(logger.Info, "ERR: %s", err)
|
|
|
|
return true
|
2021-04-24 18:47:43 +00:00
|
|
|
|
|
|
|
case <-s.terminate:
|
|
|
|
cancel()
|
|
|
|
<-done
|
|
|
|
return false
|
2020-10-03 19:10:41 +00:00
|
|
|
}
|
|
|
|
}
|