From 4e2f59ba37071da6247128e29890755a56f9665f Mon Sep 17 00:00:00 2001
From: Fijxu <fijxu@nadeko.net>
Date: Thu, 8 May 2025 00:44:36 -0400
Subject: [PATCH] feat(routes): redirect to the backend selected at the time of
 downloading

---
 src/invidious/frontend/watch_page.cr | 5 +++--
 src/invidious/routes/watch.cr        | 4 +++-
 src/invidious/views/watch.ecr        | 2 +-
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/invidious/frontend/watch_page.cr b/src/invidious/frontend/watch_page.cr
index 00b46a2f..ae1b4f7f 100644
--- a/src/invidious/frontend/watch_page.cr
+++ b/src/invidious/frontend/watch_page.cr
@@ -18,14 +18,15 @@ module Invidious::Frontend::WatchPage
     end
   end
 
-  def download_widget(locale : String, video : Video, video_assets : VideoAssets) : String
+  def download_widget(locale : String, video : Video, video_assets : VideoAssets, env : HTTP::Server::Context) : String
     if CONFIG.disabled?("downloads")
       return "<p id=\"download\">#{translate(locale, "Download is disabled")}</p>"
     end
 
     url = "/download"
     if (CONFIG.invidious_companion.present?)
-      invidious_companion = CONFIG.invidious_companion.sample
+      current_companion = env.get("current_companion").as(Int32)
+      invidious_companion = CONFIG.invidious_companion[current_companion]
       url = "#{invidious_companion.public_url}/download?check=#{invidious_companion_encrypt(video.id)}"
     end
 
diff --git a/src/invidious/routes/watch.cr b/src/invidious/routes/watch.cr
index 0d55bab7..bd1b9481 100644
--- a/src/invidious/routes/watch.cr
+++ b/src/invidious/routes/watch.cr
@@ -312,7 +312,9 @@ module Invidious::Routes::Watch
       return error_template(403, "Administrator has disabled this endpoint.")
     end
     if CONFIG.invidious_companion.present?
-      return error_template(403, "Downloads should be routed through Companion when present")
+      current_companion = env.get("current_companion").as(Int32)
+      invidious_companion = CONFIG.invidious_companion[current_companion]
+      return env.redirect invidious_companion.public_url
     end
 
     title = env.params.body["title"]? || ""
diff --git a/src/invidious/views/watch.ecr b/src/invidious/views/watch.ecr
index d37d6de0..79953121 100644
--- a/src/invidious/views/watch.ecr
+++ b/src/invidious/views/watch.ecr
@@ -203,7 +203,7 @@ we're going to need to do it here in order to allow for translations.
                 <% end %>
             <% end %>
 
-            <%= Invidious::Frontend::WatchPage.download_widget(locale, video, video_assets) %>
+            <%= Invidious::Frontend::WatchPage.download_widget(locale, video, video_assets, env) %>
 
             <p id="views"><i class="icon ion-ios-eye"></i> <%= number_with_separator(video.views) %></p>
             <p id="likes"><i class="icon ion-ios-thumbs-up"></i> <%= number_with_separator(video.likes) %></p>