Null-DRM-Official/apps/modules/rte/mpd.go

package rte

import (
	"fmt"
	"regexp"
	"strings"
	"sync"

	"drmdecryption/mpd"
)

func init() {
	mpd.Register(Name, func() mpd.Rewriter { return NewRewriter() })
}

var reBaseURL = regexp.MustCompile(`(?i)<BaseURL[^>]*>([^<]+)</BaseURL>`)

// Rewriter turns an ad-stitched upstream MPD into a synthetic manifest on the
// origin encoder's own segment grid, so a downloader keeps pulling live segments
// after the upstream session expires. All origins and KIDs come from config.
type Rewriter struct {
	cfg Config

	mu           sync.Mutex
	fallbackCh   string
	fallbackBase string
}

// NewRewriter loads the module config and returns a rewriter. A config error is
// deferred to Rewrite so registry lookup never fails at startup.
func NewRewriter() *Rewriter {
	cfg, _, err := loadConfig()
	if err != nil {
		return &Rewriter{cfg: Config{}.withDefaults()}
	}
	return &Rewriter{cfg: cfg}
}

// NewRewriterWith builds a rewriter from an already-loaded config.
func NewRewriterWith(cfg Config) *Rewriter {
	return &Rewriter{cfg: cfg}
}

func (r *Rewriter) Name() string { return Name }

// Prime records a channel/origin learned elsewhere (e.g. from the KID).
func (r *Rewriter) Prime(channel, originBase string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if channel != "" {
		r.fallbackCh = channel
	}
	if originBase != "" {
		r.fallbackBase = originBase
	}
}

// ResolveKID maps a KID back to its encoder channel and origin.
func (r *Rewriter) ResolveKID(kidHex string) (channel, originBase string, ok bool) {
	for ch, kid := range r.cfg.KIDs {
		if mpd.KIDHex(kid) == kidHex {
			return ch, r.cfg.Origins[ch], true
		}
	}
	return "", "", false
}

// Rewrite produces the synthetic MPD. upstreamXML may be empty (expired session):
// the KID fallback then supplies the channel and origin.
func (r *Rewriter) Rewrite(upstreamXML []byte, kidHint string) ([]byte, error) {
	r.mu.Lock()
	fbCh, fbBase := r.fallbackCh, r.fallbackBase
	r.mu.Unlock()

	payload, channel, originBase, err := r.rewrite(string(upstreamXML), kidHint, fbCh, fbBase)
	if channel != "" || originBase != "" {
		r.Prime(channel, originBase)
	}
	return payload, err
}

func (r *Rewriter) rewrite(xmlText, kidHint, fallbackChannel, fallbackBase string) (payload []byte, channel, originBase string, err error) {
	cfg := r.cfg
	channel, originBase = r.originFromXML(xmlText)
	resolvedKid := mpd.ExtractDefaultKID(xmlText)
	if resolvedKid == "" && kidHint != "" {
		resolvedKid, _ = mpd.KIDToUUID(kidHint)
	}
	if (channel == "" || originBase == "") && resolvedKid != "" {
		if mapped, base, ok := r.ResolveKID(mpd.KIDHex(resolvedKid)); ok {
			if channel == "" {
				channel = mapped
			}
			if originBase == "" {
				originBase = base
			}
		}
	}
	if channel == "" {
		channel = fallbackChannel
	}
	if originBase == "" {
		originBase = fallbackBase
	}
	if channel != "" && originBase == "" {
		originBase = cfg.Origins[channel]
	}
	if channel != "" && originBase != "" && r.usesEncoderTimeline(channel) {
		if resolvedKid == "" {
			resolvedKid = cfg.KIDs[channel]
		}
		if resolvedKid == "" {
			return nil, "", "", fmt.Errorf("%s: no KID for synthetic MPD channel %s", Name, channel)
		}
		payload, err = mpd.BuildSynthetic(cfg.Synthetic, channel, originBase, resolvedKid)
		return payload, channel, originBase, err
	}
	return nil, channel, originBase, fmt.Errorf("%s rewrite: unsupported MPD (no encoder channel); channel=%q", Name, channel)
}

func (r *Rewriter) usesEncoderTimeline(channel string) bool {
	ok, _ := regexp.MatchString(r.cfg.EncoderChannelRE, channel)
	return ok
}

// originFromXML finds the origin BaseURL and derives the encoder channel from it.
func (r *Rewriter) originFromXML(xmlText string) (channel, base string) {
	marker := r.cfg.OriginMarker
	if marker == "" {
		marker = ".isml"
	}
	for _, m := range reBaseURL.FindAllStringSubmatch(xmlText, -1) {
		text := strings.TrimSpace(m[1])
		idx := strings.Index(text, marker)
		if idx < 0 {
			continue
		}
		base = text[:idx] + marker + "/"
		parts := strings.Split(strings.TrimSuffix(base, "/"), "/")
		channel = strings.TrimSuffix(parts[len(parts)-1], marker)
		if channel != "" {
			return channel, base
		}
	}
	return "", ""
}