From 16e48314990d851368e57ffbe7b1afe5b01b5b15 Mon Sep 17 00:00:00 2001
From: Alex X <alexey.khit@gmail.com>
Date: Thu, 24 Oct 2024 23:31:21 +0300
Subject: [PATCH] Add the option to pass ICE servers with an async WebRTC offer
 #1408

---
 internal/api/ws/ws.go     | 28 +++++++++++------------
 internal/webrtc/webrtc.go | 48 ++++++++++++++++++++++++++-------------
 2 files changed, 46 insertions(+), 30 deletions(-)

diff --git a/internal/api/ws/ws.go b/internal/api/ws/ws.go
index 800a377d..0eb45aa7 100644
--- a/internal/api/ws/ws.go
+++ b/internal/api/ws/ws.go
@@ -1,6 +1,7 @@
 package ws
 
 import (
+	"encoding/json"
 	"io"
 	"net/http"
 	"net/url"
@@ -36,22 +37,16 @@ var log zerolog.Logger
 type Message struct {
 	Type  string `json:"type"`
 	Value any    `json:"value,omitempty"`
+	Raw   []byte `json:"-"`
 }
 
-func (m *Message) String() string {
-	if s, ok := m.Value.(string); ok {
-		return s
-	}
-	return ""
+func (m *Message) String() (value string) {
+	_ = json.Unmarshal(m.Raw, &value)
+	return
 }
 
-func (m *Message) GetString(key string) string {
-	if v, ok := m.Value.(map[string]any); ok {
-		if s, ok := v[key].(string); ok {
-			return s
-		}
-	}
-	return ""
+func (m *Message) Unmarshal(v any) error {
+	return json.Unmarshal(m.Raw, v)
 }
 
 type WSHandler func(tr *Transport, msg *Message) error
@@ -118,8 +113,11 @@ func apiWS(w http.ResponseWriter, r *http.Request) {
 	})
 
 	for {
-		msg := new(Message)
-		if err = ws.ReadJSON(msg); err != nil {
+		var raw struct {
+			Type  string          `json:"type"`
+			Value json.RawMessage `json:"value"`
+		}
+		if err = ws.ReadJSON(&raw); err != nil {
 			if !websocket.IsCloseError(err, websocket.CloseNoStatusReceived) {
 				log.Trace().Err(err).Caller().Send()
 			}
@@ -127,6 +125,8 @@ func apiWS(w http.ResponseWriter, r *http.Request) {
 			break
 		}
 
+		msg := &Message{Type: raw.Type, Raw: raw.Value}
+
 		log.Trace().Str("type", msg.Type).Msg("[api] ws msg")
 
 		if handler := wsHandlers[msg.Type]; handler != nil {
diff --git a/internal/webrtc/webrtc.go b/internal/webrtc/webrtc.go
index 8b4943c3..fe25c919 100644
--- a/internal/webrtc/webrtc.go
+++ b/internal/webrtc/webrtc.go
@@ -40,15 +40,17 @@ func Init() {
 		AddCandidate(network, candidate)
 	}
 
+	var err error
+
 	// create pionAPI with custom codecs list and custom network settings
-	serverAPI, err := webrtc.NewServerAPI(network, address, &filters)
+	serverAPI, err = webrtc.NewServerAPI(network, address, &filters)
 	if err != nil {
 		log.Error().Err(err).Caller().Send()
 		return
 	}
 
 	// use same API for WebRTC server and client if no address
-	clientAPI := serverAPI
+	clientAPI = serverAPI
 
 	if address != "" {
 		log.Info().Str("addr", cfg.Mod.Listen).Msg("[webrtc] listen")
@@ -81,11 +83,13 @@ func Init() {
 	streams.HandleFunc("webrtc", streamsHandler)
 }
 
+var serverAPI, clientAPI *pion.API
+
 var log zerolog.Logger
 
 var PeerConnection func(active bool) (*pion.PeerConnection, error)
 
-func asyncHandler(tr *ws.Transport, msg *ws.Message) error {
+func asyncHandler(tr *ws.Transport, msg *ws.Message) (err error) {
 	var stream *streams.Stream
 	var mode core.Mode
 
@@ -104,8 +108,30 @@ func asyncHandler(tr *ws.Transport, msg *ws.Message) error {
 		return errors.New(api.StreamNotFound)
 	}
 
+	var offer struct {
+		Type       string           `json:"type"`
+		SDP        string           `json:"sdp"`
+		ICEServers []pion.ICEServer `json:"ice_servers"`
+	}
+
+	// V2 - json/object exchange, V1 - raw SDP exchange
+	apiV2 := msg.Type == "webrtc"
+
+	if apiV2 {
+		if err = msg.Unmarshal(&offer); err != nil {
+			return err
+		}
+	} else {
+		offer.SDP = msg.String()
+	}
+
 	// create new PeerConnection instance
-	pc, err := PeerConnection(false)
+	var pc *pion.PeerConnection
+	if offer.ICEServers == nil {
+		pc, err = PeerConnection(false)
+	} else {
+		pc, err = serverAPI.NewPeerConnection(pion.Configuration{ICEServers: offer.ICEServers})
+	}
 	if err != nil {
 		log.Error().Err(err).Caller().Send()
 		return err
@@ -145,20 +171,10 @@ func asyncHandler(tr *ws.Transport, msg *ws.Message) error {
 		}
 	})
 
-	// V2 - json/object exchange, V1 - raw SDP exchange
-	apiV2 := msg.Type == "webrtc"
+	log.Trace().Msgf("[webrtc] offer:\n%s", offer.SDP)
 
 	// 1. SetOffer, so we can get remote client codecs
-	var offer string
-	if apiV2 {
-		offer = msg.GetString("sdp")
-	} else {
-		offer = msg.String()
-	}
-
-	log.Trace().Msgf("[webrtc] offer:\n%s", offer)
-
-	if err = conn.SetOffer(offer); err != nil {
+	if err = conn.SetOffer(offer.SDP); err != nil {
 		log.Warn().Err(err).Caller().Send()
 		return err
 	}