From ab4df7e078c0d5f9f369670cdc81c2540127e2ce Mon Sep 17 00:00:00 2001
From: Omar Roth <omarroth@protonmail.com>
Date: Sat, 18 May 2019 19:14:58 -0500
Subject: [PATCH] Fix response for proxied assets

---
 src/invidious.cr                 | 81 ++------------------------------
 src/invidious/helpers/helpers.cr | 14 ++++++
 2 files changed, 18 insertions(+), 77 deletions(-)

diff --git a/src/invidious.cr b/src/invidious.cr
index 23573fc18..1f6621848 100644
--- a/src/invidious.cr
+++ b/src/invidious.cr
@@ -5079,17 +5079,7 @@ get "/videoplayback" do |env|
 
     env.response.headers["Access-Control-Allow-Origin"] = "*"
 
-    begin
-      chunk_size = 4096
-      size = 1
-      while size > 0
-        size = IO.copy(response.body_io, env.response.output, chunk_size)
-        env.response.flush
-        Fiber.yield
-      end
-    rescue ex
-      break
-    end
+    proxy_file(response, env)
   end
 end
 
@@ -5119,28 +5109,7 @@ get "/ggpht/*" do |env|
       break
     end
 
-    chunk_size = 4096
-    size = 1
-    if response.headers.includes_word?("Content-Encoding", "gzip")
-      Gzip::Writer.open(env.response) do |deflate|
-        until size == 0
-          size = IO.copy(response.body_io, deflate)
-          env.response.flush
-        end
-      end
-    elsif response.headers.includes_word?("Content-Encoding", "deflate")
-      Flate::Writer.open(env.response) do |deflate|
-        until size == 0
-          size = IO.copy(response.body_io, deflate)
-          env.response.flush
-        end
-      end
-    else
-      until size == 0
-        size = IO.copy(response.body_io, env.response, chunk_size)
-        env.response.flush
-      end
-    end
+    proxy_file(response, env)
   end
 end
 
@@ -5182,28 +5151,7 @@ get "/sb/:id/:storyboard/:index" do |env|
       break
     end
 
-    chunk_size = 4096
-    size = 1
-    if response.headers.includes_word?("Content-Encoding", "gzip")
-      Gzip::Writer.open(env.response) do |deflate|
-        until size == 0
-          size = IO.copy(response.body_io, deflate)
-          env.response.flush
-        end
-      end
-    elsif response.headers.includes_word?("Content-Encoding", "deflate")
-      Flate::Writer.open(env.response) do |deflate|
-        until size == 0
-          size = IO.copy(response.body_io, deflate)
-          env.response.flush
-        end
-      end
-    else
-      until size == 0
-        size = IO.copy(response.body_io, env.response, chunk_size)
-        env.response.flush
-      end
-    end
+    proxy_file(response, env)
   end
 end
 
@@ -5241,28 +5189,7 @@ get "/vi/:id/:name" do |env|
       break
     end
 
-    chunk_size = 4096
-    size = 1
-    if response.headers.includes_word?("Content-Encoding", "gzip")
-      Gzip::Writer.open(env.response) do |deflate|
-        until size == 0
-          size = IO.copy(response.body_io, deflate)
-          env.response.flush
-        end
-      end
-    elsif response.headers.includes_word?("Content-Encoding", "deflate")
-      Flate::Writer.open(env.response) do |deflate|
-        until size == 0
-          size = IO.copy(response.body_io, deflate)
-          env.response.flush
-        end
-      end
-    else
-      until size == 0
-        size = IO.copy(response.body_io, env.response, chunk_size)
-        env.response.flush
-      end
-    end
+    proxy_file(response, env)
   end
 end
 
diff --git a/src/invidious/helpers/helpers.cr b/src/invidious/helpers/helpers.cr
index 2b843087c..9dac5213d 100644
--- a/src/invidious/helpers/helpers.cr
+++ b/src/invidious/helpers/helpers.cr
@@ -629,3 +629,17 @@ def cache_annotation(db, id, annotations)
     db.exec("INSERT INTO annotations VALUES ($1, $2) ON CONFLICT DO NOTHING", id, annotations)
   end
 end
+
+def proxy_file(response, env)
+  if response.headers.includes_word?("Content-Encoding", "gzip")
+    Gzip::Writer.open(env.response) do |deflate|
+      IO.copy(response.body_io, deflate)
+    end
+  elsif response.headers.includes_word?("Content-Encoding", "deflate")
+    Flate::Writer.open(env.response) do |deflate|
+      IO.copy(response.body_io, deflate)
+    end
+  else
+    IO.copy(response.body_io, env.response)
+  end
+end