Add AV1 Support

Also include example

Resolves #1670

examples/save-to-disk-av1/README.md

@@ -0,0 +1,35 @@
+# save-to-disk-av1
+save-to-disk-av1 is a simple application that shows how to save a video to disk using AV1.
+
+If you wish to save VP8 and Opus instead of AV1 see [save-to-disk](https://github.com/pion/webrtc/tree/master/examples/save-to-disk)
+
+If you wish to save VP8/Opus inside the same file see [save-to-webm](https://github.com/pion/example-webrtc-applications/tree/master/save-to-webm)
+
+## Instructions
+### Download save-to-disk-av1
+```
+export GO111MODULE=on
+go get github.com/pion/webrtc/v3/examples/save-to-disk-av1
+```
+
+### Open save-to-disk-av1 example page
+[jsfiddle.net](https://jsfiddle.net/xjcve6d3/) you should see your Webcam, two text-areas and two buttons: `Copy browser SDP to clipboard`, `Start Session`.
+
+### Run save-to-disk-av1, with your browsers SessionDescription as stdin
+In the jsfiddle the top textarea is your browser's Session Description. Press `Copy browser SDP to clipboard` or copy the base64 string manually.
+We will use this value in the next step.
+
+#### Linux/macOS
+Run `echo $BROWSER_SDP | save-to-disk-av1`
+#### Windows
+1. Paste the SessionDescription into a file.
+1. Run `save-to-disk-av1 < my_file`
+
+### Input save-to-disk-av1's SessionDescription into your browser
+Copy the text that `save-to-disk-av1` just emitted and copy into second text area
+
+### Hit 'Start Session' in jsfiddle, wait, close jsfiddle, enjoy your video!
+In the folder you ran `save-to-disk-av1` you should now have a file `output.ivf` play with your video player of choice!
+> Note: In order to correctly create the files, the remote client (JSFiddle) should be closed. The Go example will automatically close itself.
+
+Congrats, you have used Pion WebRTC! Now start building something cool

examples/save-to-disk-av1/main.go

@@ -0,0 +1,165 @@
+//go:build !js
+// +build !js
+
+package main
+
+import (
+	"fmt"
+	"os"
+	"strings"
+	"time"
+
+	"github.com/pion/interceptor"
+	"github.com/pion/rtcp"
+	"github.com/pion/webrtc/v3"
+	"github.com/pion/webrtc/v3/examples/internal/signal"
+	"github.com/pion/webrtc/v3/pkg/media"
+	"github.com/pion/webrtc/v3/pkg/media/ivfwriter"
+)
+
+func saveToDisk(i media.Writer, track *webrtc.TrackRemote) {
+	defer func() {
+		if err := i.Close(); err != nil {
+			panic(err)
+		}
+	}()
+
+	for {
+		rtpPacket, _, err := track.ReadRTP()
+		if err != nil {
+			panic(err)
+		}
+		if err := i.WriteRTP(rtpPacket); err != nil {
+			panic(err)
+		}
+	}
+}
+
+func main() {
+	// Everything below is the Pion WebRTC API! Thanks for using it ❤️.
+
+	// Create a MediaEngine object to configure the supported codec
+	m := &webrtc.MediaEngine{}
+
+	// Setup the codecs you want to use.
+	// We'll use a VP8 and Opus but you can also define your own
+	if err := m.RegisterCodec(webrtc.RTPCodecParameters{
+		RTPCodecCapability: webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeAV1, ClockRate: 90000, Channels: 0, SDPFmtpLine: "", RTCPFeedback: nil},
+		PayloadType:        96,
+	}, webrtc.RTPCodecTypeVideo); err != nil {
+		panic(err)
+	}
+
+	// Create a InterceptorRegistry. This is the user configurable RTP/RTCP Pipeline.
+	// This provides NACKs, RTCP Reports and other features. If you use `webrtc.NewPeerConnection`
+	// this is enabled by default. If you are manually managing You MUST create a InterceptorRegistry
+	// for each PeerConnection.
+	i := &interceptor.Registry{}
+
+	// Use the default set of Interceptors
+	if err := webrtc.RegisterDefaultInterceptors(m, i); err != nil {
+		panic(err)
+	}
+
+	// Create the API object with the MediaEngine
+	api := webrtc.NewAPI(webrtc.WithMediaEngine(m), webrtc.WithInterceptorRegistry(i))
+
+	// Prepare the configuration
+	config := webrtc.Configuration{}
+
+	// Create a new RTCPeerConnection
+	peerConnection, err := api.NewPeerConnection(config)
+	if err != nil {
+		panic(err)
+	}
+
+	// Allow us to receive 1 video track
+	if _, err = peerConnection.AddTransceiverFromKind(webrtc.RTPCodecTypeVideo); err != nil {
+		panic(err)
+	}
+
+	ivfFile, err := ivfwriter.New("output.ivf", ivfwriter.WithCodec(webrtc.MimeTypeAV1))
+	if err != nil {
+		panic(err)
+	}
+
+	// Set a handler for when a new remote track starts, this handler saves buffers to disk as
+	// an ivf file, since we could have multiple video tracks we provide a counter.
+	// In your application this is where you would handle/process video
+	peerConnection.OnTrack(func(track *webrtc.TrackRemote, receiver *webrtc.RTPReceiver) {
+		// Send a PLI on an interval so that the publisher is pushing a keyframe every rtcpPLIInterval
+		go func() {
+			ticker := time.NewTicker(time.Second * 3)
+			for range ticker.C {
+				errSend := peerConnection.WriteRTCP([]rtcp.Packet{&rtcp.PictureLossIndication{MediaSSRC: uint32(track.SSRC())}})
+				if errSend != nil {
+					fmt.Println(errSend)
+				}
+			}
+		}()
+
+		if strings.EqualFold(track.Codec().MimeType, webrtc.MimeTypeAV1) {
+			fmt.Println("Got AV1 track, saving to disk as output.ivf")
+			saveToDisk(ivfFile, track)
+		}
+	})
+
+	// Set the handler for ICE connection state
+	// This will notify you when the peer has connected/disconnected
+	peerConnection.OnICEConnectionStateChange(func(connectionState webrtc.ICEConnectionState) {
+		fmt.Printf("Connection State has changed %s \n", connectionState.String())
+
+		if connectionState == webrtc.ICEConnectionStateConnected {
+			fmt.Println("Ctrl+C the remote client to stop the demo")
+		} else if connectionState == webrtc.ICEConnectionStateFailed {
+			if closeErr := ivfFile.Close(); closeErr != nil {
+				panic(closeErr)
+			}
+
+			fmt.Println("Done writing media files")
+
+			// Gracefully shutdown the peer connection
+			if closeErr := peerConnection.Close(); closeErr != nil {
+				panic(closeErr)
+			}
+
+			os.Exit(0)
+		}
+	})
+
+	// Wait for the offer to be pasted
+	offer := webrtc.SessionDescription{}
+	signal.Decode(signal.MustReadStdin(), &offer)
+
+	// Set the remote SessionDescription
+	err = peerConnection.SetRemoteDescription(offer)
+	if err != nil {
+		panic(err)
+	}
+
+	// Create answer
+	answer, err := peerConnection.CreateAnswer(nil)
+	if err != nil {
+		panic(err)
+	}
+
+	// Create channel that is blocked until ICE Gathering is complete
+	gatherComplete := webrtc.GatheringCompletePromise(peerConnection)
+
+	// Sets the LocalDescription, and starts our UDP listeners
+	err = peerConnection.SetLocalDescription(answer)
+	if err != nil {
+		panic(err)
+	}
+
+	// Block until ICE Gathering is complete, disabling trickle ICE
+	// we do this because we only can exchange one signaling message
+	// in a production application you should exchange ICE Candidates via OnICECandidate
+	<-gatherComplete
+
+	// Output the answer in base64 so we can paste it in browser
+	fmt.Println(signal.Encode(*peerConnection.LocalDescription()))
+
+	// Block forever
+	select {}
+}

go.mod

@@ -12,7 +12,7 @@ require (
 	github.com/pion/logging v0.2.2
 	github.com/pion/randutil v0.1.0
 	github.com/pion/rtcp v1.2.9
-	github.com/pion/rtp v1.7.11
+	github.com/pion/rtp v1.7.12
 	github.com/pion/sctp v1.8.2
 	github.com/pion/sdp/v3 v3.0.4
 	github.com/pion/srtp/v2 v2.0.5

go.sum

@@ -59,8 +59,8 @@ github.com/pion/rtcp v1.2.9 h1:1ujStwg++IOLIEoOiIQ2s+qBuJ1VN81KW+9pMPsif+U=
 github.com/pion/rtcp v1.2.9/go.mod h1:qVPhiCzAm4D/rxb6XzKeyZiQK69yJpbUDJSF7TgrqNo=
 github.com/pion/rtp v1.7.0/go.mod h1:bDb5n+BFZxXx0Ea7E5qe+klMuqiBrP+w8XSjiWtCUko=
 github.com/pion/rtp v1.7.4/go.mod h1:bDb5n+BFZxXx0Ea7E5qe+klMuqiBrP+w8XSjiWtCUko=
-github.com/pion/rtp v1.7.11 h1:WosqH088pRIAnAoAGZjagA1H3uFtzjyD5yagQXqZ3uo=
-github.com/pion/rtp v1.7.11/go.mod h1:bDb5n+BFZxXx0Ea7E5qe+klMuqiBrP+w8XSjiWtCUko=
+github.com/pion/rtp v1.7.12 h1:Wtrx1btLYn96vQGx35UTpgRBG/MGJmIHvrGND1m219A=
+github.com/pion/rtp v1.7.12/go.mod h1:bDb5n+BFZxXx0Ea7E5qe+klMuqiBrP+w8XSjiWtCUko=
 github.com/pion/sctp v1.8.0/go.mod h1:xFe9cLMZ5Vj6eOzpyiKjT9SwGM4KpK/8Jbw5//jc+0s=
 github.com/pion/sctp v1.8.2 h1:yBBCIrUMJ4yFICL3RIvR4eh/H2BTTvlligmSTy+3kiA=
 github.com/pion/sctp v1.8.2/go.mod h1:xFe9cLMZ5Vj6eOzpyiKjT9SwGM4KpK/8Jbw5//jc+0s=

mediaengine.go

@@ -615,6 +615,8 @@ func payloaderForCodec(codec RTPCodecCapability) (rtp.Payloader, error) {
 		}, nil
 	case strings.ToLower(MimeTypeVP9):
 		return &codecs.VP9Payloader{}, nil
+	case strings.ToLower(MimeTypeAV1):
+		return &codecs.AV1Payloader{}, nil
 	case strings.ToLower(MimeTypeG722):
 		return &codecs.G722Payloader{}, nil
 	case strings.ToLower(MimeTypePCMU), strings.ToLower(MimeTypePCMA):

pkg/media/ivfwriter/ivfwriter.go

@@ -9,11 +9,21 @@ import (
 
 	"github.com/pion/rtp"
 	"github.com/pion/rtp/codecs"
+	"github.com/pion/rtp/pkg/frame"
 )
 
 var (
 	errFileNotOpened    = errors.New("file not opened")
 	errInvalidNilPacket = errors.New("invalid nil packet")
+	errCodecAlreadySet  = errors.New("codec is already set")
+	errNoSuchCodec      = errors.New("no codec for this MimeType")
+)
+
+const (
+	mimeTypeVP8 = "video/VP8"
+	mimeTypeAV1 = "video/AV1"
+
+	ivfFileHeaderSignature = "DKIF"
 )
 
 // IVFWriter is used to take RTP packets and write them to an IVF on disk
@@ -21,16 +31,23 @@ type IVFWriter struct {
 	ioWriter     io.Writer
 	count        uint64
 	seenKeyFrame bool
+
+	isVP8, isAV1 bool
+
+	// VP8
 	currentFrame []byte
+
+	// AV1
+	av1Frame frame.AV1
 }
 
 // New builds a new IVF writer
-func New(fileName string) (*IVFWriter, error) {
+func New(fileName string, opts ...Option) (*IVFWriter, error) {
 	f, err := os.Create(fileName)
 	if err != nil {
 		return nil, err
 	}
-	writer, err := NewWith(f)
+	writer, err := NewWith(f, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -39,7 +56,7 @@ func New(fileName string) (*IVFWriter, error) {
 }
 
 // NewWith initialize a new IVF writer with an io.Writer output
-func NewWith(out io.Writer) (*IVFWriter, error) {
+func NewWith(out io.Writer, opts ...Option) (*IVFWriter, error) {
 	if out == nil {
 		return nil, errFileNotOpened
 	}
@@ -48,6 +65,17 @@ func NewWith(out io.Writer) (*IVFWriter, error) {
 		ioWriter:     out,
 		seenKeyFrame: false,
 	}
+
+	for _, o := range opts {
+		if err := o(writer); err != nil {
+			return nil, err
+		}
+	}
+
+	if !writer.isAV1 && !writer.isVP8 {
+		writer.isVP8 = true
+	}
+
 	if err := writer.writeHeader(); err != nil {
 		return nil, err
 	}
@@ -56,10 +84,17 @@ func NewWith(out io.Writer) (*IVFWriter, error) {
 
 func (i *IVFWriter) writeHeader() error {
 	header := make([]byte, 32)
-	copy(header[0:], "DKIF")                        // DKIF
+	copy(header[0:], ivfFileHeaderSignature)      // DKIF
 	binary.LittleEndian.PutUint16(header[4:], 0)  // Version
 	binary.LittleEndian.PutUint16(header[6:], 32) // Header size
-	copy(header[8:], "VP80")                        // FOURCC
+
+	// FOURCC
+	if i.isVP8 {
+		copy(header[8:], "VP80")
+	} else if i.isAV1 {
+		copy(header[8:], "AV01")
+	}
+
 	binary.LittleEndian.PutUint16(header[12:], 640) // Width in pixels
 	binary.LittleEndian.PutUint16(header[14:], 480) // Height in pixels
 	binary.LittleEndian.PutUint32(header[16:], 30)  // Framerate denominator
@@ -71,15 +106,28 @@ func (i *IVFWriter) writeHeader() error {
 	return err
 }
 
+func (i *IVFWriter) writeFrame(frame []byte) error {
+	frameHeader := make([]byte, 12)
+	binary.LittleEndian.PutUint32(frameHeader[0:], uint32(len(frame))) // Frame length
+	binary.LittleEndian.PutUint64(frameHeader[4:], i.count)            // PTS
+	i.count++
+
+	if _, err := i.ioWriter.Write(frameHeader); err != nil {
+		return err
+	}
+	_, err := i.ioWriter.Write(frame)
+	return err
+}
+
 // WriteRTP adds a new packet and writes the appropriate headers for it
 func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error {
 	if i.ioWriter == nil {
 		return errFileNotOpened
-	}
-	if len(packet.Payload) == 0 {
+	} else if len(packet.Payload) == 0 {
 		return nil
 	}
 
+	if i.isVP8 {
 		vp8Packet := codecs.VP8Packet{}
 		if _, err := vp8Packet.Unmarshal(packet.Payload); err != nil {
 			return err
@@ -102,19 +150,28 @@ func (i *IVFWriter) WriteRTP(packet *rtp.Packet) error {
 			return nil
 		}
 
-	frameHeader := make([]byte, 12)
-	binary.LittleEndian.PutUint32(frameHeader[0:], uint32(len(i.currentFrame))) // Frame length
-	binary.LittleEndian.PutUint64(frameHeader[4:], i.count)                     // PTS
-
-	i.count++
-
-	if _, err := i.ioWriter.Write(frameHeader); err != nil {
+		if err := i.writeFrame(i.currentFrame); err != nil {
 			return err
-	} else if _, err := i.ioWriter.Write(i.currentFrame); err != nil {
+		}
+		i.currentFrame = nil
+	} else if i.isAV1 {
+		av1Packet := &codecs.AV1Packet{}
+		if _, err := av1Packet.Unmarshal(packet.Payload); err != nil {
 			return err
 		}
 
-	i.currentFrame = nil
+		obus, err := i.av1Frame.ReadFrames(av1Packet)
+		if err != nil {
+			return err
+		}
+
+		for j := range obus {
+			if err := i.writeFrame(obus[j]); err != nil {
+				return err
+			}
+		}
+	}
+
 	return nil
 }
 
@@ -148,3 +205,26 @@ func (i *IVFWriter) Close() error {
 
 	return nil
 }
+
+// An Option configures a SampleBuilder.
+type Option func(i *IVFWriter) error
+
+// WithCodec configures if IVFWriter is writing AV1 or VP8 packets to disk
+func WithCodec(mimeType string) Option {
+	return func(i *IVFWriter) error {
+		if i.isVP8 || i.isAV1 {
+			return errCodecAlreadySet
+		}
+
+		switch mimeType {
+		case mimeTypeVP8:
+			i.isVP8 = true
+		case mimeTypeAV1:
+			i.isAV1 = true
+		default:
+			return errNoSuchCodec
+		}
+
+		return nil
+	}
+}

pkg/media/ivfwriter/ivfwriter_test.go

@@ -187,7 +187,7 @@ func TestIVFWriter_VP8(t *testing.T) {
 	}
 
 	// first test tries to write a valid VP8 packet
-	writer, err := NewWith(addPacketTestCase[0].buffer)
+	writer, err := NewWith(addPacketTestCase[0].buffer, WithCodec(mimeTypeVP8))
 	assert.Nil(err, "IVFWriter should be created")
 	assert.NotNil(writer, "Writer shouldn't be nil")
 	assert.False(writer.seenKeyFrame, "Writer's seenKeyFrame should initialize false")
@@ -239,3 +239,60 @@ func TestIVFWriter_EmptyPayload(t *testing.T) {
 
 	assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{}}))
 }
+
+func TestIVFWriter_Errors(t *testing.T) {
+	// Creating a Writer with AV1 and VP8
+	_, err := NewWith(&bytes.Buffer{}, WithCodec(mimeTypeAV1), WithCodec(mimeTypeAV1))
+	assert.ErrorIs(t, err, errCodecAlreadySet)
+
+	// Creating a Writer with Invalid Codec
+	_, err = NewWith(&bytes.Buffer{}, WithCodec(""))
+	assert.ErrorIs(t, err, errNoSuchCodec)
+}
+
+func TestIVFWriter_AV1(t *testing.T) {
+	t.Run("Unfragmented", func(t *testing.T) {
+		buffer := &bytes.Buffer{}
+
+		writer, err := NewWith(buffer, WithCodec(mimeTypeAV1))
+		assert.NoError(t, err)
+
+		assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x00, 0x01, 0xFF}}))
+		assert.NoError(t, writer.Close())
+		assert.Equal(t, buffer.Bytes(), []byte{
+			0x44, 0x4b, 0x49, 0x46, 0x0, 0x0, 0x20,
+			0x0, 0x41, 0x56, 0x30, 0x31, 0x80, 0x2,
+			0xe0, 0x1, 0x1e, 0x0, 0x0, 0x0, 0x1, 0x0,
+			0x0, 0x0, 0x84, 0x3, 0x0, 0x0, 0x0, 0x0,
+			0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+			0x0, 0x0, 0x0, 0x0, 0x0, 0xff,
+		})
+	})
+
+	t.Run("Fragmented", func(t *testing.T) {
+		buffer := &bytes.Buffer{}
+
+		writer, err := NewWith(buffer, WithCodec(mimeTypeAV1))
+		assert.NoError(t, err)
+
+		for _, p := range [][]byte{{0x40, 0x02, 0x00, 0x01}, {0xc0, 0x02, 0x02, 0x03}, {0xc0, 0x02, 0x04, 0x04}} {
+			assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: p}))
+			assert.Equal(t, buffer.Bytes(), []byte{
+				0x44, 0x4b, 0x49, 0x46, 0x0,
+				0x0, 0x20, 0x0, 0x41, 0x56, 0x30,
+				0x31, 0x80, 0x2, 0xe0, 0x1, 0x1e,
+				0x0, 0x0, 0x0, 0x1, 0x0, 0x0,
+				0x0, 0x84, 0x3, 0x0, 0x0, 0x0, 0x0,
+				0x0, 0x0,
+			})
+		}
+		assert.NoError(t, writer.WriteRTP(&rtp.Packet{Payload: []byte{0x80, 0x01, 0x05}}))
+		assert.Equal(t, buffer.Bytes(), []byte{
+			0x44, 0x4b, 0x49, 0x46, 0x0, 0x0, 0x20, 0x0, 0x41, 0x56, 0x30, 0x31, 0x80,
+			0x2, 0xe0, 0x1, 0x1e, 0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x84, 0x3, 0x0, 0x0,
+			0x0, 0x0, 0x0, 0x0, 0x7, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+			0x0, 0x1, 0x2, 0x3, 0x4, 0x4, 0x5,
+		})
+		assert.NoError(t, writer.Close())
+	})
+}