Repository: pion/mediadevices
Pull request #76 (2020-02-19 08:05:59)

Fix race condition in codecs
CGO based codecs caused segmentation fault due to the race condition
of Close() and Read().



Patch:
diff --git a/pkg/codec/openh264/openh264.go b/pkg/codec/openh264/openh264.go
index 4a2e5a3..cd6c810 100644
--- a/pkg/codec/openh264/openh264.go
+++ b/pkg/codec/openh264/openh264.go
@@ -13,6 +13,7 @@ import (
 	"fmt"
 	"image"
 	"io"
+	"sync"
 	"unsafe"
 
 	"github.com/pion/mediadevices/pkg/codec"
@@ -27,6 +28,9 @@ type encoder struct {
 	engine *C.Encoder
 	r      video.Reader
 	buff   []byte
+
+	mu     sync.Mutex
+	closed bool
 }
 
 var _ codec.VideoEncoderBuilder = codec.VideoEncoderBuilder(NewEncoder)
@@ -58,6 +62,13 @@ func NewEncoder(r video.Reader, p prop.Media) (io.ReadCloser, error) {
 }
 
 func (e *encoder) Read(p []byte) (n int, err error) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	if e.closed {
+		return 0, io.EOF
+	}
+
 	if e.buff != nil {
 		n, err = mio.Copy(p, e.buff)
 		if err == nil {
@@ -96,6 +107,11 @@ func (e *encoder) Read(p []byte) (n int, err error) {
 }
 
 func (e *encoder) Close() error {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	e.closed = true
+
 	C.enc_free(e.engine)
 	return nil
 }
diff --git a/pkg/codec/vpx/vpx.go b/pkg/codec/vpx/vpx.go
index b9ea017..b0d9aec 100644
--- a/pkg/codec/vpx/vpx.go
+++ b/pkg/codec/vpx/vpx.go
@@ -51,6 +51,7 @@ import (
 	"fmt"
 	"image"
 	"io"
+	"sync"
 	"time"
 	"unsafe"
 
@@ -72,6 +73,9 @@ type encoder struct {
 	tStart     int
 	tLastFrame int
 	frame      []byte
+
+	mu     sync.Mutex
+	closed bool
 }
 
 func init() {
@@ -139,6 +143,13 @@ func newEncoder(r video.Reader, p prop.Media, codecIface *C.vpx_codec_iface_t) (
 }
 
 func (e *encoder) Read(p []byte) (int, error) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	if e.closed {
+		return 0, io.EOF
+	}
+
 	if e.buff != nil {
 		n, err := mio.Copy(p, e.buff)
 		if err == nil {
@@ -204,6 +215,11 @@ func (e *encoder) Read(p []byte) (int, error) {
 }
 
 func (e *encoder) Close() error {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	e.closed = true
+
 	C.free(unsafe.Pointer(e.raw))
 	defer C.free(unsafe.Pointer(e.codec))
 
diff --git a/track.go b/track.go
index 6fb49aa..4b23348 100644
--- a/track.go
+++ b/track.go
@@ -4,7 +4,7 @@ import (
 	"fmt"
 	"io"
 	"math/rand"
-	"sync/atomic"
+	"sync"
 
 	"github.com/pion/mediadevices/pkg/codec"
 	"github.com/pion/mediadevices/pkg/driver"
@@ -19,6 +19,9 @@ type Tracker interface {
 	Track() *webrtc.Track
 	LocalTrack() LocalTrack
 	Stop()
+	// OnEnded registers a handler to receive an error from the media stream track.
+	// If the error is already occured before registering, the handler will be
+	// immediately called.
 	OnEnded(func(error))
 }
 
@@ -33,7 +36,10 @@ type track struct {
 	t LocalTrack
 	s *sampler
 
-	onErrorHandler atomic.Value // func(error)
+	onErrorHandler func(error)
+	err            error
+	mu             sync.Mutex
+	endOnce        sync.Once
 }
 
 func newTrack(codecs []*webrtc.RTPCodec, trackGenerator TrackGenerator, d driver.Driver, codecName string) (*track, error) {
@@ -66,13 +72,29 @@ func newTrack(codecs []*webrtc.RTPCodec, trackGenerator TrackGenerator, d driver
 }
 
 func (t *track) OnEnded(handler func(error)) {
-	t.onErrorHandler.Store(handler)
+	t.mu.Lock()
+	t.onErrorHandler = handler
+	err := t.err
+	t.mu.Unlock()
+
+	if err != nil && handler != nil {
+		// Already errored.
+		t.endOnce.Do(func() {
+			handler(err)
+		})
+	}
 }
 
 func (t *track) onError(err error) {
-	handler := t.onErrorHandler.Load()
+	t.mu.Lock()
+	t.err = err
+	handler := t.onErrorHandler
+	t.mu.Unlock()
+
 	if handler != nil {
-		handler.(func(error))(err)
+		t.endOnce.Do(func() {
+			handler(err)
+		})
 	}
 }
 


Test patch:
diff --git a/track_test.go b/track_test.go
new file mode 100644
index 0000000..baf67c5
--- /dev/null
+++ b/track_test.go
@@ -0,0 +1,55 @@
+package mediadevices
+
+import (
+	"errors"
+	"testing"
+	"time"
+)
+
+func TestOnEnded(t *testing.T) {
+	errExpected := errors.New("an error")
+
+	t.Run("ErrorAfterRegister", func(t *testing.T) {
+		tr := &track{}
+
+		called := make(chan error, 1)
+		tr.OnEnded(func(error) {
+			called <- errExpected
+		})
+		select {
+		case <-called:
+			t.Error("OnEnded handler is unexpectedly called")
+		case <-time.After(10 * time.Millisecond):
+		}
+
+		tr.onError(errExpected)
+
+		select {
+		case err := <-called:
+			if err != errExpected {
+				t.Errorf("Expected to receive error: %v, got: %v", errExpected, err)
+			}
+		case <-time.After(10 * time.Millisecond):
+			t.Error("Timeout")
+		}
+	})
+
+	t.Run("ErrorBeforeRegister", func(t *testing.T) {
+		tr := &track{}
+
+		tr.onError(errExpected)
+
+		called := make(chan error, 1)
+		tr.OnEnded(func(err error) {
+			called <- errExpected
+		})
+		select {
+		case err := <-called:
+			if err != errExpected {
+				t.Errorf("Expected to receive error: %v, got: %v", errExpected, err)
+			}
+		case <-time.After(10 * time.Millisecond):
+			t.Error("Timeout")
+		}
+	})
+}