Add the ability to close ReadStreams

This allows us to close/remove ReadStreams from a session at anytime.

Resolves #3

.golangci.yml

@@ -11,6 +11,7 @@ linters:
     - dupl
     - gocritic
     - gochecknoglobals
+    - maligned
 
 issues:
   exclude-use-default: false

session.go

@@ -65,6 +65,17 @@ func (s *session) getOrCreateReadStream(ssrc uint32, child streamSession, proto
 	return r, false
 }
 
+func (s *session) removeReadStream(ssrc uint32) {
+	s.readStreamsLock.Lock()
+	defer s.readStreamsLock.Unlock()
+
+	if s.readStreamsClosed {
+		return
+	}
+
+	delete(s.readStreams, ssrc)
+}
+
 func (s *session) close() error {
 	if s.nextConn == nil {
 		return nil

session_srtcp.go

@@ -131,6 +131,10 @@ func (s *SessionSRTCP) decrypt(buf []byte) error {
 				return fmt.Errorf("failed to get/create ReadStreamSRTP")
 			}
 
+			// Ensure that readStream.Close() isn't called while in flight
+			readStream.mu.Lock()
+			defer readStream.mu.Unlock()
+
 			readBuf := <-readStream.readCh
 			if len(readBuf) < len(decrypted) {
 				return fmt.Errorf("input buffer was not long enough to contain decrypted RTCP")

session_srtp.go

@@ -116,6 +116,10 @@ func (s *SessionSRTP) decrypt(buf []byte) error {
 		return fmt.Errorf("failed to get/create ReadStreamSRTP")
 	}
 
+	// Ensure that readStream.Close() isn't called while in flight
+	readStream.mu.Lock()
+	defer readStream.mu.Unlock()
+
 	readBuf := <-readStream.readCh
 	decrypted, err := s.remoteContext.decryptRTP(readBuf, buf, h)
 	if err != nil {

stream_srtcp.go

@@ -2,6 +2,7 @@ package srtp
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/pions/rtcp"
 )
@@ -13,6 +14,11 @@ type readResultSRTCP struct {
 
 // ReadStreamSRTCP handles decryption for a single RTCP SSRC
 type ReadStreamSRTCP struct {
+	mu sync.Mutex
+
+	isInited bool
+	isClosed chan bool
+
 	session   *SessionSRTCP
 	ssrc      uint32
 	readCh    chan []byte
@@ -25,12 +31,18 @@ func (r *ReadStreamSRTCP) ReadRTCP(payload []byte) (int, *rtcp.Header, error) {
 	case <-r.session.closed:
 		return 0, nil, fmt.Errorf("SRTCP session is closed")
 	case r.readCh <- payload:
+	case <-r.isClosed:
+		return 0, nil, fmt.Errorf("SRTCP read stream is closed")
 	}
 
 	select {
 	case <-r.session.closed:
 		return 0, nil, fmt.Errorf("SRTCP session is closed")
-	case res := <-r.readRetCh:
+	case res, ok := <-r.readRetCh:
+		if !ok {
+			return 0, nil, fmt.Errorf("SRTCP read stream is closed")
+		}
+
 		return res.len, res.header, nil
 	}
 }
@@ -41,26 +53,56 @@ func (r *ReadStreamSRTCP) Read(b []byte) (int, error) {
 	case <-r.session.closed:
 		return 0, fmt.Errorf("SRTCP session is closed")
 	case r.readCh <- b:
+	case <-r.isClosed:
+		return 0, fmt.Errorf("SRTCP read stream is closed")
 	}
 
 	select {
 	case <-r.session.closed:
 		return 0, fmt.Errorf("SRTCP session is closed")
-	case res := <-r.readRetCh:
+	case res, ok := <-r.readRetCh:
+		if !ok {
+			return 0, fmt.Errorf("SRTCP read stream is closed")
+		}
 		return res.len, nil
 	}
 }
 
+// Close removes the ReadStream from the session and cleans up any associated state
+func (r *ReadStreamSRTCP) Close() error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if !r.isInited {
+		return fmt.Errorf("ReadStreamSRTCP has not been inited")
+	}
+
+	select {
+	case <-r.isClosed:
+		return fmt.Errorf("ReadStreamSRTCP is already closed")
+	default:
+		close(r.readRetCh)
+		r.session.removeReadStream(r.ssrc)
+		return nil
+	}
+}
+
 func (r *ReadStreamSRTCP) init(child streamSession, ssrc uint32) error {
 	sessionSRTCP, ok := child.(*SessionSRTCP)
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
 	if !ok {
 		return fmt.Errorf("ReadStreamSRTCP init failed type assertion")
+	} else if r.isInited {
+		return fmt.Errorf("ReadStreamSRTCP has already been inited")
 	}
 
 	r.session = sessionSRTCP
 	r.ssrc = ssrc
 	r.readCh = make(chan []byte)
 	r.readRetCh = make(chan readResultSRTCP)
+	r.isInited = true
+	r.isClosed = make(chan bool)
 	return nil
 }
 
@@ -84,7 +126,7 @@ func (w *WriteStreamSRTCP) WriteRTCP(header *rtcp.Header, payload []byte) (int,
 	return w.session.write(append(headerRaw, payload...))
 }
 
-// Write encrypts and writes a full RTP packets to the nextConn
+// Write encrypts and writes a full RTCP packets to the nextConn
 func (w *WriteStreamSRTCP) Write(b []byte) (int, error) {
 	return w.session.write(b)
 }

stream_srtp.go

@@ -2,6 +2,7 @@ package srtp
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/pions/rtp"
 )
@@ -13,6 +14,11 @@ type readResultSRTP struct {
 
 // ReadStreamSRTP handles decryption for a single RTP SSRC
 type ReadStreamSRTP struct {
+	mu sync.Mutex
+
+	isInited bool
+	isClosed chan bool
+
 	session   *SessionSRTP
 	ssrc      uint32
 	readCh    chan []byte
@@ -25,12 +31,17 @@ func (r *ReadStreamSRTP) ReadRTP(payload []byte) (int, *rtp.Header, error) {
 	case <-r.session.closed:
 		return 0, nil, fmt.Errorf("SRTP session is closed")
 	case r.readCh <- payload:
+	case <-r.isClosed:
+		return 0, nil, fmt.Errorf("SRTP read stream is closed")
 	}
 
 	select {
 	case <-r.session.closed:
 		return 0, nil, fmt.Errorf("SRTP session is closed")
-	case res := <-r.readRetCh:
+	case res, ok := <-r.readRetCh:
+		if !ok {
+			return 0, nil, fmt.Errorf("SRTP read stream is closed")
+		}
 		return res.len, res.header, nil
 	}
 }
@@ -41,26 +52,56 @@ func (r *ReadStreamSRTP) Read(b []byte) (int, error) {
 	case <-r.session.closed:
 		return 0, fmt.Errorf("SRTP session is closed")
 	case r.readCh <- b:
+	case <-r.isClosed:
+		return 0, fmt.Errorf("SRTP read stream is closed")
 	}
 
 	select {
 	case <-r.session.closed:
 		return 0, fmt.Errorf("SRTP session is closed")
-	case res := <-r.readRetCh:
+	case res, ok := <-r.readRetCh:
+		if !ok {
+			return 0, fmt.Errorf("SRTP read stream is closed")
+		}
 		return res.len, nil
 	}
 }
 
+// Close removes the ReadStream from the session and cleans up any associated state
+func (r *ReadStreamSRTP) Close() error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if !r.isInited {
+		return fmt.Errorf("ReadStreamSRTP has not been inited")
+	}
+
+	select {
+	case <-r.isClosed:
+		return fmt.Errorf("ReadStreamSRTP is already closed")
+	default:
+		close(r.readRetCh)
+		r.session.removeReadStream(r.ssrc)
+		return nil
+	}
+}
+
 func (r *ReadStreamSRTP) init(child streamSession, ssrc uint32) error {
 	sessionSRTP, ok := child.(*SessionSRTP)
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
 	if !ok {
 		return fmt.Errorf("ReadStreamSRTP init failed type assertion")
+	} else if r.isInited {
+		return fmt.Errorf("ReadStreamSRTP has already been inited")
 	}
 
 	r.session = sessionSRTP
 	r.ssrc = ssrc
 	r.readCh = make(chan []byte)
 	r.readRetCh = make(chan readResultSRTP)
+	r.isInited = true
+	r.isClosed = make(chan bool)
 	return nil
 }