dogtamer/server/web.go

package server

import (
	"context"
	"fmt"
	"html/template"
	"net/http"
	"time"

	"git.t-juice.club/torjus/dogtamer/m3u"
	"github.com/go-chi/chi/v5"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"go.uber.org/zap"
)

type WebServer struct {
	Logger        *zap.SugaredLogger
	EnableMetrics bool
	ListenAddr    string
	ctx           context.Context
	rtmpServer    *RTMPServer
	httpServer    *http.Server
}

func NewWebServer(ctx context.Context, rs *RTMPServer) *WebServer {
	return &WebServer{
		ctx:        ctx,
		rtmpServer: rs,
		Logger:     zap.NewNop().Sugar(),
		ListenAddr: ":8077",
	}
}

func (ws *WebServer) Serve() error {
	r := chi.NewRouter()
	r.Get("/", ws.IndexHandler)
	if ws.EnableMetrics {
		r.Handle("/metrics", promhttp.Handler())
	}
	r.Get("/playlist/{name}", ws.PlaylistHandler)
	ws.httpServer = &http.Server{
		Addr:    ws.ListenAddr,
		Handler: r,
	}

	go func() {
		<-ws.ctx.Done()
		ws.Logger.Debugw("HTTP shutdown signal received.")
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = ws.httpServer.Shutdown(shutdownCtx)
	}()

	err := ws.httpServer.ListenAndServe()
	if err != nil && err != http.ErrServerClosed {
		ws.Logger.Warnw("HTTP Server stopped with error.", "error", err)
		return err
	}
	ws.Logger.Info("HTTP server stopped.")

	return nil
}

func (ws *WebServer) IndexHandler(w http.ResponseWriter, r *http.Request) {

	var data struct {
		Streams []struct {
			Name string
			Path template.URL
		}
	}
	streams := ws.rtmpServer.List()
	for _, stream := range streams {
		data.Streams = append(data.Streams, struct {
			Name string
			Path template.URL
		}{Name: stream.Name, Path: template.URL(fmt.Sprintf("/playlist/%s", stream.Name))})
	}

	tmpl := template.Must(template.ParseFiles("server/templates/index.html"))

	w.Header().Add("Content-Type", "text/html")

	tmpl.Execute(w, data)
}

func (ws *WebServer) PlaylistHandler(w http.ResponseWriter, r *http.Request) {
	streamName := chi.URLParam(r, "name")
	if streamName == "" {
		w.WriteHeader(http.StatusBadRequest)
		return
	}
	streamInfo, err := ws.rtmpServer.GetInfo(streamName)
	if err != nil {
		if err == ErrNoSuchItem {
			ws.Logger.Debugw("Client requested non-existing playlist", "stream_name", streamName, "error", err)
			w.WriteHeader(http.StatusNotFound)
			return
		}
		ws.Logger.Warnw("Error getting stream info", "stream_name", streamName, "error", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	var p m3u.Playlist
	p.Add(&m3u.PlaylistItem{Name: streamInfo.Name, Time: -1, Path: streamInfo.Path})

	w.Header().Add("Content-Type", "application/mpegurl")
	w.Header().Add("Content-Disposition", "attachment; filename=stream.m3u")
	w.WriteHeader(http.StatusOK)

	_, err = p.WriteTo(w)
	if err != nil {
		ws.Logger.Warnw("error generating playlist", "stream_name", streamName, "error", err)
	}
}