From f7dbf2bdd4f38fed72ad823be1bc86b727aafdb0 Mon Sep 17 00:00:00 2001 From: Omar Roth Date: Wed, 3 Jul 2019 13:22:52 -0500 Subject: [PATCH] Add 'pipe' for proxying assets --- src/invidious/helpers/helpers.cr | 41 ++++++++++++++++++++++---------- 1 file changed, 28 insertions(+), 13 deletions(-) diff --git a/src/invidious/helpers/helpers.cr b/src/invidious/helpers/helpers.cr index 039ac55b..89fdfe91 100644 --- a/src/invidious/helpers/helpers.cr +++ b/src/invidious/helpers/helpers.cr @@ -638,29 +638,44 @@ def cache_annotation(db, id, annotations) end def proxy_file(response, env) - if !response.body_io? - return - end - if response.headers.includes_word?("Content-Encoding", "gzip") Gzip::Writer.open(env.response) do |deflate| - copy_in_chunks(response.body_io, deflate) + response.pipe(deflate) end elsif response.headers.includes_word?("Content-Encoding", "deflate") Flate::Writer.open(env.response) do |deflate| - copy_in_chunks(response.body_io, deflate) + response.pipe(deflate) end else - copy_in_chunks(response.body_io, env.response) + response.pipe(env.response) end end -# https://stackoverflow.com/a/44802810 <3 -def copy_in_chunks(input, output, chunk_size = 4096) - size = 1 - while size > 0 - size = IO.copy(input, output, chunk_size) - Fiber.yield +class HTTP::Client::Response + def pipe(io) + HTTP.serialize_body(io, headers, @body, @body_io, @version) + end +end + +# Supports serialize_body without first writing headers +module HTTP + def self.serialize_body(io, headers, body, body_io, version) + if body + io << body + elsif body_io + content_length = content_length(headers) + if content_length + copied = IO.copy(body_io, io) + if copied != content_length + raise ArgumentError.new("Content-Length header is #{content_length} but body had #{copied} bytes") + end + elsif Client::Response.supports_chunked?(version) + headers["Transfer-Encoding"] = "chunked" + serialize_chunked_body(io, body_io) + else + io << body + end + end end end