core/whip/pion_whep.go

package whip

import (
	"fmt"
	"strings"
	"sync"

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

// whepCodecParams describes the exact RTP codec the WHIP publisher (OBS) is sending.
// Using these instead of RegisterDefaultCodecs ensures the browser negotiates
// the same payload type numbers OBS uses, which is mandatory because
// TrackLocalStaticRTP does NOT rewrite the PT in packets — it only rewrites SSRC.
// Mismatch: OBS sends PT 96, browser expects PT 102 → browser silently discards
// all H264 packets → no video. Audio works by coincidence (Opus is PT 111 in both).
type whepCodecParams struct {
	VideoMime      string // "video/H264"
	VideoClockRate uint32
	VideoFmtp      string // e.g. "level-asymmetry-allowed=1;packetization-mode=1;profile-level-id=640034"
	VideoPT        uint8  // OBS's H264 payload type, typically 96
	AudioMime      string // "audio/opus"
	AudioClockRate uint32
	AudioChannels  uint16
	AudioPT        uint8  // OBS's Opus payload type, typically 111
}

// 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.
// params must match the publisher's codec configuration exactly so that the
// browser negotiates the same payload type numbers the publisher uses.
func newWHEPPionSession(id, offerSDP string, params whepCodecParams) (*whepPionSession, string, error) {
	m := &webrtc.MediaEngine{}

	// Register the exact video codec OBS is sending.
	// Using the publisher's PT (e.g. 96) forces the SDP answer to contain
	// a=rtpmap:96 H264/90000, so the browser maps PT 96 → H264 and accepts
	// the raw RTP packets that arrive unchanged from OBS.
	videoFmtp := normalizeH264Fmtp(params.VideoMime, params.VideoFmtp)
	if err := m.RegisterCodec(webrtc.RTPCodecParameters{
		RTPCodecCapability: webrtc.RTPCodecCapability{
			MimeType:    params.VideoMime,
			ClockRate:   params.VideoClockRate,
			SDPFmtpLine: videoFmtp,
			RTCPFeedback: []webrtc.RTCPFeedback{
				{Type: "nack", Parameter: "pli"},
			},
		},
		PayloadType: webrtc.PayloadType(params.VideoPT),
	}, webrtc.RTPCodecTypeVideo); err != nil {
		return nil, "", fmt.Errorf("whep: register video codec: %w", err)
	}

	audioFmtp := ""
	if strings.Contains(strings.ToLower(params.AudioMime), "opus") {
		audioFmtp = "minptime=10;useinbandfec=1"
	}
	if err := m.RegisterCodec(webrtc.RTPCodecParameters{
		RTPCodecCapability: webrtc.RTPCodecCapability{
			MimeType:    params.AudioMime,
			ClockRate:   params.AudioClockRate,
			Channels:    params.AudioChannels,
			SDPFmtpLine: audioFmtp,
		},
		PayloadType: webrtc.PayloadType(params.AudioPT),
	}, webrtc.RTPCodecTypeAudio); err != nil {
		return nil, "", fmt.Errorf("whep: register audio codec: %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: params.VideoMime, ClockRate: params.VideoClockRate, SDPFmtpLine: videoFmtp},
		"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: params.AudioMime, ClockRate: params.AudioClockRate, Channels: params.AudioChannels},
		"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
}

// normalizeH264Fmtp ensures H264 fmtp always has level-asymmetry-allowed=1 and
// packetization-mode=1. level-asymmetry-allowed lets the browser decode a
// higher profile than it advertised. packetization-mode=1 is required for FU-A
// fragmented NALUs (how OBS sends large H264 frames).
func normalizeH264Fmtp(mime, fmtp string) string {
	if !strings.Contains(strings.ToLower(mime), "h264") {
		return fmtp
	}
	if fmtp == "" {
		return "level-asymmetry-allowed=1;packetization-mode=1;profile-level-id=42e01f"
	}
	if !strings.Contains(fmtp, "packetization-mode") {
		fmtp += ";packetization-mode=1"
	}
	if !strings.Contains(fmtp, "level-asymmetry-allowed") {
		fmtp = "level-asymmetry-allowed=1;" + fmtp
	}
	return fmtp
}

// 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)
	})
}