silo-server/internal/proxy/server.go

package proxy

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"mime"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/go-chi/cors"

	"github.com/Silo-Server/silo-server/internal/downloadprepare"
	"github.com/Silo-Server/silo-server/internal/downloads"
	"github.com/Silo-Server/silo-server/internal/nodeconfig"
	"github.com/Silo-Server/silo-server/internal/nodesessions"
	"github.com/Silo-Server/silo-server/internal/playback"
	"github.com/Silo-Server/silo-server/internal/streamtoken"
)

// Server is the HTTP handler for proxy mode.
type Server struct {
	watcher              *nodeconfig.Watcher
	tracker              *nodesessions.Tracker
	httpClient           *http.Client
	artifactMissReporter remoteArtifactMissReporter
	egress               *egressMeter
	// subCache stores full-track PGS (.sup) extracts under the transcode dir
	// so repeat selections skip the whole-file ffmpeg demux.
	subCache *playback.SubtitleCache
	// Download limits are node-local once egress is delegated. Rebuild the
	// manager when hot-reloaded settings change so existing transfers retain
	// their original limiter while new transfers use the new values.
	downloadBandwidthMu sync.Mutex
	downloadBandwidth   *downloads.BandwidthManager
	downloadServerBPS   int64
	downloadUserBPS     int64
}

type remoteArtifactMissReporter interface {
	ReportRemoteArtifactMissing(ctx context.Context, artifactID, originNodeURL, originArtifactID string) error
}

// NewServer creates a new proxy server backed by a config watcher and session
// tracker.
func NewServer(watcher *nodeconfig.Watcher, tracker *nodesessions.Tracker) *Server {
	return &Server{
		watcher: watcher,
		tracker: tracker,
		// No overall timeout — stream bodies are long-lived. Hung nodes are
		// bounded by the transport's response-header timeout instead.
		httpClient: &http.Client{
			Transport: newStreamTransport(),
			CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
				return http.ErrUseLastResponse
			},
		},
		egress: newEgressMeter(),
		subCache: playback.NewSubtitleCache(func() string {
			return watcher.Config().Playback.TranscodeDir
		}),
	}
}

// SetRemoteArtifactMissReporter wires the authoritative database transition
// used when an origin returns 404 after the API's proxy preflight. It must be
// called during construction, before the server begins handling requests.
func (s *Server) SetRemoteArtifactMissReporter(reporter remoteArtifactMissReporter) {
	s.artifactMissReporter = reporter
}

// newStreamTransport tunes the proxy→transcode-node connection pool. Many
// concurrent viewers fan their segment fetches through one proxy→node pair,
// and Go's default of 2 idle connections per host causes constant connection
// churn (and TLS re-handshakes) under load. The response-header timeout
// bounds requests to a hung node; the longest legitimate server-side wait is
// the 30s manifest-readiness poll on the transcode node.
func newStreamTransport() *http.Transport {
	t := http.DefaultTransport.(*http.Transport).Clone()
	t.MaxIdleConns = 128
	t.MaxIdleConnsPerHost = 32
	t.ResponseHeaderTimeout = 60 * time.Second
	return t
}

// Handler returns the chi.Router with all proxy routes mounted.
func (s *Server) Handler() http.Handler {
	r := chi.NewRouter()
	// hls.js uses XHR for manifest/segment fetches which are subject to
	// CORS when the proxy runs on a different origin than the web app.
	r.Use(cors.Handler(cors.Options{
		AllowedOrigins: []string{"*"},
		AllowedMethods: []string{"GET", "HEAD", "OPTIONS"},
		AllowedHeaders: []string{
			"Accept", "Authorization", "Content-Type", "Range",
			"If-Match", "If-Modified-Since", "If-None-Match", "If-Range", "If-Unmodified-Since",
		},
		MaxAge: 86400,
	}))
	r.Get("/api/v1/health", s.handleHealth)
	r.Group(func(r chi.Router) {
		// Streaming and download bytes count toward the node's measured egress.
		r.Use(s.meterEgress)
		r.Head("/stream/direct/{token}", s.handleDirectPlay)
		r.Get("/stream/direct/{token}", s.handleDirectPlay)
		r.Head("/stream/remux/{token}", s.handleRemux)
		r.Get("/stream/remux/{token}", s.handleRemux)
		r.Head("/stream/transcode/{token}/master.m3u8", s.handleTranscodeManifest)
		r.Get("/stream/transcode/{token}/master.m3u8", s.handleTranscodeManifest)
		r.Get("/stream/transcode/{token}/segment/{name}", s.handleTranscodeSegment)
		r.Get("/stream/subtitles/{token}/{track}/fonts", s.handleSubtitleFonts)
		r.Get("/stream/subtitles/{token}/{track}", s.handleSubtitle)
		r.Head("/downloads/file/{token}", s.handleDownloadFile)
		r.Get("/downloads/file/{token}", s.handleDownloadFile)
	})

	// Admin routes — bearer-auth protected.
	r.Group(func(r chi.Router) {
		r.Use(s.requireBearer)
		r.Post("/admin/force-reload", s.handleForceReload)
		r.Get("/status", s.handleStatus)
	})
	return r
}

type healthResponse struct {
	Status     string `json:"status"`
	ActiveJobs int    `json:"active_jobs"`
	EgressKbps int    `json:"egress_kbps"`
}

func (s *Server) handleHealth(w http.ResponseWriter, _ *http.Request) {
	activeJobs := 0
	if s.tracker != nil {
		activeJobs = s.tracker.ActiveCount()
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(healthResponse{
		Status:     "ok",
		ActiveJobs: activeJobs,
		EgressKbps: s.egress.RateKbps(),
	})
}

// requireBearer checks Authorization: Bearer {secret} for admin endpoints.
func (s *Server) requireBearer(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		cfg := s.watcher.Config()
		auth := r.Header.Get("Authorization")
		if !strings.HasPrefix(auth, "Bearer ") || strings.TrimPrefix(auth, "Bearer ") != cfg.Auth.JWTSecret {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		next.ServeHTTP(w, r)
	})
}

// verifyToken extracts and validates the stream token from the URL.
func (s *Server) verifyToken(w http.ResponseWriter, r *http.Request) *streamtoken.Claims {
	cfg := s.watcher.Config()
	tokenStr := chi.URLParam(r, "token")
	claims, err := streamtoken.Verify(tokenStr, cfg.Auth.JWTSecret)
	if err != nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return nil
	}
	return claims
}

func (s *Server) handleDirectPlay(w http.ResponseWriter, r *http.Request) {
	claims := s.verifyToken(w, r)
	if claims == nil {
		return
	}

	info := sessionInfo(s.tracker, claims, "direct_play")
	s.tracker.Track(r.Context(), info)
	defer s.tracker.Remove(r.Context(), claims.SessionID)

	http.ServeFile(w, r, claims.MediaPath)
}

func (s *Server) handleDownloadFile(w http.ResponseWriter, r *http.Request) {
	claims := s.verifyToken(w, r)
	if claims == nil {
		return
	}
	remoteArtifact := claims.DownloadArtifactID != "" && strings.TrimSpace(claims.TranscodeNode) != ""
	if claims.PlayMethod != streamtoken.PlayMethodDownload || (strings.TrimSpace(claims.MediaPath) == "" && !remoteArtifact) {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}

	// HEAD is a capability/path preflight, not an active transfer. Counting it
	// would briefly consume job capacity and could make a health report retire
	// the API's reservation before the client starts its GET.
	if s.tracker != nil && r.Method != http.MethodHead {
		info := sessionInfo(s.tracker, claims, "download")
		s.tracker.Track(r.Context(), info)
		defer s.tracker.Remove(context.WithoutCancel(r.Context()), claims.SessionID)
	}
	if remoteArtifact {
		s.relayDownloadArtifact(w, r, claims)
		return
	}

	f, err := os.Open(claims.MediaPath)
	if err != nil {
		http.NotFound(w, r)
		return
	}
	defer func() { _ = f.Close() }()
	stat, err := f.Stat()
	if err != nil {
		http.Error(w, "download unavailable", http.StatusInternalServerError)
		return
	}

	filename := filepath.Base(claims.MediaPath)
	if disposition := mime.FormatMediaType("attachment", map[string]string{"filename": filename}); disposition != "" {
		w.Header().Set("Content-Disposition", disposition)
	}
	w.Header().Set("Content-Type", playback.MimeFromExtension(claims.MediaPath))
	reader := io.ReadSeeker(f)
	if bandwidth := s.downloadBandwidthManager(); bandwidth != nil {
		reader = bandwidth.ThrottledReader(r.Context(), f, claims.UserID)
	}
	http.ServeContent(w, r, stat.Name(), stat.ModTime(), reader)
}

func (s *Server) relayDownloadArtifact(w http.ResponseWriter, r *http.Request, claims *streamtoken.Claims) {
	if !downloadprepare.ValidArtifactID(claims.DownloadArtifactID) {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	cfg := s.watcher.Config()
	if cfg == nil || strings.TrimSpace(cfg.Auth.JWTSecret) == "" {
		http.Error(w, "download unavailable", http.StatusServiceUnavailable)
		return
	}
	client := downloadprepare.HTTPPreparer{Client: s.httpClient}
	resp, err := client.Open(r.Context(), claims.TranscodeNode, cfg.Auth.JWTSecret, claims.DownloadArtifactID, r.Method, r.Header)
	if err != nil {
		slog.WarnContext(r.Context(), "download artifact relay failed", "component", "proxy", "artifact_id", claims.DownloadArtifactID, "node", claims.TranscodeNode, "error", err)
		http.Error(w, "download unavailable", http.StatusBadGateway)
		return
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode == http.StatusNotFound {
		if s.artifactMissReporter != nil && strings.TrimSpace(claims.DownloadArtifactRowID) != "" {
			reportCtx, cancel := context.WithTimeout(context.WithoutCancel(r.Context()), 5*time.Second)
			err := s.artifactMissReporter.ReportRemoteArtifactMissing(
				reportCtx, claims.DownloadArtifactRowID, claims.TranscodeNode, claims.DownloadArtifactID,
			)
			cancel()
			if err != nil {
				slog.WarnContext(r.Context(), "report missing remote download artifact", "component", "proxy", "artifact_id", claims.DownloadArtifactRowID, "error", err)
			}
		}
		http.NotFound(w, r)
		return
	}
	if !downloadprepare.RelayStatusAllowed(resp.StatusCode) {
		http.Error(w, "download unavailable", http.StatusBadGateway)
		return
	}
	downloadprepare.CopyResponseHeaders(w.Header(), resp.Header)
	if filename := filepath.Base(strings.TrimSpace(claims.DownloadFilename)); filename != "" && filename != "." {
		if disposition := mime.FormatMediaType("attachment", map[string]string{"filename": filename}); disposition != "" {
			w.Header().Set("Content-Disposition", disposition)
		}
	}
	w.WriteHeader(resp.StatusCode)
	if r.Method == http.MethodHead {
		return
	}
	var reader io.Reader = resp.Body
	if bandwidth := s.downloadBandwidthManager(); bandwidth != nil {
		reader = bandwidth.ThrottledStreamReader(r.Context(), reader, claims.UserID)
	}
	if _, err := io.Copy(w, reader); err != nil && r.Context().Err() == nil {
		slog.WarnContext(r.Context(), "download artifact relay interrupted", "component", "proxy", "artifact_id", claims.DownloadArtifactID, "error", err)
	}
}

func (s *Server) downloadBandwidthManager() *downloads.BandwidthManager {
	cfg := s.watcher.Config()
	if cfg == nil {
		return nil
	}
	serverBPS := cfg.Download.ServerBandwidthBPS
	userBPS := cfg.Download.UserBandwidthBPS
	s.downloadBandwidthMu.Lock()
	defer s.downloadBandwidthMu.Unlock()
	if s.downloadBandwidth == nil || serverBPS != s.downloadServerBPS || userBPS != s.downloadUserBPS {
		s.downloadBandwidth = downloads.NewBandwidthManager(serverBPS, userBPS)
		s.downloadServerBPS = serverBPS
		s.downloadUserBPS = userBPS
	}
	return s.downloadBandwidth
}

func (s *Server) handleRemux(w http.ResponseWriter, r *http.Request) {
	claims := s.verifyToken(w, r)
	if claims == nil {
		return
	}

	info := sessionInfo(s.tracker, claims, "remux")
	s.tracker.Track(r.Context(), info)
	defer s.tracker.Remove(r.Context(), claims.SessionID)

	seekSeconds := 0.0
	if seekStr := r.URL.Query().Get("seek"); seekStr != "" {
		if v, err := strconv.ParseFloat(seekStr, 64); err == nil {
			seekSeconds = v
		}
	}
	// Honor the Dolby Vision mode frozen in the token (empty decodes as the
	// legacy auto behavior for old tokens), mirroring how the integrated
	// server's stream handler serves the same claims.
	_ = playback.ServeRemuxWithOptions(w, r, claims.MediaPath, "mp4", seekSeconds, claims.TranscodeAudio, claims.AudioTrackIndex, claims.DVProfile, playback.RemuxServeOptions{
		DVMode:                 playback.RemuxDVMode(claims.RemuxDVMode),
		FFmpegPath:             s.watcher.Config().Playback.FFmpegPath,
		ContentType:            playback.RemuxContentType(claims.AudioOnly),
		AudioOnly:              claims.AudioOnly,
		TargetAudioChannels:    claims.TargetAudioChannels,
		TargetAudioBitrateKbps: claims.TargetAudioBitrateKbps,
	})
}

func (s *Server) handleTranscodeManifest(w http.ResponseWriter, r *http.Request) {
	claims := s.verifyToken(w, r)
	if claims == nil {
		return
	}
	s.touchTranscodeSession(r, claims)
	s.proxyToTranscodeNode(w, r, claims, "/transcode/"+transcodeTransportIDFromClaims(claims)+"/master.m3u8")
}

func (s *Server) handleTranscodeSegment(w http.ResponseWriter, r *http.Request) {
	claims := s.verifyToken(w, r)
	if claims == nil {
		return
	}
	s.touchTranscodeSession(r, claims)
	name := chi.URLParam(r, "name")
	s.proxyToTranscodeNode(w, r, claims, "/transcode/"+transcodeTransportIDFromClaims(claims)+"/segment/"+name)
}

func transcodeTransportIDFromClaims(claims *streamtoken.Claims) string {
	if claims != nil && claims.TranscodeTransportID != "" {
		return claims.TranscodeTransportID
	}
	if claims == nil {
		return ""
	}
	return claims.SessionID
}

// touchTranscodeSession keeps HLS sessions visible in the active stream count.
// Unlike direct play and remux, transcode playback reaches the proxy as many
// short manifest/segment requests, so the session is tracked by recent
// activity instead of request lifetime.
func (s *Server) touchTranscodeSession(r *http.Request, claims *streamtoken.Claims) {
	s.tracker.Touch(r.Context(), sessionInfo(s.tracker, claims, "transcode"))
}

// sessionInfo builds the node-session tracker record for a verified token,
// copying the numeric ownership keys the node-session tracker needs.
func sessionInfo(tr *nodesessions.Tracker, claims *streamtoken.Claims, kind string) nodesessions.SessionInfo {
	return nodesessions.SessionInfo{
		SessionID:   claims.SessionID,
		NodeURL:     tr.NodeURL(),
		NodeName:    tr.NodeName(),
		Type:        kind,
		StartedAt:   time.Now().UTC().Format(time.RFC3339),
		AuthUserID:  claims.UserID,
		ProfileID:   claims.ProfileID,
		MediaFileID: claims.MediaFileID,
	}
}

func (s *Server) handleSubtitle(w http.ResponseWriter, r *http.Request) {
	claims := s.verifyToken(w, r)
	if claims == nil {
		return
	}
	cfg := s.watcher.Config()
	trackParam := chi.URLParam(r, "track")
	trackIndex, requestedFormat, err := playback.ParseSubtitleTrackParam(trackParam)
	if err != nil {
		http.Error(w, "invalid subtitle index", http.StatusBadRequest)
		return
	}

	// When the URL requests SUP format (e.g. /subtitles/{token}/2.sup),
	// serve the PGS track as a raw .sup elementary stream for client-side
	// bitmap rendering (libpgs). Unlike the buffered text paths below, this
	// serves the cached full-track extract when present, and otherwise
	// streams ffmpeg output directly (the client renders progressively as
	// data arrives) while teeing it into the cache for the next request.
	// Clients that manage their own sliding window opt in with ?windowed=1
	// (+ ?position=/?duration=), mirroring the API stream handler; windowed
	// requests extract only the requested slice — from the cached full
	// track when one exists (warming it in the background when not).
	if requestedFormat == "sup" {
		allowWindow, seek, duration := playback.PGSWindowRequest(r.URL.Query())
		err := s.subCache.ServeSUPExtract(w, r, playback.StreamExtractOpts{
			InputPath:       claims.MediaPath,
			TrackIndex:      trackIndex,
			SourceCodec:     "hdmv_pgs_subtitle", // .sup URLs are only generated for PGS tracks
			SeekSeconds:     seek,
			DurationSeconds: duration,
			AllowWindow:     allowWindow,
			FFmpegPath:      cfg.Playback.FFmpegPath,
		}, playback.StreamExtractSubtitle)
		if err != nil && r.Context().Err() == nil {
			// Headers already committed — log and let the client see a
			// truncated response.
			slog.ErrorContext(r.Context(), "stream subtitle (sup)", "component", "proxy", "error", err, "track", trackIndex,
				"path", claims.MediaPath, "playback_session_id", claims.SessionID)
		}
		return
	}

	// When the URL requests ASS format (e.g. /subtitles/{token}/2.ass),
	// extract as raw ASS to preserve styling for client-side rendering.
	if requestedFormat == "ass" {
		data, err := playback.ExtractSubtitleWithFormat(r.Context(), claims.MediaPath, trackIndex, "ass", cfg.Playback.FFmpegPath)
		if err != nil {
			slog.ErrorContext(r.Context(), "extract subtitle (ass)", "component", "proxy", "error", err, "track", trackIndex, "path", claims.MediaPath, "playback_session_id", claims.SessionID)
			http.Error(w, "subtitle extraction failed", http.StatusInternalServerError)
			return
		}
		playback.ServeSubtitle(w, data, "ass")
		return
	}

	data, format, err := playback.ExtractSubtitle(r.Context(), claims.MediaPath, trackIndex, cfg.Playback.FFmpegPath)
	if err != nil {
		slog.ErrorContext(r.Context(), "extract subtitle", "component", "proxy", "error", err, "track", trackIndex, "path", claims.MediaPath, "playback_session_id", claims.SessionID)
		http.Error(w, "subtitle extraction failed", http.StatusInternalServerError)
		return
	}

	vtt, err := playback.ConvertToVTT(data, format)
	if err != nil {
		slog.ErrorContext(r.Context(), "convert to vtt", "component", "proxy", "error", err, "playback_session_id", claims.SessionID)
		http.Error(w, "subtitle conversion failed", http.StatusInternalServerError)
		return
	}

	playback.ServeSubtitle(w, vtt, "vtt")
}

func (s *Server) handleSubtitleFonts(w http.ResponseWriter, r *http.Request) {
	claims := s.verifyToken(w, r)
	if claims == nil {
		return
	}
	cfg := s.watcher.Config()
	trackParam := chi.URLParam(r, "track")
	trackIndex, _, err := playback.ParseSubtitleTrackParam(trackParam)
	if err != nil {
		http.Error(w, "invalid subtitle index", http.StatusBadRequest)
		return
	}

	fonts, err := playback.ExtractAttachedSubtitleFonts(r.Context(), claims.MediaPath, cfg.Playback.FFmpegPath)
	if err != nil {
		slog.ErrorContext(r.Context(), "extract subtitle fonts", "component", "proxy", "error", err, "track", trackIndex, "path", claims.MediaPath, "playback_session_id", claims.SessionID)
		http.Error(w, "subtitle font extraction failed", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Cache-Control", "no-store")
	if err := json.NewEncoder(w).Encode(playback.EncodeSubtitleFontBundle(fonts)); err != nil {
		slog.WarnContext(r.Context(), "subtitle font response encode failed", "component", "proxy", "error", err, "playback_session_id", claims.SessionID)
	}
}

// proxyToTranscodeNode forwards the request to the transcode node specified in the claims.
func (s *Server) proxyToTranscodeNode(w http.ResponseWriter, r *http.Request, claims *streamtoken.Claims, path string) {
	cfg := s.watcher.Config()
	if claims.TranscodeNode == "" {
		http.Error(w, "no transcode node in token", http.StatusBadRequest)
		return
	}

	targetURL := claims.TranscodeNode + path
	if rawQuery := r.URL.RawQuery; rawQuery != "" {
		targetURL = fmt.Sprintf("%s?%s", targetURL, rawQuery)
	}

	req, err := http.NewRequestWithContext(r.Context(), r.Method, targetURL, nil)
	if err != nil {
		http.Error(w, "internal error", http.StatusInternalServerError)
		return
	}
	req.Header.Set("Authorization", "Bearer "+cfg.Auth.JWTSecret)
	// Forward the verified stream token so the transcode node can self-reconstruct
	// a lost session after its OWN restart: the token carries the full byte-affecting
	// recipe, so the node can re-spawn ffmpeg seeked to the requested segment instead
	// of 404ing (the integrated server already does this from the same token). The
	// node re-verifies the token independently before trusting it.
	if token := chi.URLParam(r, "token"); token != "" {
		req.Header.Set("X-Silo-Stream-Token", token)
	}

	resp, err := s.httpClient.Do(req)
	if err != nil {
		slog.ErrorContext(r.Context(), "proxy to transcode node", "component", "proxy", "error", err, "url", targetURL, "playback_session_id", claims.SessionID)
		http.Error(w, "transcode node unavailable", http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	// Copy response headers
	for k, vv := range resp.Header {
		for _, v := range vv {
			w.Header().Add(k, v)
		}
	}
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}

func (s *Server) handleForceReload(w http.ResponseWriter, r *http.Request) {
	if err := s.watcher.ForceReload(r.Context()); err != nil {
		http.Error(w, "reload failed: "+err.Error(), http.StatusInternalServerError)
		return
	}
	slog.InfoContext(r.Context(), "proxy force reload completed", slog.String("component", "proxy"))
	w.WriteHeader(http.StatusNoContent)
}

type statusResponse struct {
	ActiveSessions int `json:"active_sessions"`
}

func (s *Server) handleStatus(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(statusResponse{
		ActiveSessions: s.tracker.ActiveCount(),
	})
}