From 6cb834a18d770f859719102af78a018cc53e03f1 Mon Sep 17 00:00:00 2001
From: Omar Roth <omarroth@hotmail.com>
Date: Mon, 17 Sep 2018 09:38:52 -0500
Subject: [PATCH] Add support for 304 in thumbnails

---
 src/invidious.cr | 45 ++++++++++++++++++++++++++++++++-------------
 1 file changed, 32 insertions(+), 13 deletions(-)

diff --git a/src/invidious.cr b/src/invidious.cr
index a5a5a2299..5c809eab7 100644
--- a/src/invidious.cr
+++ b/src/invidious.cr
@@ -3003,26 +3003,45 @@ get "/vi/:id/:name" do |env|
   end
   url = "/vi/#{id}/#{name}"
 
-  client.get(url) do |response|
-    env.response.status_code = response.status_code
+  headers = env.request.headers
+  headers.delete("Host")
+  headers.delete("Cookie")
+  headers.delete("User-Agent")
+  headers.delete("Referer")
 
+  client.get(url, headers) do |response|
+    env.response.status_code = response.status_code
+    puts response.headers.inspect
     response.headers.each do |key, value|
       env.response.headers[key] = value
     end
 
-    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
+    if response.status_code == 304
       break
     end
+
+    chunk_size = 4096
+    size = chunk_size
+    if response.headers.includes_word?("Content-Encoding", "gzip")
+      Gzip::Writer.open(env.response) do |deflate|
+        until size < chunk_size
+          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 < chunk_size
+          size = IO.copy(response.body_io, deflate)
+          env.response.flush
+        end
+      end
+    else
+      until size < chunk_size
+        size = IO.copy(response.body_io, env.response, chunk_size)
+        env.response.flush
+      end
+    end
   end
 end