From 9f9dc7e844a9241f876a8d703ddac15ca9c8d674 Mon Sep 17 00:00:00 2001
From: Alexey Khit <alexey.khit@gmail.com>
Date: Wed, 23 Aug 2023 14:08:15 +0300
Subject: [PATCH] Add support custom timeout for RTSP source

---
 internal/rtsp/rtsp.go |  1 +
 pkg/core/helpers.go   |  4 +++-
 pkg/rtmp/producer.go  |  2 +-
 pkg/rtsp/client.go    |  6 +++++-
 pkg/rtsp/conn.go      | 17 +++++++++++------
 pkg/tcp/dial.go       |  7 +++----
 6 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/internal/rtsp/rtsp.go b/internal/rtsp/rtsp.go
index 10023a41..8f7d278f 100644
--- a/internal/rtsp/rtsp.go
+++ b/internal/rtsp/rtsp.go
@@ -102,6 +102,7 @@ func rtspHandler(rawURL string) (core.Producer, error) {
 		query := streams.ParseQuery(rawQuery)
 		conn.Backchannel = query.Get("backchannel") == "1"
 		conn.Media = query.Get("media")
+		conn.Timeout = core.Atoi(query.Get("timeout"))
 		conn.Transport = query.Get("transport")
 	}
 
diff --git a/pkg/core/helpers.go b/pkg/core/helpers.go
index ca288834..c65a84df 100644
--- a/pkg/core/helpers.go
+++ b/pkg/core/helpers.go
@@ -63,7 +63,9 @@ func Between(s, sub1, sub2 string) string {
 }
 
 func Atoi(s string) (i int) {
-	i, _ = strconv.Atoi(s)
+	if s != "" {
+		i, _ = strconv.Atoi(s)
+	}
 	return
 }
 
diff --git a/pkg/rtmp/producer.go b/pkg/rtmp/producer.go
index c7e42c01..45665aa6 100644
--- a/pkg/rtmp/producer.go
+++ b/pkg/rtmp/producer.go
@@ -14,7 +14,7 @@ func Dial(rawURL string) (core.Producer, error) {
 		return nil, err
 	}
 
-	conn, err := tcp.Dial(u, "1935")
+	conn, err := tcp.Dial(u, "1935", core.ConnDialTimeout)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/rtsp/client.go b/pkg/rtsp/client.go
index 82518ede..186d69bb 100644
--- a/pkg/rtsp/client.go
+++ b/pkg/rtsp/client.go
@@ -31,7 +31,11 @@ func (c *Conn) Dial() (err error) {
 	var conn net.Conn
 
 	if c.Transport == "" {
-		conn, err = tcp.Dial(c.URL, "554")
+		timeout := core.ConnDialTimeout
+		if c.Timeout != 0 {
+			timeout = time.Second * time.Duration(c.Timeout)
+		}
+		conn, err = tcp.Dial(c.URL, "554", timeout)
 	} else {
 		conn, err = websocket.Dial(c.Transport)
 	}
diff --git a/pkg/rtsp/conn.go b/pkg/rtsp/conn.go
index 85bd53a8..df9b7ec2 100644
--- a/pkg/rtsp/conn.go
+++ b/pkg/rtsp/conn.go
@@ -26,6 +26,7 @@ type Conn struct {
 	Media       string
 	PacketSize  uint16
 	SessionName string
+	Timeout     int
 	Transport   string // custom transport support, ex. RTSP over WebSocket
 
 	Medias    []*core.Media
@@ -108,13 +109,17 @@ func (c *Conn) Handle() (err error) {
 		}
 		keepaliveTS = time.Now().Add(keepaliveDT)
 
-		// polling frames from remote RTSP Server (ex Camera)
-		if len(c.receivers) > 0 {
-			// if we receiving video/audio from camera
-			timeout = time.Second * 5
+		if c.Timeout == 0 {
+			// polling frames from remote RTSP Server (ex Camera)
+			if len(c.receivers) > 0 {
+				// if we receiving video/audio from camera
+				timeout = time.Second * 5
+			} else {
+				// if we only send audio to camera
+				timeout = time.Second * 30
+			}
 		} else {
-			// if we only send audio to camera
-			timeout = time.Second * 30
+			timeout = time.Second * time.Duration(c.Timeout)
 		}
 
 	case core.ModePassiveProducer:
diff --git a/pkg/tcp/dial.go b/pkg/tcp/dial.go
index bb604234..75d2f7a4 100644
--- a/pkg/tcp/dial.go
+++ b/pkg/tcp/dial.go
@@ -6,12 +6,11 @@ import (
 	"net"
 	"net/url"
 	"strings"
-
-	"github.com/AlexxIT/go2rtc/pkg/core"
+	"time"
 )
 
 // Dial - for RTSP(S|X) and RTMP(S|X)
-func Dial(u *url.URL, port string) (net.Conn, error) {
+func Dial(u *url.URL, port string, timeout time.Duration) (net.Conn, error) {
 	var hostname string // without port
 	if i := strings.IndexByte(u.Host, ':'); i > 0 {
 		hostname = u.Host[:i]
@@ -34,7 +33,7 @@ func Dial(u *url.URL, port string) (net.Conn, error) {
 		return nil, errors.New("unsupported scheme: " + u.Scheme)
 	}
 
-	conn, err := net.DialTimeout("tcp", u.Host, core.ConnDialTimeout)
+	conn, err := net.DialTimeout("tcp", u.Host, timeout)
 	if err != nil {
 		return nil, err
 	}