fix: stream is nil course panic

a495433 · May 30, 2017 · f6a8749 · f6a8749
1 parent 09976a1
commit f6a8749
Showing 1 changed file with 17 additions and 14 deletions.
diff --git a/protocol/rtmp/stream.go b/protocol/rtmp/stream.go
@@ -27,22 +27,25 @@ func NewRtmpStream() *RtmpStream {
 
 func (rs *RtmpStream) HandleReader(r av.ReadCloser) {
 	info := r.Info()
-	var s *Stream
-	ok := rs.streams.Has(info.Key)
-	if !ok {
-		s = NewStream()
-		rs.streams.Set(info.Key, s)
-	} else {
-		s.TransStop()
-		id := s.ID()
-		if id != EmptyID && id != info.UID {
-			ns := NewStream()
-			s.Copy(ns)
-			s = ns
-			rs.streams.Set(info.Key, ns)
+	var stream *Stream
+	i, ok := rs.streams.Get(info.Key)
+	if ok {
+		if s, ok := i.(*Stream); ok {
+			s.TransStop()
+			id := s.ID()
+			if id != EmptyID && id != info.UID {
+				ns := NewStream()
+				s.Copy(ns)
+				s = ns
+				rs.streams.Set(info.Key, ns)
+			}
 		}
+	} else {
+		stream = NewStream()
+		rs.streams.Set(info.Key, stream)
 	}
-	s.AddReader(r)
+
+	stream.AddReader(r)
 }
 
 func (rs *RtmpStream) HandleWriter(w av.WriteCloser) {