core/whip/pion_whep.go

package whip

import (
	"fmt"
	"sync"

	"github.com/pion/interceptor"
	"github.com/pion/webrtc/v3"
)

// whepPionSession is a single WHEP subscriber backed by a pion PeerConnection.
// It receives plain RTP packets (already decrypted by the WHIP publisher side)
// via WriteVideo/WriteAudio and forwards them to the remote browser via
// ICE + DTLS-SRTP.
type whepPionSession struct {
	pc         *webrtc.PeerConnection
	videoTrack *webrtc.TrackLocalStaticRTP
	audioTrack *webrtc.TrackLocalStaticRTP
	id         string
	done       chan struct{}
	once       sync.Once
}

// newWHEPPionSession creates a new WHEP subscriber PeerConnection.
// offerSDP is the SDP offer from the browser.
// videoMime: e.g. "video/H264"; audioMime: e.g. "audio/opus".
// Returns the session and the SDP answer to send back to the browser.
func newWHEPPionSession(id, offerSDP, videoMime, audioMime string) (*whepPionSession, string, error) {
	m := &webrtc.MediaEngine{}
	if err := m.RegisterDefaultCodecs(); err != nil {
		return nil, "", fmt.Errorf("whep: register codecs: %w", err)
	}

	ir := &interceptor.Registry{}
	if err := webrtc.RegisterDefaultInterceptors(m, ir); err != nil {
		return nil, "", fmt.Errorf("whep: register interceptors: %w", err)
	}

	api := webrtc.NewAPI(webrtc.WithMediaEngine(m), webrtc.WithInterceptorRegistry(ir))
	pc, err := api.NewPeerConnection(webrtc.Configuration{})
	if err != nil {
		return nil, "", fmt.Errorf("whep: new peer connection: %w", err)
	}

	videoTrack, err := webrtc.NewTrackLocalStaticRTP(
		webrtc.RTPCodecCapability{MimeType: videoMime},
		"video", "whep-video-"+id,
	)
	if err != nil {
		pc.Close()
		return nil, "", fmt.Errorf("whep: create video track: %w", err)
	}

	audioTrack, err := webrtc.NewTrackLocalStaticRTP(
		webrtc.RTPCodecCapability{MimeType: audioMime},
		"audio", "whep-audio-"+id,
	)
	if err != nil {
		pc.Close()
		return nil, "", fmt.Errorf("whep: create audio track: %w", err)
	}

	if _, err = pc.AddTrack(videoTrack); err != nil {
		pc.Close()
		return nil, "", fmt.Errorf("whep: add video track: %w", err)
	}
	if _, err = pc.AddTrack(audioTrack); err != nil {
		pc.Close()
		return nil, "", fmt.Errorf("whep: add audio track: %w", err)
	}

	if err = pc.SetRemoteDescription(webrtc.SessionDescription{
		Type: webrtc.SDPTypeOffer,
		SDP:  offerSDP,
	}); err != nil {
		pc.Close()
		return nil, "", fmt.Errorf("whep: set remote description: %w", err)
	}

	answer, err := pc.CreateAnswer(nil)
	if err != nil {
		pc.Close()
		return nil, "", fmt.Errorf("whep: create answer: %w", err)
	}

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

	sess := &whepPionSession{
		pc:         pc,
		videoTrack: videoTrack,
		audioTrack: audioTrack,
		id:         id,
		done:       make(chan struct{}),
	}

	pc.OnConnectionStateChange(func(state webrtc.PeerConnectionState) {
		if state == webrtc.PeerConnectionStateClosed ||
			state == webrtc.PeerConnectionStateFailed ||
			state == webrtc.PeerConnectionStateDisconnected {
			sess.Close()
		}
	})

	return sess, pc.LocalDescription().SDP, nil
}

// ID returns the unique subscriber identifier.
func (s *whepPionSession) ID() string { return s.id }

// WriteVideo forwards a decrypted RTP video packet to the subscriber via DTLS-SRTP.
func (s *whepPionSession) WriteVideo(pkt []byte) {
	select {
	case <-s.done:
		return
	default:
		s.videoTrack.Write(pkt) //nolint:errcheck
	}
}

// WriteAudio forwards a decrypted RTP audio packet to the subscriber via DTLS-SRTP.
func (s *whepPionSession) WriteAudio(pkt []byte) {
	select {
	case <-s.done:
		return
	default:
		s.audioTrack.Write(pkt) //nolint:errcheck
	}
}

// Done returns a channel that is closed when the subscriber disconnects.
func (s *whepPionSession) Done() <-chan struct{} { return s.done }

// Close tears down the subscriber PeerConnection.
func (s *whepPionSession) Close() {
	s.once.Do(func() {
		s.pc.Close()
		close(s.done)
	})
}