From ac74c7ed950e524b050c6e3349940e1a2714ce5d Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Thu, 19 Nov 2020 12:25:06 +0100
Subject: [PATCH] fix freeze when proxying streams with tcp

---
 go.mod                        |  2 +-
 go.sum                        |  4 ++--
 internal/sourcertmp/source.go | 14 +++-----------
 internal/sourcertsp/source.go | 22 ++++------------------
 4 files changed, 10 insertions(+), 32 deletions(-)

diff --git a/go.mod b/go.mod
index 4466d6b7..b8b467d3 100644
--- a/go.mod
+++ b/go.mod
@@ -5,7 +5,7 @@ go 1.15
 require (
 	github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 // indirect
 	github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d // indirect
-	github.com/aler9/gortsplib v0.0.0-20201116080759-cf6f734972f3
+	github.com/aler9/gortsplib v0.0.0-20201119110120-5019561d3fae
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/fsnotify/fsnotify v1.4.9
 	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51
diff --git a/go.sum b/go.sum
index 38bd2fbd..698adb42 100644
--- a/go.sum
+++ b/go.sum
@@ -2,8 +2,8 @@ github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 h1:JYp7IbQjafo
 github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d h1:UQZhZ2O0vMHr2cI+DC1Mbh0TJxzA3RcLoMsFw+aXw7E=
 github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d/go.mod h1:rBZYJk541a8SKzHPHnH3zbiI+7dagKZ0cgpgrD7Fyho=
-github.com/aler9/gortsplib v0.0.0-20201116080759-cf6f734972f3 h1:qgWLtuOANssYR0tuvA9Ty+T4lYPmNe2l6AJWhEIA6D0=
-github.com/aler9/gortsplib v0.0.0-20201116080759-cf6f734972f3/go.mod h1:6yKsTNIrCapRz90WHQtyFV/rKK0TT+QapxUXNqSJi9M=
+github.com/aler9/gortsplib v0.0.0-20201119110120-5019561d3fae h1:FF6+/D0sjbx90ayB6kR3OqFTrynC/2eLIOdY0jB5/io=
+github.com/aler9/gortsplib v0.0.0-20201119110120-5019561d3fae/go.mod h1:6yKsTNIrCapRz90WHQtyFV/rKK0TT+QapxUXNqSJi9M=
 github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
diff --git a/internal/sourcertmp/source.go b/internal/sourcertmp/source.go
index fab89aa6..37fc79ae 100644
--- a/internal/sourcertmp/source.go
+++ b/internal/sourcertmp/source.go
@@ -228,6 +228,7 @@ func (s *Source) runInner() bool {
 
 	s.parent.Log("rtmp source ready")
 	s.parent.OnSourceSetReady(tracks)
+	defer s.parent.OnSourceSetNotReady()
 
 	readerDone := make(chan error)
 	go func() {
@@ -286,26 +287,17 @@ func (s *Source) runInner() bool {
 		}
 	}()
 
-	var ret bool
-
-outer:
 	for {
 		select {
 		case <-s.terminate:
 			nconn.Close()
 			<-readerDone
-			ret = false
-			break outer
+			return false
 
 		case err := <-readerDone:
 			nconn.Close()
 			s.parent.Log("rtmp source ERR: %s", err)
-			ret = true
-			break outer
+			return true
 		}
 	}
-
-	s.parent.OnSourceSetNotReady()
-
-	return ret
 }
diff --git a/internal/sourcertsp/source.go b/internal/sourcertsp/source.go
index 5786326f..f76801b9 100644
--- a/internal/sourcertsp/source.go
+++ b/internal/sourcertsp/source.go
@@ -138,37 +138,23 @@ func (s *Source) runInner() bool {
 
 	s.parent.Log("rtsp source ready")
 	s.parent.OnSourceSetReady(tracks)
+	defer s.parent.OnSourceSetNotReady()
 
-	readerDone := make(chan error, 1)
-
-	conn.OnFrame(func(trackId int, streamType gortsplib.StreamType, content []byte, err error) {
-		if err != nil {
-			readerDone <- err
-			return
-		}
-
+	readerDone := conn.OnFrame(func(trackId int, streamType gortsplib.StreamType, content []byte) {
 		s.parent.OnFrame(trackId, streamType, content)
 	})
 
-	var ret bool
-
-outer:
 	for {
 		select {
 		case <-s.terminate:
 			conn.Close()
 			<-readerDone
-			ret = false
-			break outer
+			return false
 
 		case err := <-readerDone:
 			conn.Close()
 			s.parent.Log("rtsp source ERR: %s", err)
-			ret = true
-			break outer
+			return true
 		}
 	}
-
-	s.parent.OnSourceSetNotReady()
-	return ret
 }