ministream/server/server.go

package server

import (
	"bytes"
	"embed"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"

	"github.com/go-chi/chi/v5"
	"github.com/pion/sdp/v3"
	"github.com/pion/webrtc/v4"
)

//go:embed static
var static embed.FS

type Server struct {
	users   *UserStore
	streams *StreamStore
	http.Server
}

func NewServer(store *UserStore) *Server {
	srv := &Server{
		users:   store,
		streams: NewStreamStore(),
	}

	r := chi.NewRouter()
	r.Use(corsMiddleware)
	r.Get("/", srv.StaticHandler)
	r.Get("/{name}", srv.StaticHandler)
	r.Post("/whip", http.HandlerFunc(srv.WhipHandler))
	r.Get("/whip", srv.ListHandler)
	r.Options("/whip", srv.OptionsHandler)
	r.Delete("/whip/{streamKey}", srv.DeleteHandler)
	r.Patch("/whip/{streamKey}", srv.PatchHandler)
	r.Post("/whip/{streamKey}", srv.PostOfferHandler)

	srv.Handler = r

	return srv
}

func corsMiddleware(next http.Handler) http.Handler {
	fn := func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Access-Control-Allow-Origin", "*")
		next.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}

func (s *Server) OptionsHandler(w http.ResponseWriter, r *http.Request) {
	slog.Info("Got OPTIONS")
}

func (s *Server) PatchHandler(w http.ResponseWriter, r *http.Request) {
	slog.Info("Got PATCH!")
}

func (s *Server) StaticHandler(w http.ResponseWriter, r *http.Request) {
	name := chi.URLParam(r, "name")
	if name == "" {
		name = "index.html"
	}

	path := fmt.Sprintf("static/%s", name)

	if strings.Contains(name, ".js") {
		w.Header().Add("Content-Type", "text/javascript")
	}
	if strings.Contains(name, ".css") {
		w.Header().Add("Content-Type", "text/css")
	}

	f, err := static.Open(path)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	defer f.Close()

	io.Copy(w, f)
}

func (s *Server) PostOfferHandler(w http.ResponseWriter, r *http.Request) {
	defer r.Body.Close()
	streamKey := chi.URLParam(r, "streamKey")

	stream, err := s.streams.Get(streamKey)
	if err != nil {
		slog.Warn("Unable to fetch stream.", "error", err, "stream_key", streamKey)
		w.WriteHeader(http.StatusNotFound)
		return
	}

	var buf bytes.Buffer
	io.Copy(&buf, r.Body)

	offer := &webrtc.SessionDescription{Type: webrtc.SDPTypeOffer, SDP: buf.String()}

	answer, err := stream.AddListener(offer)
	if err != nil {
		slog.Warn("Unable to add peer.", "error", err, "stream_key", streamKey, "buf", buf.String())
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	io.WriteString(w, answer.SDP)

	slog.Info("Got listener for stream.", "stream_key", streamKey)
}

func (s *Server) DeleteHandler(w http.ResponseWriter, r *http.Request) {
	streamKey := chi.URLParam(r, "streamKey")
	s.streams.Delete(streamKey)
	slog.Info("Deleted stream.", "streamKey", streamKey)
	w.WriteHeader(http.StatusOK)
}

func (s *Server) ListHandler(w http.ResponseWriter, r *http.Request) {
	streams := s.streams.List()
	enc := json.NewEncoder(w)
	enc.Encode(&streams)
}

func (s *Server) WhipHandler(w http.ResponseWriter, r *http.Request) {
	defer r.Body.Close()

	streamKey, err := streamKeyFromHeader(r.Header.Get("Authorization"))
	if err != nil {
		slog.Warn("Incoming request with invalid auth header.")
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	var buf bytes.Buffer

	io.Copy(&buf, r.Body)

	offer := &webrtc.SessionDescription{Type: webrtc.SDPTypeOffer, SDP: buf.String()}

	offerSdp := sdp.SessionDescription{}
	if err := offerSdp.Unmarshal(buf.Bytes()); err != nil {
		panic(err)
	}

	slog.Info("Got SDP.", "id", offerSdp.Origin.SessionID, "ip", offerSdp.Origin.UnicastAddress)

	answer, err := s.streams.Add(streamKey, offer)
	if err != nil {
		slog.Error("Error adding stream.", "error", err)
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	w.Header().Add("Location", fmt.Sprintf("/whip/%s", streamKey))
	w.Header().Add("Link", "stun:stun.l.google.com:19302; rel=\"ice-server\";")
	w.WriteHeader(http.StatusCreated)
	if _, err := io.WriteString(w, answer.SDP); err != nil {
		slog.Error("Error writing response.", "error", err)
	}
	slog.Info("Wrote answer.")
}

func streamKeyFromHeader(header string) (string, error) {
	split := strings.Split(header, " ")
	if len(split) != 2 {
		return "", fmt.Errorf("invalid header")
	}

	return split[1], nil
}