From a1d61e05ce60d51f6ed29bea9289e348cc9d5a73 Mon Sep 17 00:00:00 2001
From: Emilien <4016501+unixfox@users.noreply.github.com>
Date: Sat, 14 Jun 2025 18:01:44 +0200
Subject: [PATCH] chore: add the suggestions

---
 src/invidious/routes/before_all.cr          |  1 +
 src/invidious/routes/companion.cr           |  6 +-----
 src/invidious/yt_backend/connection_pool.cr | 15 ++++++++++-----
 src/invidious/yt_backend/youtube_api.cr     | 16 ++++------------
 4 files changed, 16 insertions(+), 22 deletions(-)

diff --git a/src/invidious/routes/before_all.cr b/src/invidious/routes/before_all.cr
index b5269668..63b935ec 100644
--- a/src/invidious/routes/before_all.cr
+++ b/src/invidious/routes/before_all.cr
@@ -63,6 +63,7 @@ module Invidious::Routes::BeforeAll
                 "/videoplayback",
                 "/latest_version",
                 "/download",
+                "/companion/",
               }.any? { |r| env.request.resource.starts_with? r }
 
     if env.request.cookies.has_key? "SID"
diff --git a/src/invidious/routes/companion.cr b/src/invidious/routes/companion.cr
index 23b62e9d..cd7ed422 100644
--- a/src/invidious/routes/companion.cr
+++ b/src/invidious/routes/companion.cr
@@ -28,10 +28,6 @@ module Invidious::Routes::Companion
       env.response.headers[key] = value
     end
 
-    if response.status_code >= 300
-      return env.response.headers.delete("Transfer-Encoding")
-    end
-
-    return proxy_file(response, env)
+    return IO.copy response.body_io, env.response
   end
 end
diff --git a/src/invidious/yt_backend/connection_pool.cr b/src/invidious/yt_backend/connection_pool.cr
index 97ce7c40..45455a8a 100644
--- a/src/invidious/yt_backend/connection_pool.cr
+++ b/src/invidious/yt_backend/connection_pool.cr
@@ -46,13 +46,18 @@ struct YoutubeConnectionPool
   end
 end
 
-class CompanionWrapper
+# Packages a `HTTP::Client` to an Invidious companion instance alongside the configuration for that instance.
+#  
+# This is used as the resource for the `CompanionPool` as to allow the ability to
+# proxy the requests to Invidious companion from Invidious directly.
+# Instead of setting up routes in a reverse proxy.
+struct CompanionWrapper
   property client : HTTP::Client
   property companion : Config::CompanionConfig
 
   def initialize(companion : Config::CompanionConfig)
     @companion = companion
-    @client = HTTP::Client.new(companion.private_url)
+    @client = make_client(companion.private_url, use_http_proxy: false)
   end
 
   def close
@@ -73,7 +78,7 @@ struct CompanionConnectionPool
 
     @pool = DB::Pool(CompanionWrapper).new(options) do
       companion = CONFIG.invidious_companion.sample
-      client = make_client(companion.private_url, use_http_proxy: false)
+      make_client(companion.private_url, use_http_proxy: false)
       CompanionWrapper.new(companion: companion)
     end
   end
@@ -84,10 +89,10 @@ struct CompanionConnectionPool
     begin
       response = yield wrapper
     rescue ex
-      wrapper.client.close
+      wrapper.close
 
       companion = CONFIG.invidious_companion.sample
-      client = make_client(companion.private_url, use_http_proxy: false)
+      make_client(companion.private_url, use_http_proxy: false)
       wrapper = CompanionWrapper.new(companion: companion)
 
       response = yield wrapper
diff --git a/src/invidious/yt_backend/youtube_api.cr b/src/invidious/yt_backend/youtube_api.cr
index 455aa836..b1775070 100644
--- a/src/invidious/yt_backend/youtube_api.cr
+++ b/src/invidious/yt_backend/youtube_api.cr
@@ -695,25 +695,17 @@ module YoutubeAPI
     # Send the POST request
 
     begin
-      response_body = ""
+      response_body = Hash(String, JSON::Any).new
 
       COMPANION_POOL.client do |wrapper|
         companion_base_url = wrapper.companion.private_url.path
-        puts "Using companion: #{wrapper.companion.private_url}"
 
-        response = wrapper.client.post(companion_base_url + endpoint, headers: headers, body: data.to_json)
-        response_body = response.body
-
-        if response.status_code != 200
-          raise Exception.new(
-            "Error while communicating with Invidious companion: " \
-            "status code: #{response.status_code} and body: #{response_body.dump}"
-          )
+        wrapper.client.post("#{companion_base_url}#{endpoint}", headers: headers, body: data.to_json) do | response |
+          response_body = JSON.parse(response.body_io).as_h
         end
       end
 
-      # Convert result to Hash
-      return JSON.parse(response_body).as_h
+      return response_body
     rescue ex
       raise InfoException.new("Error while communicating with Invidious companion: " + (ex.message || "no extra info found"))
     end