datarhei-dragonfork-core/app/webrtc/handler.go

package webrtc

import (
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/labstack/echo/v4"
	"github.com/pion/webrtc/v4"

	corewebrtc "github.com/datarhei/core/v16/core/webrtc"
)

// Default per-stream peer cap when the caller passes 0. The total cap
// (passed to NewHandler) is enforced separately and takes precedence.
const defaultMaxPeersPerStream = 8

// WebRTCStats is the JSON response for GET /webrtc/stats.
type WebRTCStats struct {
	// ActiveStreams is the number of running FFmpeg processes with a
	// registered WHEP egress pair (video + audio Sources).
	ActiveStreams int `json:"active_streams"`

	// ActivePeers is the total count of live WHEP subscriber sessions
	// (each call to Subscribe that has not yet been torn down).
	ActivePeers int64 `json:"active_peers"`

	// ActivePublishers is the total count of live WHIP ingest sessions
	// (each call to WHIPHandler.Publish that has not yet been unpublished).
	ActivePublishers int64 `json:"active_publishers"`

	// UDPPortsInUse is an approximation of the number of UDP ports
	// allocated for ICE traffic. When using ephemeral ports (default)
	// each stream uses two ports (one video, one audio).
	UDPPortsInUse int `json:"udp_ports_in_use"`
}

// Handler exposes the subsystem's WHEP Echo handlers. Wire them into
// the /api/v3 group (or a sibling group) via Handler.Register.
type Handler struct {
	sub  *Subsystem
	whip *WHIPHandler // optional; enables active_publishers in /webrtc/stats

	mu            sync.Mutex
	peersByStream map[string]map[string]*corewebrtc.Peer // streamID -> resource -> peer
	peerStream    map[string]string                      // resource -> streamID (reverse index)
	count         int64                                  // atomic
	maxCapTotal   int64
	maxCapPerStrm int64

	met *webrtcMetrics
}

// NewHandler wraps the subsystem in an Echo-compatible HTTP handler.
func NewHandler(s *Subsystem, maxPeers int) *Handler {
	return NewHandlerWithCaps(s, maxPeers, 0)
}

// NewHandlerWithCaps is NewHandler plus an explicit per-stream cap.
func NewHandlerWithCaps(s *Subsystem, maxPeers, maxPeersPerStream int) *Handler {
	total := int64(maxPeers)
	if total <= 0 {
		total = int64(corewebrtc.DefaultConfig().MaxPeersTotal)
	}
	perStream := int64(maxPeersPerStream)
	if perStream <= 0 {
		perStream = defaultMaxPeersPerStream
	}
	h := &Handler{
		sub:           s,
		peersByStream: make(map[string]map[string]*corewebrtc.Peer),
		peerStream:    make(map[string]string),
		maxCapTotal:   total,
		maxCapPerStrm: perStream,
	}
	if s != nil {
		s.SetTeardownHook(h.tearDownStreamPeers)
	}
	return h
}

// SetWHIPHandler links the WHIP ingest handler so that /webrtc/stats
// can report active_publishers. Pass nil to disable that field (returns 0).
func (h *Handler) SetWHIPHandler(wh *WHIPHandler) {
	h.whip = wh
}

// Register mounts the WHEP routes and the shared stats route on the
// provided Echo group.
//
// Routes registered:
//
//	GET     /webrtc/stats
//	OPTIONS /whep/:id
//	OPTIONS /whep/:id/:resource
//	POST    /whep/:id
//	DELETE  /whep/:id/:resource
//	PATCH   /whep/:id/:resource
func (h *Handler) Register(g *echo.Group) {
	g.GET("/webrtc/stats", h.StatsHandler)
	g.OPTIONS("/whep/:id", h.preflight)
	g.OPTIONS("/whep/:id/:resource", h.preflight)
	g.POST("/whep/:id", h.Subscribe)
	g.DELETE("/whep/:id/:resource", h.Unsubscribe)
	g.PATCH("/whep/:id/:resource", h.Trickle)
}

// StatsHandler handles GET /webrtc/stats. Returns a JSON snapshot of
// the current WebRTC subsystem state.
//
// @Summary WebRTC subsystem stats
// @Description Returns a live snapshot: active egress streams, subscriber peer count, ingest publisher count, and approximate UDP port usage.
// @Tags v16.16.0
// @ID webrtc-3-stats
// @Produce json
// @Success 200 {object} WebRTCStats
// @Router /api/v3/webrtc/stats [get]
func (h *Handler) StatsHandler(c echo.Context) error {
	sc := 0
	if h.sub != nil {
		sc = h.sub.StreamCount()
	}

	var publishers int64
	if h.whip != nil {
		publishers = h.whip.PublisherCount()
	}

	stats := WebRTCStats{
		ActiveStreams:     sc,
		ActivePeers:      atomic.LoadInt64(&h.count),
		ActivePublishers: publishers,
		UDPPortsInUse:    sc * 2,
	}
	return c.JSON(http.StatusOK, stats)
}

// Subscribe handles POST /whep/:id.
func (h *Handler) Subscribe(c echo.Context) error {
	addCORS(c)
	t0 := time.Now()

	id := c.Param("id")
	if id == "" {
		h.recordRequest("subscribe", "", http.StatusBadRequest, t0)
		return c.String(http.StatusBadRequest, "missing stream id")
	}

	if atomic.LoadInt64(&h.count) >= h.maxCapTotal {
		if h.met != nil {
			h.met.capRejections.WithLabelValues("", "global").Inc()
		}
		h.recordRequest("subscribe", id, http.StatusServiceUnavailable, t0)
		return c.String(http.StatusServiceUnavailable, corewebrtc.ErrPeerCapReached.Error())
	}

	stream, ok := h.sub.lookup(id)
	if !ok {
		h.recordRequest("subscribe", id, http.StatusNotFound, t0)
		return c.String(http.StatusNotFound, corewebrtc.ErrStreamNotFound.Error())
	}

	h.mu.Lock()
	if int64(len(h.peersByStream[id])) >= h.maxCapPerStrm {
		h.mu.Unlock()
		if h.met != nil {
			h.met.capRejections.WithLabelValues(id, "stream").Inc()
		}
		h.recordRequest("subscribe", id, http.StatusServiceUnavailable, t0)
		return c.String(http.StatusServiceUnavailable, "webrtc: per-stream peer cap reached")
	}
	h.mu.Unlock()

	body, err := io.ReadAll(c.Request().Body)
	if err != nil {
		h.recordRequest("subscribe", id, http.StatusBadRequest, t0)
		return c.String(http.StatusBadRequest, "read body: "+err.Error())
	}
	if len(body) == 0 || !strings.HasPrefix(string(body), "v=") {
		h.recordRequest("subscribe", id, http.StatusBadRequest, t0)
		return c.String(http.StatusBadRequest, corewebrtc.ErrInvalidSDP.Error())
	}
	if err := requireH264AndOpus(string(body)); err != nil {
		if h.met != nil {
			if cme, ok2 := err.(*codecMismatchError); ok2 {
				for _, kind := range cme.missing {
					h.met.codecMismatches.WithLabelValues(id, strings.ToLower(kind)).Inc()
				}
			}
		}
		h.recordRequest("subscribe", id, http.StatusNotAcceptable, t0)
		return c.String(http.StatusNotAcceptable, err.Error())
	}

	offer := webrtc.SessionDescription{Type: webrtc.SDPTypeOffer, SDP: string(body)}
	peer, err := h.sub.factory.CreatePeerFromSources(c.Request().Context(), stream.video, stream.audio, offer)
	if err != nil {
		switch err {
		case corewebrtc.ErrICETimeout:
			h.recordRequest("subscribe", id, http.StatusGatewayTimeout, t0)
			return c.String(http.StatusGatewayTimeout, err.Error())
		case corewebrtc.ErrCodecMismatch:
			h.recordRequest("subscribe", id, http.StatusNotAcceptable, t0)
			return c.String(http.StatusNotAcceptable, err.Error())
		default:
			h.recordRequest("subscribe", id, http.StatusInternalServerError, t0)
			return c.String(http.StatusInternalServerError, "create peer: "+err.Error())
		}
	}

	rid := peer.ResourceID()
	h.mu.Lock()
	if h.peersByStream[id] == nil {
		h.peersByStream[id] = make(map[string]*corewebrtc.Peer)
	}
	h.peersByStream[id][rid] = peer
	h.peerStream[rid] = id
	h.mu.Unlock()
	atomic.AddInt64(&h.count, 1)

	go h.awaitPeerClose(rid, peer)
	go h.trackICE(id, peer, time.Now())

	h.recordRequest("subscribe", id, http.StatusCreated, t0)

	// RFC 9429 §4.3: emit one Link header per configured ICE server.
	for _, uri := range h.sub.ICEServerURIs() {
		c.Response().Header().Add("Link", "<"+uri+`>; rel="ice-server"`)
	}
	c.Response().Header().Set("Content-Type", "application/sdp")
	c.Response().Header().Set("Location", "/whep/"+id+"/"+rid)
	c.Response().Header().Set("ETag", `"`+rid+`"`)
	return c.String(http.StatusCreated, peer.Answer().SDP)
}

// Unsubscribe handles DELETE /whep/:id/:resource.
func (h *Handler) Unsubscribe(c echo.Context) error {
	addCORS(c)
	t0 := time.Now()

	resource := c.Param("resource")
	if resource == "" {
		h.recordRequest("unsubscribe", "", http.StatusBadRequest, t0)
		return c.String(http.StatusBadRequest, "missing resource id")
	}

	h.mu.Lock()
	streamID := h.peerStream[resource]
	var peer *corewebrtc.Peer
	if streamID != "" {
		peer = h.peersByStream[streamID][resource]
		delete(h.peersByStream[streamID], resource)
		if len(h.peersByStream[streamID]) == 0 {
			delete(h.peersByStream, streamID)
		}
		delete(h.peerStream, resource)
	}
	h.mu.Unlock()

	if peer != nil {
		_ = peer.Close()
	}
	if streamID != "" {
		atomic.AddInt64(&h.count, -1)
	}

	h.recordRequest("unsubscribe", streamID, http.StatusNoContent, t0)
	return c.NoContent(http.StatusNoContent)
}

// Trickle handles PATCH /whep/:id/:resource.
func (h *Handler) Trickle(c echo.Context) error {
	addCORS(c)
	t0 := time.Now()

	resource := c.Param("resource")
	if resource == "" {
		h.recordRequest("trickle", "", http.StatusBadRequest, t0)
		return c.String(http.StatusBadRequest, "missing resource id")
	}

	h.mu.Lock()
	streamID := h.peerStream[resource]
	var peer *corewebrtc.Peer
	if streamID != "" {
		peer = h.peersByStream[streamID][resource]
	}
	h.mu.Unlock()
	if peer == nil {
		h.recordRequest("trickle", streamID, http.StatusNotFound, t0)
		return c.NoContent(http.StatusNotFound)
	}

	body, err := io.ReadAll(c.Request().Body)
	if err != nil {
		h.recordRequest("trickle", streamID, http.StatusBadRequest, t0)
		return c.String(http.StatusBadRequest, "read body: "+err.Error())
	}
	for _, line := range strings.Split(string(body), "\n") {
		line = strings.TrimSpace(line)
		if !strings.HasPrefix(line, "a=candidate:") {
			continue
		}
		cand := strings.TrimPrefix(line, "a=")
		_ = peer.AddICECandidate(webrtc.ICECandidateInit{Candidate: cand})
	}

	h.recordRequest("trickle", streamID, http.StatusNoContent, t0)
	return c.NoContent(http.StatusNoContent)
}

func (h *Handler) recordRequest(route, streamID string, code int, t0 time.Time) {
	if h.met == nil {
		return
	}
	codeStr := fmt.Sprintf("%d", code)
	h.met.whepRequests.WithLabelValues(route, codeStr, streamID).Inc()
	h.met.whepRequestDuration.WithLabelValues(route, streamID).Observe(time.Since(t0).Seconds())
}

func (h *Handler) preflight(c echo.Context) error {
	addCORS(c)
	return c.NoContent(http.StatusNoContent)
}

// Close tears down every active peer (e.g., during Core shutdown).
func (h *Handler) Close() {
	h.mu.Lock()
	peers := make([]*corewebrtc.Peer, 0)
	for _, m := range h.peersByStream {
		for _, p := range m {
			peers = append(peers, p)
		}
	}
	h.peersByStream = make(map[string]map[string]*corewebrtc.Peer)
	h.peerStream = make(map[string]string)
	h.mu.Unlock()

	for _, p := range peers {
		if p != nil {
			_ = p.Close()
		}
	}
	atomic.StoreInt64(&h.count, 0)
}

func (h *Handler) awaitPeerClose(resource string, peer *corewebrtc.Peer) {
	<-peer.Done()
	h.mu.Lock()
	streamID := h.peerStream[resource]
	_, present := h.peerStream[resource]
	if present {
		delete(h.peerStream, resource)
		if streamID != "" {
			delete(h.peersByStream[streamID], resource)
			if len(h.peersByStream[streamID]) == 0 {
				delete(h.peersByStream, streamID)
			}
		}
	}
	h.mu.Unlock()
	if present {
		atomic.AddInt64(&h.count, -1)
	}
}

func (h *Handler) tearDownStreamPeers(streamID string) {
	h.mu.Lock()
	bucket := h.peersByStream[streamID]
	hadPeers := len(bucket) > 0
	peers := make([]*corewebrtc.Peer, 0, len(bucket))
	for _, p := range bucket {
		peers = append(peers, p)
	}
	h.mu.Unlock()

	for _, p := range peers {
		if p != nil {
			_ = p.Close()
		}
	}

	if hadPeers && h.met != nil {
		h.met.ffmpegLegFailures.WithLabelValues(streamID, "video").Inc()
		h.met.ffmpegLegFailures.WithLabelValues(streamID, "audio").Inc()
	}
}

func addCORS(c echo.Context) {
	hh := c.Response().Header()
	hh.Set("Access-Control-Allow-Origin", "*")
	hh.Set("Access-Control-Allow-Methods", "POST, DELETE, PATCH, OPTIONS")
	hh.Set("Access-Control-Allow-Headers", "Content-Type, Authorization, If-Match, If-None-Match")
	hh.Set("Access-Control-Expose-Headers", "Location, ETag, Link")
}

func requireH264AndOpus(sdp string) error {
	lower := strings.ToLower(sdp)
	hasH264 := strings.Contains(lower, "h264/90000") || strings.Contains(lower, " h264/")
	hasOpus := strings.Contains(lower, "opus/48000") || strings.Contains(lower, " opus/")
	if hasH264 && hasOpus {
		return nil
	}
	missing := []string{}
	if !hasH264 {
		missing = append(missing, "H264")
	}
	if !hasOpus {
		missing = append(missing, "Opus")
	}
	return &codecMismatchError{missing: missing}
}

type codecMismatchError struct{ missing []string }

func (e *codecMismatchError) Error() string {
	return "webrtc: codec mismatch -- offer is missing: " + strings.Join(e.missing, ", ")
}