From 208bb2d72ffd9559b6df64f1c038392bfa7083ea Mon Sep 17 00:00:00 2001
From: Omar Roth <omarroth@protonmail.com>
Date: Sun, 26 May 2019 09:41:12 -0500
Subject: [PATCH] Catch connection reset when proxying files

---
 src/invidious.cr                 | 20 ++++++++++++++++----
 src/invidious/helpers/helpers.cr | 21 +++++++++------------
 2 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/src/invidious.cr b/src/invidious.cr
index adb0b659..af0c549d 100644
--- a/src/invidious.cr
+++ b/src/invidious.cr
@@ -4964,11 +4964,12 @@ get "/videoplayback" do |env|
   end
 
   client = make_client(URI.parse(host), proxies, region)
+  begin
   client.get(url, headers) do |response|
     env.response.status_code = response.status_code
 
     response.headers.each do |key, value|
-      if !{"Access-Control-Allow-Origin", "Alt-Svc"}.includes? key
+        if !{"Access-Control-Allow-Origin", "Alt-Svc", "Server"}.includes? key
         env.response.headers[key] = value
       end
     end
@@ -4996,6 +4997,8 @@ get "/videoplayback" do |env|
 
     proxy_file(response, env)
   end
+  rescue ex
+end
 end
 
 # We need this so the below route works as expected
@@ -5014,9 +5017,10 @@ get "/ggpht/*" do |env|
     end
   end
 
+  begin
   client.get(url, headers) do |response|
     response.headers.each do |key, value|
-      if !{"Access-Control-Allow-Origin", "Alt-Svc"}.includes? key
+        if !{"Access-Control-Allow-Origin", "Alt-Svc", "Server"}.includes? key
         env.response.headers[key] = value
       end
     end
@@ -5029,6 +5033,8 @@ get "/ggpht/*" do |env|
 
     proxy_file(response, env)
   end
+  rescue ex
+  end
 end
 
 options "/sb/:id/:storyboard/:index" do |env|
@@ -5059,10 +5065,11 @@ get "/sb/:id/:storyboard/:index" do |env|
     end
   end
 
+  begin
   client.get(url, headers) do |response|
     env.response.status_code = response.status_code
     response.headers.each do |key, value|
-      if !{"Access-Control-Allow-Origin", "Alt-Svc"}.includes? key
+        if !{"Access-Control-Allow-Origin", "Alt-Svc", "Server"}.includes? key
         env.response.headers[key] = value
       end
     end
@@ -5075,6 +5082,8 @@ get "/sb/:id/:storyboard/:index" do |env|
 
     proxy_file(response, env)
   end
+  rescue ex
+  end
 end
 
 get "/vi/:id/:name" do |env|
@@ -5101,10 +5110,11 @@ get "/vi/:id/:name" do |env|
     end
   end
 
+  begin
   client.get(url, headers) do |response|
     env.response.status_code = response.status_code
     response.headers.each do |key, value|
-      if !{"Access-Control-Allow-Origin", "Alt-Svc"}.includes? key
+        if !{"Access-Control-Allow-Origin", "Alt-Svc", "Server"}.includes? key
         env.response.headers[key] = value
       end
     end
@@ -5117,6 +5127,8 @@ get "/vi/:id/:name" do |env|
 
     proxy_file(response, env)
   end
+  rescue ex
+  end
 end
 
 # Undocumented, creates anonymous playlist with specified 'video_ids'
diff --git a/src/invidious/helpers/helpers.cr b/src/invidious/helpers/helpers.cr
index a9f27d5e..476038c7 100644
--- a/src/invidious/helpers/helpers.cr
+++ b/src/invidious/helpers/helpers.cr
@@ -639,19 +639,16 @@ def proxy_file(response, env)
     return
   end
 
-  begin
-    if response.headers.includes_word?("Content-Encoding", "gzip")
-      Gzip::Writer.open(env.response) do |deflate|
-        copy_in_chunks(response.body_io, deflate)
-      end
-    elsif response.headers.includes_word?("Content-Encoding", "deflate")
-      Flate::Writer.open(env.response) do |deflate|
-        copy_in_chunks(response.body_io, deflate)
-      end
-    else
-      copy_in_chunks(response.body_io, env.response)
+  if response.headers.includes_word?("Content-Encoding", "gzip")
+    Gzip::Writer.open(env.response) do |deflate|
+      copy_in_chunks(response.body_io, deflate)
     end
-  rescue ex
+  elsif response.headers.includes_word?("Content-Encoding", "deflate")
+    Flate::Writer.open(env.response) do |deflate|
+      copy_in_chunks(response.body_io, deflate)
+    end
+  else
+    copy_in_chunks(response.body_io, env.response)
   end
 end