From 6c04f9199f1fe5e468eab8b12151098fa38be69b Mon Sep 17 00:00:00 2001
From: WithoutPants <53250216+WithoutPants@users.noreply.github.com>
Date: Mon, 10 Oct 2022 10:11:51 +1100
Subject: [PATCH] Add apikey to streams (#2981)

---
 internal/api/resolver_model_scene.go |  3 ++-
 internal/api/urlbuilders/scene.go    | 18 ++++++++++++-----
 internal/manager/scene.go            | 30 +++++++++++++++++++---------
 3 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/internal/api/resolver_model_scene.go b/internal/api/resolver_model_scene.go
index 3be095340..624b08b50 100644
--- a/internal/api/resolver_model_scene.go
+++ b/internal/api/resolver_model_scene.go
@@ -168,7 +168,7 @@ func (r *sceneResolver) Paths(ctx context.Context, obj *models.Scene) (*ScenePat
 	builder.APIKey = config.GetAPIKey()
 	screenshotPath := builder.GetScreenshotURL(obj.UpdatedAt)
 	previewPath := builder.GetStreamPreviewURL()
-	streamPath := builder.GetStreamURL()
+	streamPath := builder.GetStreamURL().String()
 	webpPath := builder.GetStreamPreviewImageURL()
 	vttPath := builder.GetSpriteVTTURL()
 	spritePath := builder.GetSpriteURL()
@@ -357,6 +357,7 @@ func (r *sceneResolver) SceneStreams(ctx context.Context, obj *models.Scene) ([]
 
 	baseURL, _ := ctx.Value(BaseURLCtxKey).(string)
 	builder := urlbuilders.NewSceneURLBuilder(baseURL, obj.ID)
+	builder.APIKey = config.GetAPIKey()
 
 	return manager.GetSceneStreamPaths(obj, builder.GetStreamURL(), config.GetMaxStreamingTranscodeSize())
 }
diff --git a/internal/api/urlbuilders/scene.go b/internal/api/urlbuilders/scene.go
index 650fa0218..014daec95 100644
--- a/internal/api/urlbuilders/scene.go
+++ b/internal/api/urlbuilders/scene.go
@@ -2,6 +2,7 @@ package urlbuilders
 
 import (
 	"fmt"
+	"net/url"
 	"strconv"
 	"time"
 )
@@ -19,12 +20,19 @@ func NewSceneURLBuilder(baseURL string, sceneID int) SceneURLBuilder {
 	}
 }
 
-func (b SceneURLBuilder) GetStreamURL() string {
-	var apiKeyParam string
-	if b.APIKey != "" {
-		apiKeyParam = fmt.Sprintf("?apikey=%s", b.APIKey)
+func (b SceneURLBuilder) GetStreamURL() *url.URL {
+	u, err := url.Parse(fmt.Sprintf("%s/scene/%s/stream", b.BaseURL, b.SceneID))
+	if err != nil {
+		// shouldn't happen
+		panic(err)
 	}
-	return fmt.Sprintf("%s/scene/%s/stream%s", b.BaseURL, b.SceneID, apiKeyParam)
+
+	if b.APIKey != "" {
+		v := u.Query()
+		v.Set("apikey", b.APIKey)
+		u.RawQuery = v.Encode()
+	}
+	return u
 }
 
 func (b SceneURLBuilder) GetStreamPreviewURL() string {
diff --git a/internal/manager/scene.go b/internal/manager/scene.go
index f9693787b..5539ad231 100644
--- a/internal/manager/scene.go
+++ b/internal/manager/scene.go
@@ -2,6 +2,7 @@ package manager
 
 import (
 	"fmt"
+	"net/url"
 
 	"github.com/stashapp/stash/internal/manager/config"
 	"github.com/stashapp/stash/pkg/ffmpeg"
@@ -58,15 +59,20 @@ type SceneStreamEndpoint struct {
 	Label    *string `json:"label"`
 }
 
-func makeStreamEndpoint(streamURL string, streamingResolution models.StreamingResolutionEnum, mimeType, label string) *SceneStreamEndpoint {
+func makeStreamEndpoint(streamURL *url.URL, streamingResolution models.StreamingResolutionEnum, mimeType, label string) *SceneStreamEndpoint {
+	urlCopy := *streamURL
+	v := urlCopy.Query()
+	v.Set("resolution", streamingResolution.String())
+	urlCopy.RawQuery = v.Encode()
+
 	return &SceneStreamEndpoint{
-		URL:      fmt.Sprintf("%s?resolution=%s", streamURL, streamingResolution.String()),
+		URL:      urlCopy.String(),
 		MimeType: &mimeType,
 		Label:    &label,
 	}
 }
 
-func GetSceneStreamPaths(scene *models.Scene, directStreamURL string, maxStreamingTranscodeSize models.StreamingResolutionEnum) ([]*SceneStreamEndpoint, error) {
+func GetSceneStreamPaths(scene *models.Scene, directStreamURL *url.URL, maxStreamingTranscodeSize models.StreamingResolutionEnum) ([]*SceneStreamEndpoint, error) {
 	if scene == nil {
 		return nil, fmt.Errorf("nil scene")
 	}
@@ -93,10 +99,16 @@ func GetSceneStreamPaths(scene *models.Scene, directStreamURL string, maxStreami
 	// don't care if we can't get the container
 	container, _ := GetVideoFileContainer(pf)
 
+	replaceSuffix := func(suffix string) *url.URL {
+		urlCopy := *directStreamURL
+		urlCopy.Path += suffix
+		return &urlCopy
+	}
+
 	if HasTranscode(scene, config.GetInstance().GetVideoFileNamingAlgorithm()) || ffmpeg.IsValidAudioForContainer(audioCodec, container) {
 		label := "Direct stream"
 		ret = append(ret, &SceneStreamEndpoint{
-			URL:      directStreamURL,
+			URL:      directStreamURL.String(),
 			MimeType: &mimeMp4,
 			Label:    &label,
 		})
@@ -106,7 +118,7 @@ func GetSceneStreamPaths(scene *models.Scene, directStreamURL string, maxStreami
 	if container == ffmpeg.Matroska {
 		label := "mkv"
 		ret = append(ret, &SceneStreamEndpoint{
-			URL: directStreamURL + ".mkv",
+			URL: replaceSuffix(".mkv").String(),
 			// set mkv to mp4 to trick the client, since many clients won't try mkv
 			MimeType: &mimeMp4,
 			Label:    &label,
@@ -131,8 +143,8 @@ func GetSceneStreamPaths(scene *models.Scene, directStreamURL string, maxStreami
 	var webmStreams []*SceneStreamEndpoint
 	var mp4Streams []*SceneStreamEndpoint
 
-	webmURL := directStreamURL + ".webm"
-	mp4URL := directStreamURL + ".mp4"
+	webmURL := replaceSuffix(".webm")
+	mp4URL := replaceSuffix(".mp4")
 
 	if includeSceneStreamPath(pf, models.StreamingResolutionEnumFourK, maxStreamingTranscodeSize) {
 		webmStreams = append(webmStreams, makeStreamEndpoint(webmURL, models.StreamingResolutionEnumFourK, mimeMp4, webmLabelFourK))
@@ -164,7 +176,7 @@ func GetSceneStreamPaths(scene *models.Scene, directStreamURL string, maxStreami
 
 	defaultStreams := []*SceneStreamEndpoint{
 		{
-			URL:      directStreamURL + ".webm",
+			URL:      replaceSuffix(".webm").String(),
 			MimeType: &mimeWebm,
 			Label:    &labelWebm,
 		},
@@ -173,7 +185,7 @@ func GetSceneStreamPaths(scene *models.Scene, directStreamURL string, maxStreami
 	ret = append(ret, defaultStreams...)
 
 	hls := SceneStreamEndpoint{
-		URL:      directStreamURL + ".m3u8",
+		URL:      replaceSuffix(".m3u8").String(),
 		MimeType: &mimeHLS,
 		Label:    &labelHLS,
 	}