Proxy: de-duplicate the code that copies headers

src/invidious.cr

@@ -71,8 +71,6 @@ CHARS_SAFE         = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz012345
 TEST_IDS           = {"AgbeGFYluEA", "BaW_jenozKc", "a9LDPn-MO4I", "ddFvjfvPnqk", "iqKdEhx-dD4"}
 MAX_ITEMS_PER_PAGE = 1500
 
-REQUEST_HEADERS_WHITELIST  = {"accept", "accept-encoding", "cache-control", "content-length", "if-none-match", "range"}
-RESPONSE_HEADERS_BLACKLIST = {"access-control-allow-origin", "alt-svc", "server"}
 HTTP_CHUNK_SIZE = 10485760 # ~10MB
 
 CURRENT_BRANCH  = {{ "#{`git branch | sed -n '/* /s///p'`.strip}" }}

src/invidious/routes/images.cr

@@ -6,22 +6,14 @@ module Invidious::Routes::Images
     headers = HTTP::Headers.new
     headers[":authority"] = "yt3.ggpht.com" if CONFIG.use_quic
 
-    REQUEST_HEADERS_WHITELIST.each do |header|
+    MediaProxy.copy_request_headers(from: env.request.headers, to: headers)
-      if env.request.headers[header]?
-        headers[header] = env.request.headers[header]
-      end
-    end
 
     # We're encapsulating this into a proc in order to easily reuse this
     # portion of the code for each request block below.
     request_proc = ->(response : HTTP::Client::Response) {
       env.response.status_code = response.status_code
-      response.headers.each do |key, value|
-        if !RESPONSE_HEADERS_BLACKLIST.includes?(key.downcase)
-          env.response.headers[key] = value
-        end
-      end
 
+      MediaProxy.copy_response_headers(from: response.headers, to: env.response.headers)
       env.response.headers["Access-Control-Allow-Origin"] = "*"
 
       if response.status_code >= 300
@@ -64,20 +56,12 @@ module Invidious::Routes::Images
     headers = HTTP::Headers.new
     headers[":authority"] = "#{authority}.ytimg.com" if CONFIG.use_quic
 
-    REQUEST_HEADERS_WHITELIST.each do |header|
+    MediaProxy.copy_request_headers(from: env.request.headers, to: headers)
-      if env.request.headers[header]?
-        headers[header] = env.request.headers[header]
-      end
-    end
 
     request_proc = ->(response : HTTP::Client::Response) {
       env.response.status_code = response.status_code
-      response.headers.each do |key, value|
-        if !RESPONSE_HEADERS_BLACKLIST.includes?(key.downcase)
-          env.response.headers[key] = value
-        end
-      end
 
+      MediaProxy.copy_response_headers(from: response.headers, to: env.response.headers)
       env.response.headers["Connection"] = "close"
       env.response.headers["Access-Control-Allow-Origin"] = "*"
 
@@ -112,20 +96,12 @@ module Invidious::Routes::Images
     headers = HTTP::Headers.new
     headers[":authority"] = "i9.ytimg.com" if CONFIG.use_quic
 
-    REQUEST_HEADERS_WHITELIST.each do |header|
+    MediaProxy.copy_request_headers(from: env.request.headers, to: headers)
-      if env.request.headers[header]?
-        headers[header] = env.request.headers[header]
-      end
-    end
 
     request_proc = ->(response : HTTP::Client::Response) {
       env.response.status_code = response.status_code
-      response.headers.each do |key, value|
-        if !RESPONSE_HEADERS_BLACKLIST.includes?(key.downcase)
-          env.response.headers[key] = value
-        end
-      end
 
+      MediaProxy.copy_response_headers(from: response.headers, to: env.response.headers)
       env.response.headers["Access-Control-Allow-Origin"] = "*"
 
       if response.status_code >= 300 && response.status_code != 404
@@ -152,21 +128,13 @@ module Invidious::Routes::Images
 
   def self.yts_image(env)
     headers = HTTP::Headers.new
-    REQUEST_HEADERS_WHITELIST.each do |header|
+    MediaProxy.copy_request_headers(from: env.request.headers, to: headers)
-      if env.request.headers[header]?
-        headers[header] = env.request.headers[header]
-      end
-    end
 
     begin
       YT_POOL.client &.get(env.request.resource, headers) do |response|
         env.response.status_code = response.status_code
-        response.headers.each do |key, value|
-          if !RESPONSE_HEADERS_BLACKLIST.includes?(key.downcase)
-            env.response.headers[key] = value
-          end
-        end
 
+        MediaProxy.copy_response_headers(from: response.headers, to: env.response.headers)
         env.response.headers["Access-Control-Allow-Origin"] = "*"
 
         if response.status_code >= 300 && response.status_code != 404
@@ -209,20 +177,12 @@ module Invidious::Routes::Images
 
     url = "/vi/#{id}/#{name}"
 
-    REQUEST_HEADERS_WHITELIST.each do |header|
+    MediaProxy.copy_request_headers(from: env.request.headers, to: headers)
-      if env.request.headers[header]?
-        headers[header] = env.request.headers[header]
-      end
-    end
 
     request_proc = ->(response : HTTP::Client::Response) {
       env.response.status_code = response.status_code
-      response.headers.each do |key, value|
-        if !RESPONSE_HEADERS_BLACKLIST.includes?(key.downcase)
-          env.response.headers[key] = value
-        end
-      end
 
+      MediaProxy.copy_response_headers(from: response.headers, to: env.response.headers)
       env.response.headers["Access-Control-Allow-Origin"] = "*"
 
       if response.status_code >= 300 && response.status_code != 404

src/invidious/routes/video_playback.cr

@@ -29,11 +29,7 @@ module Invidious::Routes::VideoPlayback
     url = "/videoplayback?#{query_params}"
 
     headers = HTTP::Headers.new
-    REQUEST_HEADERS_WHITELIST.each do |header|
+    MediaProxy.copy_request_headers(from: env.request.headers, to: headers)
-      if env.request.headers[header]?
-        headers[header] = env.request.headers[header]
-      end
-    end
 
     # See: https://github.com/iv-org/invidious/issues/3302
     range_header = env.request.headers["Range"]?
@@ -92,12 +88,7 @@ module Invidious::Routes::VideoPlayback
 
       begin
         client.get(url, headers) do |resp|
-          resp.headers.each do |key, value|
+          MediaProxy.copy_response_headers(from: resp.headers, to: env.response.headers)
-            if !RESPONSE_HEADERS_BLACKLIST.includes?(key.downcase)
-              env.response.headers[key] = value
-            end
-          end
-
           env.response.headers["Access-Control-Allow-Origin"] = "*"
 
           if location = resp.headers["Location"]?
@@ -150,12 +141,8 @@ module Invidious::Routes::VideoPlayback
                 env.response.status_code = resp.status_code
               end
 
-              resp.headers.each do |key, value|
+              MediaProxy.copy_response_headers(from: resp.headers, to: env.response.headers)
-                if !RESPONSE_HEADERS_BLACKLIST.includes?(key.downcase) && key.downcase != "content-range"
+              env.response.headers.delete("Content-Range") # Important!
-                  env.response.headers[key] = value
-                end
-              end
-
               env.response.headers["Access-Control-Allow-Origin"] = "*"
 
               if location = resp.headers["Location"]?

src/invidious/yt_backend/media_proxy.cr

@@ -0,0 +1,38 @@
+module Invidious::MediaProxy
+  extend self
+
+  # -------------------
+  #  Constants
+  # -------------------
+
+  private REQUEST_HEADERS_WHITELIST = {
+    "accept", "accept-encoding", "cache-control",
+    "content-length", "if-none-match", "range",
+  }
+
+  private RESPONSE_HEADERS_BLACKLIST = {
+    "access-control-allow-origin", "alt-svc", "server",
+  }
+
+  # -------------------
+  #  Headers functions
+  # -------------------
+
+  # Copy only the selected headers from the client to youtube servers
+  # (in general, from `env.request` to a temporary `HTTP::Headers` object).
+  def copy_request_headers(*, from : HTTP::Headers, to : HTTP::Headers)
+    REQUEST_HEADERS_WHITELIST.each do |header|
+      to[header] = from[header] if from[header]?
+    end
+  end
+
+  # Copy only the selected headers from youtube servers to the client
+  # (generally, from a response block to `env.response`).
+  def copy_response_headers(*, from : HTTP::Headers, to : HTTP::Headers)
+    from.each do |key, value|
+      if !RESPONSE_HEADERS_BLACKLIST.includes?(key.downcase)
+        to[key] = value
+      end
+    end
+  end
+end