datarhei-dragonfork-core/core/webrtc/peer.go

package webrtc

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"sync"

	"github.com/pion/rtp"
	"github.com/pion/webrtc/v4"
)

// PeerFactory builds PeerConnections from a shared Pion API instance
// configured from Config.
type PeerFactory struct {
	api       *webrtc.API
	rtcConfig webrtc.Configuration
}

// NewPeerFactory initializes a Pion API with the codec set we support
// (H.264 + Opus) and applies the provided Config.
func NewPeerFactory(c Config) (*PeerFactory, error) {
	if err := c.Validate(); err != nil {
		return nil, err
	}

	me := &webrtc.MediaEngine{}
	if err := me.RegisterDefaultCodecs(); err != nil {
		return nil, fmt.Errorf("webrtc: register default codecs: %w", err)
	}

	rtcConfig, se, err := BuildICEConfig(c)
	if err != nil {
		return nil, err
	}

	opts := []func(*webrtc.API){webrtc.WithMediaEngine(me)}
	if se != nil {
		opts = append(opts, webrtc.WithSettingEngine(*se))
	}
	api := webrtc.NewAPI(opts...)

	return &PeerFactory{api: api, rtcConfig: rtcConfig}, nil
}

// Peer wraps a Pion PeerConnection bound to a Source's subscription.
type Peer struct {
	resourceID string
	pc         *webrtc.PeerConnection
	answer     webrtc.SessionDescription
	source     *Source
	sub        chan *rtp.Packet
	done       chan struct{}
	once       sync.Once
}

// CreatePeer builds a PeerConnection, sets the remote offer, generates an
// answer, attaches video+audio tracks fed from src, and blocks until ICE
// gathering completes or ctx expires.
func (f *PeerFactory) CreatePeer(ctx context.Context, src *Source, offer webrtc.SessionDescription) (*Peer, error) {
	pc, err := f.api.NewPeerConnection(f.rtcConfig)
	if err != nil {
		return nil, fmt.Errorf("webrtc: new peer connection: %w", err)
	}

	videoTrack, err := webrtc.NewTrackLocalStaticRTP(
		webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeH264},
		"video", "dragonfork")
	if err != nil {
		_ = pc.Close()
		return nil, fmt.Errorf("webrtc: new video track: %w", err)
	}
	audioTrack, err := webrtc.NewTrackLocalStaticRTP(
		webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeOpus},
		"audio", "dragonfork")
	if err != nil {
		_ = pc.Close()
		return nil, fmt.Errorf("webrtc: new audio track: %w", err)
	}
	if _, err := pc.AddTrack(videoTrack); err != nil {
		_ = pc.Close()
		return nil, fmt.Errorf("webrtc: add video track: %w", err)
	}
	if _, err := pc.AddTrack(audioTrack); err != nil {
		_ = pc.Close()
		return nil, fmt.Errorf("webrtc: add audio track: %w", err)
	}

	if err := pc.SetRemoteDescription(offer); err != nil {
		_ = pc.Close()
		return nil, fmt.Errorf("webrtc: set remote: %w", err)
	}
	answer, err := pc.CreateAnswer(nil)
	if err != nil {
		_ = pc.Close()
		return nil, fmt.Errorf("webrtc: create answer: %w", err)
	}

	gatherComplete := webrtc.GatheringCompletePromise(pc)
	if err := pc.SetLocalDescription(answer); err != nil {
		_ = pc.Close()
		return nil, fmt.Errorf("webrtc: set local: %w", err)
	}

	select {
	case <-gatherComplete:
	case <-ctx.Done():
		_ = pc.Close()
		return nil, ErrICETimeout
	}

	sub := src.Subscribe(64)

	p := &Peer{
		resourceID: newResourceID(),
		pc:         pc,
		answer:     *pc.LocalDescription(),
		source:     src,
		sub:        sub,
		done:       make(chan struct{}),
	}

	pc.OnConnectionStateChange(func(st webrtc.PeerConnectionState) {
		if st == webrtc.PeerConnectionStateFailed ||
			st == webrtc.PeerConnectionStateDisconnected ||
			st == webrtc.PeerConnectionStateClosed {
			_ = p.Close()
		}
	})

	go forwardRTP(p.done, sub, videoTrack, audioTrack)

	return p, nil
}

// Answer returns the locally-created SDP answer. Valid after CreatePeer.
func (p *Peer) Answer() webrtc.SessionDescription { return p.answer }

// ResourceID returns the stable resource id used in the WHEP Location header.
func (p *Peer) ResourceID() string { return p.resourceID }

// Close tears down the peer connection and unsubscribes from the source.
// Safe to call multiple times.
func (p *Peer) Close() error {
	var err error
	p.once.Do(func() {
		close(p.done)
		p.source.Unsubscribe(p.sub)
		err = p.pc.Close()
	})
	return err
}

func newResourceID() string {
	b := make([]byte, 8)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}