diff --git a/spec/invidious/subtitle_cache_spec.cr b/spec/invidious/subtitle_cache_spec.cr index a365a9b1c..be12e5502 100644 --- a/spec/invidious/subtitle_cache_spec.cr +++ b/spec/invidious/subtitle_cache_spec.cr @@ -11,6 +11,13 @@ require "../../src/invidious/config" CONFIG.invidious_companion_key = "1234567890123456" +private def subtitle_response(status_code : Int32, content_type : String, body : String) + headers = HTTP::Headers.new + headers["Content-Type"] = content_type + headers["Access-Control-Allow-Origin"] = "*" + Invidious::SubtitleCache::Response.new(status_code, headers, body) +end + Spectator.describe Invidious::SubtitleCache do describe "check verification" do it "verifies valid tokens" do @@ -77,22 +84,22 @@ Spectator.describe Invidious::SubtitleCache do fetch_count = 0 vtt = "WEBVTT\n\n00:00:00.000 --> 00:00:01.000\nHello" - entry1, status1 = cache.get_or_fetch("video:abc") do + result1 = cache.get_or_fetch("video:abc") do fetch_count += 1 - {200, "text/vtt", vtt} + subtitle_response(200, "text/vtt", vtt) end - expect(status1).to eq("miss") - expect(entry1.try &.body).to eq(vtt) + expect(result1.cache_status).to eq("miss") + expect(result1.entry.try &.body).to eq(vtt) expect(fetch_count).to eq(1) - entry2, status2 = cache.get_or_fetch("video:abc") do + result2 = cache.get_or_fetch("video:abc") do fetch_count += 1 - {200, "text/vtt", vtt} + subtitle_response(200, "text/vtt", vtt) end - expect(status2).to eq("hit") - expect(entry2.try &.body).to eq(vtt) + expect(result2.cache_status).to eq("hit") + expect(result2.entry.try &.body).to eq(vtt) expect(fetch_count).to eq(1) end @@ -104,12 +111,12 @@ Spectator.describe Invidious::SubtitleCache do 5.times do spawn do - entry, status = cache.get_or_fetch("video:concurrent") do + result = cache.get_or_fetch("video:concurrent") do fetch_count += 1 sleep 0.05.seconds - {200, "text/vtt", vtt} + subtitle_response(200, "text/vtt", vtt) end - done_ch.send(status) + done_ch.send(result.cache_status) end end @@ -159,12 +166,28 @@ Spectator.describe Invidious::SubtitleCache do it "does not cache failed or non-vtt responses" do cache = Invidious::SubtitleCache.new - entry, status = cache.get_or_fetch("video:bad") do - {404, "application/json", "{\"error\":\"not found\"}"} + result = cache.get_or_fetch("video:bad") do + subtitle_response(404, "application/json", "{\"error\":\"not found\"}") end - expect(status).to eq("bypass") - expect(entry).to be_nil + expect(result.cache_status).to eq("bypass") + expect(result.entry).to be_nil + expect(result.response.try &.body).to eq("{\"error\":\"not found\"}") + expect(result.response.try &.headers["Access-Control-Allow-Origin"]).to eq("*") + expect(cache.size).to eq(0) + end + + it "does not cache oversized fetch responses" do + cache = Invidious::SubtitleCache.new + oversized = "WEBVTT\n" + ("x" * (Invidious::SubtitleCache::MAX_ENTRY_BYTES + 1)) + + result = cache.get_or_fetch("video:oversized") do + subtitle_response(200, "text/vtt", oversized) + end + + expect(result.cache_status).to eq("bypass") + expect(result.entry).to be_nil + expect(result.response.try &.body).to eq(oversized) expect(cache.size).to eq(0) end end diff --git a/src/invidious/helpers/subtitle_cache.cr b/src/invidious/helpers/subtitle_cache.cr index ecd167662..fa010e8ae 100644 --- a/src/invidious/helpers/subtitle_cache.cr +++ b/src/invidious/helpers/subtitle_cache.cr @@ -15,6 +15,24 @@ module Invidious end end + struct Response + getter status_code : Int32 + getter headers : HTTP::Headers + getter body : String + + def initialize(@status_code : Int32, @headers : HTTP::Headers, @body : String) + end + end + + struct FetchResult + getter entry : Entry? + getter response : Response? + getter cache_status : String + + def initialize(@entry : Entry?, @response : Response?, @cache_status : String) + end + end + DEFAULT_MAX_ENTRIES = 256 DEFAULT_MAX_BYTES = 64 * 1024 * 1024 # 64 MiB DEFAULT_TTL = 6.hours # 21600 seconds @@ -31,7 +49,7 @@ module Invidious @entries = Hash(String, Entry).new @total_bytes = 0 @mutex = Mutex.new - @in_flight = Hash(String, Array(::Channel(Entry?))).new + @in_flight = Hash(String, Array(::Channel(FetchResult))).new end def self.valid_vtt?(body : String) : Bool @@ -75,15 +93,15 @@ module Invidious end def put(key : String, body : String, content_type : String) : Entry? - return nil unless SubtitleCache.valid_vtt?(body) - return nil if body.bytesize > MAX_ENTRY_BYTES - @mutex.synchronize do put_internal(key, body, content_type) end end private def put_internal(key : String, body : String, content_type : String) : Entry? + return nil unless SubtitleCache.valid_vtt?(body) + return nil if body.bytesize > MAX_ENTRY_BYTES || body.bytesize > @max_bytes + if old = @entries.delete(key) @total_bytes -= old.size end @@ -104,48 +122,40 @@ module Invidious entry end - def get_or_fetch(key : String, &fetch_block : -> Tuple(Int32, String, String)?) : Tuple(Entry?, String) - wait_ch : ::Channel(Entry?)? = nil + def get_or_fetch(key : String, &fetch_block : -> Response?) : FetchResult + wait_ch : ::Channel(FetchResult)? = nil @mutex.synchronize do if entry = get_internal(key) - return {entry, "hit"} + return FetchResult.new(entry, nil, "hit") end if waiting_list = @in_flight[key]? - ch = ::Channel(Entry?).new(1) + ch = ::Channel(FetchResult).new(1) waiting_list << ch wait_ch = ch else - @in_flight[key] = Array(::Channel(Entry?)).new + @in_flight[key] = Array(::Channel(FetchResult)).new end end if ch = wait_ch - entry = ch.receive - if entry - return {entry, "hit"} - else - return {nil, "bypass"} - end + return ch.receive end # Primary fetcher for this key - fetch_result = begin + response = begin fetch_block.call rescue nil end cached_entry : Entry? = nil - waiting_channels = Array(::Channel(Entry?)).new + waiting_channels = Array(::Channel(FetchResult)).new @mutex.synchronize do - if fetch_result - status, content_type, body = fetch_result - if status == 200 && SubtitleCache.valid_vtt?(body) - cached_entry = put_internal(key, body, content_type) - end + if response && response.status_code == 200 + cached_entry = put_internal(key, response.body, response.headers["Content-Type"]? || "text/vtt; charset=utf-8") end if channels = @in_flight.delete(key) @@ -153,15 +163,14 @@ module Invidious end end + primary_result = FetchResult.new(cached_entry, response, cached_entry ? "miss" : "bypass") + waiter_result = FetchResult.new(cached_entry, response, cached_entry ? "hit" : "bypass") + waiting_channels.each do |w_ch| - w_ch.send(cached_entry) + w_ch.send(waiter_result) end - if cached_entry - {cached_entry, "miss"} - else - {nil, "bypass"} - end + primary_result end end end diff --git a/src/invidious/routes/companion.cr b/src/invidious/routes/companion.cr index d359a73ef..a99317824 100644 --- a/src/invidious/routes/companion.cr +++ b/src/invidious/routes/companion.cr @@ -82,36 +82,37 @@ module Invidious::Routes::Companion cache_key = "video:#{video_id}|label:#{label}|lang:#{lang}|tlang:#{tlang}" - entry, cache_status = SUBTITLE_CACHE.get_or_fetch(cache_key) do + result = SUBTITLE_CACHE.get_or_fetch(cache_key) do fetch_from_companion(url, env.request.headers) end - if entry + if entry = result.entry env.response.status_code = 200 + env.response.headers["Access-Control-Allow-Origin"] = "*" env.response.headers["Content-Type"] = entry.content_type env.response.headers["Cache-Control"] = "private, max-age=21600" - env.response.headers["X-Invidious-Subtitle-Cache"] = cache_status + env.response.headers["X-Invidious-Subtitle-Cache"] = result.cache_status env.response.print entry.body return - else - begin - COMPANION_POOL.client do |wrapper| - wrapper.client.get(url, env.request.headers) do |resp| - env.response.headers["X-Invidious-Subtitle-Cache"] = cache_status - return self.proxy_companion(env, resp) - end - end - rescue ex + elsif response = result.response + env.response.status_code = response.status_code + response.headers.each do |key, value| + env.response.headers[key] = value end + env.response.headers["X-Invidious-Subtitle-Cache"] = result.cache_status + env.response.print response.body end end - private def self.fetch_from_companion(url : String, headers : HTTP::Headers) : Tuple(Int32, String, String)? + 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 - content_type = resp.headers["Content-Type"]? || "text/vtt; charset=utf-8" - return {resp.status_code, content_type, body} + response_headers = HTTP::Headers.new + resp.headers.each do |key, value| + response_headers[key] = value + end + return Invidious::SubtitleCache::Response.new(resp.status_code, response_headers, body) end end rescue ex