diff --git a/spec/invidious/subtitle_cache_spec.cr b/spec/invidious/subtitle_cache_spec.cr index be12e5502..8245e1d23 100644 --- a/spec/invidious/subtitle_cache_spec.cr +++ b/spec/invidious/subtitle_cache_spec.cr @@ -49,7 +49,9 @@ Spectator.describe Invidious::SubtitleCache do expect(Invidious::SubtitleCache.valid_vtt?("WEBVTT\n\n00:00:00.000 --> 00:00:01.000\nHello")).to be_true expect(Invidious::SubtitleCache.valid_vtt?("\uFEFFWEBVTT\n\n00:00:00.000 --> 00:00:01.000\nHello")).to be_true expect(Invidious::SubtitleCache.valid_vtt?(" \nWEBVTT\n\n")).to be_true + expect(Invidious::SubtitleCache.valid_vtt?("WEBVTT")).to be_true expect(Invidious::SubtitleCache.valid_vtt?("")).to be_false + expect(Invidious::SubtitleCache.valid_vtt?("WEBVTT-error")).to be_false expect(Invidious::SubtitleCache.valid_vtt?("{\"error\": \"not found\"}")).to be_false expect(Invidious::SubtitleCache.valid_vtt?("")).to be_false end @@ -190,5 +192,13 @@ Spectator.describe Invidious::SubtitleCache do expect(result.response.try &.body).to eq(oversized) expect(cache.size).to eq(0) end + + it "stops reading bodies beyond the entry size limit" do + valid_body = "WEBVTT\n" + expect(Invidious::SubtitleCache.read_limited_body(IO::Memory.new(valid_body))).to eq(valid_body) + + oversized_body = "x" * (Invidious::SubtitleCache::MAX_ENTRY_BYTES + 1) + expect(Invidious::SubtitleCache.read_limited_body(IO::Memory.new(oversized_body))).to be_nil + end end end diff --git a/src/invidious/helpers/subtitle_cache.cr b/src/invidious/helpers/subtitle_cache.cr index fa010e8ae..9895fc0f0 100644 --- a/src/invidious/helpers/subtitle_cache.cr +++ b/src/invidious/helpers/subtitle_cache.cr @@ -37,6 +37,7 @@ module Invidious DEFAULT_MAX_BYTES = 64 * 1024 * 1024 # 64 MiB DEFAULT_TTL = 6.hours # 21600 seconds MAX_ENTRY_BYTES = 2 * 1024 * 1024 # 2 MiB + READ_CHUNK_BYTES = 64 * 1024 getter max_entries : Int32 getter max_bytes : Int32 @@ -55,7 +56,26 @@ module Invidious def self.valid_vtt?(body : String) : Bool return false if body.empty? trimmed = body.lstrip("\uFEFF \t\r\n") - trimmed.starts_with?("WEBVTT") + return false unless trimmed.starts_with?("WEBVTT") + + header_end = "WEBVTT".size + header_end == trimmed.size || " \t\r\n".includes?(trimmed[header_end]) + end + + def self.read_limited_body(input : IO, limit : Int32 = MAX_ENTRY_BYTES) : String? + output = IO::Memory.new + buffer = Bytes.new(READ_CHUNK_BYTES) + + while output.bytesize < limit + remaining = limit - output.bytesize + chunk_size = Math.min(READ_CHUNK_BYTES, remaining) + bytes_read = input.read(buffer[0, chunk_size]) + break if bytes_read == 0 + output.write(buffer[0, bytes_read]) + end + + return nil if input.read_byte + output.to_s end def size : Int32 diff --git a/src/invidious/routes/companion.cr b/src/invidious/routes/companion.cr index a99317824..f13d6e07b 100644 --- a/src/invidious/routes/companion.cr +++ b/src/invidious/routes/companion.cr @@ -107,7 +107,8 @@ module Invidious::Routes::Companion private def self.fetch_from_companion(url : String, headers : HTTP::Headers) : Invidious::SubtitleCache::Response? COMPANION_POOL.client do |wrapper| wrapper.client.get(url, headers) do |resp| - body = resp.body_io.gets_to_end + body = Invidious::SubtitleCache.read_limited_body(resp.body_io) + return nil unless body response_headers = HTTP::Headers.new resp.headers.each do |key, value| response_headers[key] = value