RTMP source: apply read and write timeouts to connection initialization

internal/path/path.go

@@ -420,6 +420,7 @@ func (pa *Path) startExternalSource() {
 		pa.source = sourcertmp.New(
 			pa.conf.Source,
 			pa.readTimeout,
+			pa.writeTimeout,
 			&pa.sourceWg,
 			pa.stats,
 			pa)

internal/rtmp/client.go

@@ -1,18 +1,44 @@
 package rtmp
 
 import (
+	"bufio"
+	"context"
+	"net"
+	"net/url"
+
 	"github.com/notedit/rtmp/format/rtmp"
 )
 
-// Dial connects to a server in reading mode.
+// DialContext connects to a server in reading mode.
-func Dial(address string) (*Conn, error) {
+func DialContext(ctx context.Context, address string) (*Conn, error) {
-	rconn, nconn, err := rtmp.NewClient().Dial(address, rtmp.PrepareReading)
+	// https://github.com/aler9/rtmp/blob/master/format/rtmp/client.go#L74
+
+	u, err := url.Parse(address)
+	if err != nil {
+		return nil, err
+	}
+	host := rtmp.UrlGetHost(u)
+
+	var d net.Dialer
+	nconn, err := d.DialContext(ctx, "tcp", host)
 	if err != nil {
 		return nil, err
 	}
 
+	rw := &bufio.ReadWriter{
+		Reader: bufio.NewReaderSize(nconn, 4096),
+		Writer: bufio.NewWriterSize(nconn, 4096),
+	}
+	rconn := rtmp.NewConn(rw)
+	rconn.URL = u
+
 	return &Conn{
 		rconn: rconn,
 		nconn: nconn,
 	}, nil
 }
+
+// ClientHandshake performs the handshake of a client-side connection.
+func (c *Conn) ClientHandshake() error {
+	return c.rconn.Prepare(rtmp.StageGotPublishOrPlayCommand, rtmp.PrepareReading)
+}

internal/sourcertmp/source.go

@@ -1,6 +1,7 @@
 package sourcertmp
 
 import (
+	"context"
 	"fmt"
 	"sync"
 	"sync/atomic"
@@ -34,6 +35,7 @@ type Parent interface {
 type Source struct {
 	ur           string
 	readTimeout  time.Duration
+	writeTimeout time.Duration
 	wg           *sync.WaitGroup
 	stats        *stats.Stats
 	parent       Parent
@@ -45,12 +47,14 @@ type Source struct {
 // New allocates a Source.
 func New(ur string,
 	readTimeout time.Duration,
+	writeTimeout time.Duration,
 	wg *sync.WaitGroup,
 	stats *stats.Stats,
 	parent Parent) *Source {
 	s := &Source{
 		ur:           ur,
 		readTimeout:  readTimeout,
+		writeTimeout: writeTimeout,
 		wg:           wg,
 		stats:        stats,
 		parent:       parent,
@@ -106,47 +110,37 @@ func (s *Source) run() {
 }
 
 func (s *Source) runInner() bool {
-	s.log(logger.Info, "connecting")
+	ctx, cancel := context.WithCancel(context.Background())
 
-	var conn *rtmp.Conn
+	done := make(chan error)
-	var err error
-	dialDone := make(chan struct{}, 1)
 	go func() {
-		defer close(dialDone)
+		done <- func() error {
-		conn, err = rtmp.Dial(s.ur)
+			s.log(logger.Debug, "connecting")
-	}()
 
-	select {
+			ctx2, cancel2 := context.WithTimeout(ctx, s.readTimeout)
-	case <-dialDone:
+			defer cancel2()
-	case <-s.terminate:
-		return false
-	}
 
+			conn, err := rtmp.DialContext(ctx2, s.ur)
 			if err != nil {
-		s.log(logger.Info, "ERR: %s", err)
+				return err
-		return true
 			}
 
-	var videoTrack *gortsplib.Track
+			readDone := make(chan error)
-	var audioTrack *gortsplib.Track
-	metadataDone := make(chan struct{})
 			go func() {
-		defer close(metadataDone)
+				readDone <- func() error {
 					conn.NetConn().SetReadDeadline(time.Now().Add(s.readTimeout))
-		videoTrack, audioTrack, err = conn.ReadMetadata()
+					conn.NetConn().SetWriteDeadline(time.Now().Add(s.writeTimeout))
-	}()
+					err = conn.ClientHandshake()
-
+					if err != nil {
-	select {
+						return err
-	case <-metadataDone:
-	case <-s.terminate:
-		conn.NetConn().Close()
-		<-metadataDone
-		return false
 					}
 
+					conn.NetConn().SetWriteDeadline(time.Time{})
+
+					conn.NetConn().SetReadDeadline(time.Now().Add(s.readTimeout))
+					videoTrack, audioTrack, err := conn.ReadMetadata()
 					if err != nil {
-		s.log(logger.Info, "ERR: %s", err)
+						return err
-		return true
 					}
 
 					var tracks gortsplib.Tracks
@@ -185,9 +179,6 @@ func (s *Source) runInner() bool {
 						<-res
 					}()
 
-	readerDone := make(chan error)
-	go func() {
-		readerDone <- func() error {
 					rtcpSenders := rtcpsenderset.New(tracks, res.SP.OnFrame)
 					defer rtcpSenders.Close()
 
@@ -257,14 +248,27 @@ func (s *Source) runInner() bool {
 			}()
 
 			select {
-	case <-s.terminate:
+			case err := <-readDone:
 				conn.NetConn().Close()
-		<-readerDone
+				return err
-		return false
 
-	case err := <-readerDone:
+			case <-ctx.Done():
-		s.log(logger.Info, "ERR: %s", err)
 				conn.NetConn().Close()
+				<-readDone
+				return nil
+			}
+		}()
+	}()
+
+	select {
+	case err := <-done:
+		cancel()
+		s.log(logger.Info, "ERR: %s", err)
 		return true
+
+	case <-s.terminate:
+		cancel()
+		<-done
+		return false
 	}
 }

internal/sourcertsp/source.go

@@ -121,7 +121,7 @@ func (s *Source) run() {
 }
 
 func (s *Source) runInner() bool {
-	s.log(logger.Info, "connecting")
+	s.log(logger.Debug, "connecting")
 
 	var conn *gortsplib.ClientConn
 	var err error