From 8ec62941146697faeddc63386866ad7f89ceecc5 Mon Sep 17 00:00:00 2001
From: Fijxu <fijxu@nadeko.net>
Date: Wed, 2 Apr 2025 20:58:16 -0300
Subject: [PATCH] add CSP based on backend selected by the user

---
 src/invidious/helpers/backend_info.cr | 29 +++++++++------------------
 src/invidious/routes/before_all.cr    |  3 +--
 2 files changed, 10 insertions(+), 22 deletions(-)

diff --git a/src/invidious/helpers/backend_info.cr b/src/invidious/helpers/backend_info.cr
index 2e7b1c39..ab61d500 100644
--- a/src/invidious/helpers/backend_info.cr
+++ b/src/invidious/helpers/backend_info.cr
@@ -2,12 +2,11 @@ module BackendInfo
   extend self
   @@exvpp_url : Array(String) = Array.new(CONFIG.invidious_companion.size, "")
   @@status : Array(Int32) = Array.new(CONFIG.invidious_companion.size, 0)
-  @@extra_media_csp : String = ""
-  @@extra_connect_csp : String = ""
+  @@extra_media_csp : Array(String) = Array.new(CONFIG.invidious_companion.size, "")
+  @@extra_connect_csp : Array(String) = Array.new(CONFIG.invidious_companion.size, "")
 
   def check_backends
     check_companion()
-    generate_csp()
   end
 
   private def check_companion
@@ -16,6 +15,7 @@ module BackendInfo
         begin
           response = HTTP::Client.get "#{companion.private_url}/healthz"
           if response.status_code == 200
+            generate_csp(companion.public_url.to_s, index)
             check_videoplayback_proxy(companion, index)
           else
             @@status[index] = 0
@@ -40,6 +40,7 @@ module BackendInfo
         begin
           exvpp_health = HTTP::Client.get "#{exvpp_url}/health"
           if exvpp_health.status_code == 200
+            generate_csp(exvpp_url, index)
             @@status[index] = 2
             return
           else
@@ -53,21 +54,9 @@ module BackendInfo
     end
   end
 
-  private def generate_csp
-    @@extra_media_csp = ""
-    @@extra_connect_csp = ""
-
-    CONFIG.invidious_companion.each do |companion|
-      @@extra_media_csp += " #{companion.public_url}"
-      @@extra_connect_csp += " #{companion.public_url}"
-    end
-    exvpp_urls = self.get_exvpp
-    exvpp_urls.each do |exvpp_url|
-      if !exvpp_url.empty?
-        @@extra_media_csp += " #{exvpp_url}"
-        @@extra_connect_csp += " #{exvpp_url}"
-      end
-    end
+  private def generate_csp(item, index)
+    @@extra_media_csp[index] += " #{item}"
+    @@extra_connect_csp[index] += " #{item}"
   end
 
   def get_status
@@ -78,7 +67,7 @@ module BackendInfo
     return @@exvpp_url
   end
 
-  def get_csp
-    return @@extra_media_csp, @@extra_connect_csp
+  def get_csp(backend : Int32)
+    return @@extra_media_csp[backend], @@extra_connect_csp[backend]
   end
 end
diff --git a/src/invidious/routes/before_all.cr b/src/invidious/routes/before_all.cr
index e1edddbb..7ae7239e 100644
--- a/src/invidious/routes/before_all.cr
+++ b/src/invidious/routes/before_all.cr
@@ -50,8 +50,7 @@ module Invidious::Routes::BeforeAll
 
         env.set "current_companion", current_companion
       end
-
-      extra_media_csp, extra_connect_csp = BackendInfo.get_csp
+      extra_media_csp, extra_connect_csp = BackendInfo.get_csp(env.get("current_companion").as(Int32))
     end
 
     if !CONFIG.external_videoplayback_proxy.empty?