From 7ec753f8a361a9d4e590b5868d719cfc805669cd Mon Sep 17 00:00:00 2001
From: Alessandro Ros <aler9.dev@gmail.com>
Date: Thu, 4 May 2023 20:46:48 +0200
Subject: [PATCH] support reading AV1 tracks with WebRTC (#1768)

---
 README.md                    |  2 +-
 internal/core/webrtc_conn.go | 58 ++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index f5ba25ec..09f9bd71 100644
--- a/README.md
+++ b/README.md
@@ -25,7 +25,7 @@ And can be read from the server with:
 |RTSP|UDP, UDP-Multicast, TCP, RTSPS|AV1, VP9, VP8, H265, H264, MPEG-4 Video (H263, Xvid), MPEG-2 Video, M-JPEG and any RTP-compatible codec|Opus,  MPEG-4 Audio (AAC), MPEG-2 Audio (MP3), G722, G711, LPCM and any RTP-compatible codec|
 |RTMP|RTMP, RTMPS, Enhanced RTMP|H264|MPEG-4 Audio (AAC), MPEG-2 Audio (MP3)|
 |HLS|Low-Latency HLS, MP4-based HLS, legacy HLS|H265, H264|Opus, MPEG-4 Audio (AAC)|
-|WebRTC||VP9, VP8, H264|Opus, G722, G711|
+|WebRTC||AV1, VP9, VP8, H264|Opus, G722, G711|
 
 Features:
 
diff --git a/internal/core/webrtc_conn.go b/internal/core/webrtc_conn.go
index f5776221..25af5240 100644
--- a/internal/core/webrtc_conn.go
+++ b/internal/core/webrtc_conn.go
@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/bluenviron/gortsplib/v3/pkg/formats"
+	"github.com/bluenviron/gortsplib/v3/pkg/formats/rtpav1"
 	"github.com/bluenviron/gortsplib/v3/pkg/formats/rtph264"
 	"github.com/bluenviron/gortsplib/v3/pkg/formats/rtpvp8"
 	"github.com/bluenviron/gortsplib/v3/pkg/formats/rtpvp9"
@@ -51,6 +52,18 @@ func newPeerConnection(configuration webrtc.Configuration,
 		return nil, err
 	}
 
+	err := m.RegisterCodec(webrtc.RTPCodecParameters{
+		RTPCodecCapability: webrtc.RTPCodecCapability{
+			MimeType:  webrtc.MimeTypeAV1,
+			ClockRate: 90000,
+		},
+		PayloadType: 96,
+	},
+		webrtc.RTPCodecTypeVideo)
+	if err != nil {
+		return nil, err
+	}
+
 	i := &interceptor.Registry{}
 	if err := webrtc.RegisterDefaultInterceptors(m, i); err != nil {
 		return nil, err
@@ -560,6 +573,51 @@ outer:
 }
 
 func (c *webRTCConn) createVideoTrack(medias media.Medias) (*webRTCTrack, error) {
+	var av1Format *formats.AV1
+	av1Media := medias.FindFormat(&av1Format)
+
+	if av1Format != nil {
+		webRTCTrak, err := webrtc.NewTrackLocalStaticRTP(
+			webrtc.RTPCodecCapability{
+				MimeType:  webrtc.MimeTypeAV1,
+				ClockRate: 90000,
+			},
+			"av1",
+			"rtspss",
+		)
+		if err != nil {
+			return nil, err
+		}
+
+		encoder := &rtpav1.Encoder{
+			PayloadType:    96,
+			PayloadMaxSize: webrtcPayloadMaxSize,
+		}
+		encoder.Init()
+
+		return &webRTCTrack{
+			media:       av1Media,
+			format:      av1Format,
+			webRTCTrack: webRTCTrak,
+			cb: func(unit formatprocessor.Unit, ctx context.Context, writeError chan error) {
+				tunit := unit.(*formatprocessor.UnitAV1)
+
+				if tunit.OBUs == nil {
+					return
+				}
+
+				packets, err := encoder.Encode(tunit.OBUs, tunit.PTS)
+				if err != nil {
+					panic(err)
+				}
+
+				for _, pkt := range packets {
+					webRTCTrak.WriteRTP(pkt)
+				}
+			},
+		}, nil
+	}
+
 	var vp9Format *formats.VP9
 	vp9Media := medias.FindFormat(&vp9Format)