demo live stream

server/api.go

@@ -1,11 +1,22 @@
 package main
 
 import (
+	"context"
+	"errors"
 	"io"
+	"net/http"
+
+	"github.com/gorilla/websocket"
 
 	"github.com/pocketbase/pocketbase/core"
 )
 
+var liveWSUpgrader = websocket.Upgrader{
+	CheckOrigin: func(_ *http.Request) bool {
+		return true
+	},
+}
+
 func registerAPI(se *core.ServeEvent) {
 	group := se.Router.Group("/api/cctv")
 
@@ -31,11 +42,73 @@ func registerAPI(se *core.ServeEvent) {
 	})
 
 	group.GET("/live/{streamId}", func(e *core.RequestEvent) error {
+		return e.BadRequestError("Use /api/cctv/live/ws/{streamId} for live streaming", nil)
+	})
+
+	group.GET("/live/ws/{streamId}", func(e *core.RequestEvent) error {
 		streamId := e.Request.PathValue("streamId")
 		if streamId == "" {
 			return e.BadRequestError("Missing stream ID", nil)
 		}
-		ingestService.SubscribeLive(e.Request.Context(), streamId)
-		return nil
+
+		conn, err := liveWSUpgrader.Upgrade(e.Response, e.Request, nil)
+		if err != nil {
+			return e.InternalServerError("Failed to upgrade websocket connection", err)
+		}
+		defer conn.Close()
+
+		ctx, cancel := context.WithCancel(e.Request.Context())
+		defer cancel()
+
+		stream, err := ingestService.SubscribeLive(ctx, streamId)
+		if err != nil {
+			writeLiveWSError(conn, err.Error())
+			return nil
+		}
+
+		clientDone := make(chan struct{})
+		go func() {
+			defer close(clientDone)
+			for {
+				if _, _, err := conn.ReadMessage(); err != nil {
+					return
+				}
+			}
+		}()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return nil
+			case <-clientDone:
+				return nil
+			case chunk, ok := <-stream:
+				if !ok {
+					writeLiveWSError(conn, "live stream ended")
+					return nil
+				}
+				if chunk == nil {
+					continue
+				}
+				payload, err := io.ReadAll(chunk)
+				if err != nil {
+					writeLiveWSError(conn, "failed to read stream chunk")
+					return nil
+				}
+				if err := conn.WriteMessage(websocket.BinaryMessage, payload); err != nil {
+					if errors.Is(err, websocket.ErrCloseSent) {
+						return nil
+					}
+					return nil
+				}
+			}
+		}
 	})
 }
+
+func writeLiveWSError(conn *websocket.Conn, msg string) {
+	payload := make([]byte, len(msg)+1)
+	payload[0] = 0x03
+	copy(payload[1:], []byte(msg))
+	_ = conn.WriteMessage(websocket.BinaryMessage, payload)
+}