From 98b521b058fe5345a0ec90c44af30ae0e8094843 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Sun, 1 Nov 2020 20:02:42 +0100
Subject: [PATCH] implement Dialer

---
 connclientdial.go                        | 108 ---------------
 dialer.go                                | 163 +++++++++++++++++++++++
 connclientdial_test.go => dialer_test.go |   8 +-
 3 files changed, 167 insertions(+), 112 deletions(-)
 delete mode 100644 connclientdial.go
 create mode 100644 dialer.go
 rename connclientdial_test.go => dialer_test.go (97%)

diff --git a/connclientdial.go b/connclientdial.go
deleted file mode 100644
index 87988f2d..00000000
--- a/connclientdial.go
+++ /dev/null
@@ -1,108 +0,0 @@
-package gortsplib
-
-import (
-	"github.com/aler9/gortsplib/base"
-)
-
-// DialRead connects to the address and starts reading all tracks.
-func DialRead(address string, proto StreamProtocol) (*ConnClient, error) {
-	u, err := base.ParseURL(address)
-	if err != nil {
-		return nil, err
-	}
-
-	conn, err := NewConnClient(ConnClientConf{Host: u.Host()})
-	if err != nil {
-		return nil, err
-	}
-
-	_, err = conn.Options(u)
-	if err != nil {
-		conn.Close()
-		return nil, err
-	}
-
-	tracks, _, err := conn.Describe(u)
-	if err != nil {
-		conn.Close()
-		return nil, err
-	}
-
-	if proto == StreamProtocolUDP {
-		for _, track := range tracks {
-			_, err := conn.SetupUDP(u, TransportModePlay, track, 0, 0)
-			if err != nil {
-				return nil, err
-			}
-		}
-
-	} else {
-		for _, track := range tracks {
-			_, err := conn.SetupTCP(u, TransportModePlay, track)
-			if err != nil {
-				conn.Close()
-				return nil, err
-			}
-		}
-	}
-
-	_, err = conn.Play(u)
-	if err != nil {
-		conn.Close()
-		return nil, err
-	}
-
-	return conn, nil
-}
-
-// DialPublish connects to the address and starts publishing the tracks.
-func DialPublish(address string, proto StreamProtocol, tracks Tracks) (*ConnClient, error) {
-	u, err := base.ParseURL(address)
-	if err != nil {
-		return nil, err
-	}
-
-	conn, err := NewConnClient(ConnClientConf{Host: u.Host()})
-	if err != nil {
-		return nil, err
-	}
-
-	_, err = conn.Options(u)
-	if err != nil {
-		conn.Close()
-		return nil, err
-	}
-
-	_, err = conn.Announce(u, tracks)
-	if err != nil {
-		conn.Close()
-		return nil, err
-	}
-
-	if proto == StreamProtocolUDP {
-		for _, track := range tracks {
-			_, err = conn.SetupUDP(u, TransportModeRecord, track, 0, 0)
-			if err != nil {
-				conn.Close()
-				return nil, err
-			}
-		}
-
-	} else {
-		for _, track := range tracks {
-			_, err = conn.SetupTCP(u, TransportModeRecord, track)
-			if err != nil {
-				conn.Close()
-				return nil, err
-			}
-		}
-	}
-
-	_, err = conn.Record(u)
-	if err != nil {
-		conn.Close()
-		return nil, err
-	}
-
-	return conn, nil
-}
diff --git a/dialer.go b/dialer.go
new file mode 100644
index 00000000..c4774112
--- /dev/null
+++ b/dialer.go
@@ -0,0 +1,163 @@
+package gortsplib
+
+import (
+	"net"
+	"time"
+
+	"github.com/aler9/gortsplib/base"
+)
+
+// DefaultDialer is the default dialer, used by DialRead and DialPublish.
+var DefaultDialer = &Dialer{}
+
+// DialRead connects to the address and starts reading all tracks.
+func DialRead(address string, proto StreamProtocol) (*ConnClient, error) {
+	return DefaultDialer.DialRead(address, proto)
+}
+
+// DialPublish connects to the address and starts publishing the tracks.
+func DialPublish(address string, proto StreamProtocol, tracks Tracks) (*ConnClient, error) {
+	return DefaultDialer.DialPublish(address, proto, tracks)
+}
+
+// Dialer allows to connect to a server and read or publish tracks.
+type Dialer struct {
+	// (optional) timeout of read operations.
+	// It defaults to 10 seconds
+	ReadTimeout time.Duration
+
+	// (optional) timeout of write operations.
+	// It defaults to 5 seconds
+	WriteTimeout time.Duration
+
+	// (optional) read buffer count.
+	// If greater than 1, allows to pass buffers to routines different than the one
+	// that is reading frames.
+	// It defaults to 1
+	ReadBufferCount int
+
+	// (optional) function used to initialize the TCP client.
+	// It defaults to net.DialTimeout
+	DialTimeout func(network, address string, timeout time.Duration) (net.Conn, error)
+
+	// (optional) function used to initialize UDP listeners.
+	// It defaults to net.ListenPacket
+	ListenPacket func(network, address string) (net.PacketConn, error)
+}
+
+// DialRead connects to the address and starts reading all tracks.
+func (d *Dialer) DialRead(address string, proto StreamProtocol) (*ConnClient, error) {
+	u, err := base.ParseURL(address)
+	if err != nil {
+		return nil, err
+	}
+
+	conn, err := NewConnClient(ConnClientConf{
+		Host:            u.Host(),
+		ReadTimeout:     d.ReadTimeout,
+		WriteTimeout:    d.WriteTimeout,
+		ReadBufferCount: d.ReadBufferCount,
+		DialTimeout:     d.DialTimeout,
+		ListenPacket:    d.ListenPacket,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = conn.Options(u)
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+
+	tracks, _, err := conn.Describe(u)
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+
+	if proto == StreamProtocolUDP {
+		for _, track := range tracks {
+			_, err := conn.SetupUDP(u, TransportModePlay, track, 0, 0)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+	} else {
+		for _, track := range tracks {
+			_, err := conn.SetupTCP(u, TransportModePlay, track)
+			if err != nil {
+				conn.Close()
+				return nil, err
+			}
+		}
+	}
+
+	_, err = conn.Play(u)
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+
+	return conn, nil
+}
+
+// DialPublish connects to the address and starts publishing the tracks.
+func (d *Dialer) DialPublish(address string, proto StreamProtocol, tracks Tracks) (*ConnClient, error) {
+	u, err := base.ParseURL(address)
+	if err != nil {
+		return nil, err
+	}
+
+	conn, err := NewConnClient(ConnClientConf{
+		Host:            u.Host(),
+		ReadTimeout:     d.ReadTimeout,
+		WriteTimeout:    d.WriteTimeout,
+		ReadBufferCount: d.ReadBufferCount,
+		DialTimeout:     d.DialTimeout,
+		ListenPacket:    d.ListenPacket,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = conn.Options(u)
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+
+	_, err = conn.Announce(u, tracks)
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+
+	if proto == StreamProtocolUDP {
+		for _, track := range tracks {
+			_, err = conn.SetupUDP(u, TransportModeRecord, track, 0, 0)
+			if err != nil {
+				conn.Close()
+				return nil, err
+			}
+		}
+
+	} else {
+		for _, track := range tracks {
+			_, err = conn.SetupTCP(u, TransportModeRecord, track)
+			if err != nil {
+				conn.Close()
+				return nil, err
+			}
+		}
+	}
+
+	_, err = conn.Record(u)
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+
+	return conn, nil
+}
diff --git a/connclientdial_test.go b/dialer_test.go
similarity index 97%
rename from connclientdial_test.go
rename to dialer_test.go
index af506730..65513671 100644
--- a/connclientdial_test.go
+++ b/dialer_test.go
@@ -58,7 +58,7 @@ func (c *container) wait() int {
 	return int(code)
 }
 
-func TestConnClientDialReadUDP(t *testing.T) {
+func TestDialerReadUDP(t *testing.T) {
 	cnt1, err := newContainer("rtsp-simple-server", "server", []string{})
 	require.NoError(t, err)
 	defer cnt1.close()
@@ -97,7 +97,7 @@ func TestConnClientDialReadUDP(t *testing.T) {
 	conn.CloseUDPListeners()
 }
 
-func TestConnClientDialReadTCP(t *testing.T) {
+func TestDialerReadTCP(t *testing.T) {
 	cnt1, err := newContainer("rtsp-simple-server", "server", []string{})
 	require.NoError(t, err)
 	defer cnt1.close()
@@ -129,7 +129,7 @@ func TestConnClientDialReadTCP(t *testing.T) {
 	require.Equal(t, StreamTypeRtp, typ)
 }
 
-func TestConnClientDialPublishUDP(t *testing.T) {
+func TestDialerPublishUDP(t *testing.T) {
 	for _, server := range []string{
 		"rtsp-simple-server",
 		"ffmpeg",
@@ -228,7 +228,7 @@ func TestConnClientDialPublishUDP(t *testing.T) {
 	}
 }
 
-func TestConnClientDialPublishTCP(t *testing.T) {
+func TestDialerPublishTCP(t *testing.T) {
 	for _, server := range []string{
 		"rtsp-simple-server",
 		"ffmpeg",