diff --git a/spec/invidious/subtitle_cache_spec.cr b/spec/invidious/subtitle_cache_spec.cr new file mode 100644 index 000000000..a365a9b1c --- /dev/null +++ b/spec/invidious/subtitle_cache_spec.cr @@ -0,0 +1,171 @@ +require "../spec_helper" +require "../../src/invidious/user/preferences" +require "../../src/invidious/jobs/base_job" +require "../../src/invidious/jobs/*" +require "../../src/invidious/jobs" +require "../../src/invidious/config" + +{% unless @top_level.has_constant?(:CONFIG) %} + CONFIG = Config.from_yaml(File.open("config/config.example.yml")) +{% end %} + +CONFIG.invidious_companion_key = "1234567890123456" + +Spectator.describe Invidious::SubtitleCache do + describe "check verification" do + it "verifies valid tokens" do + token = invidious_companion_encrypt("test_video_123") + expect(invidious_companion_verify_check(token, "test_video_123")).to be_true + end + + it "rejects tokens with mismatched video id" do + token = invidious_companion_encrypt("test_video_123") + expect(invidious_companion_verify_check(token, "other_video_456")).to be_false + end + + it "rejects expired tokens older than 6 hours" do + old_ts = Time.utc.to_unix - (7 * 3600) + encrypted = encrypt_ecb_without_salt("#{old_ts}|test_video_123", CONFIG.invidious_companion_key) + token = Base64.urlsafe_encode(encrypted) + expect(invidious_companion_verify_check(token, "test_video_123")).to be_false + end + + it "rejects malformed or empty tokens" do + expect(invidious_companion_verify_check("", "test_video_123")).to be_false + expect(invidious_companion_verify_check("invalid-base64-!@#$", "test_video_123")).to be_false + expect(invidious_companion_verify_check("AAAA", "test_video_123")).to be_false + end + end + + describe "SubtitleCache operations" do + it "correctly identifies valid WebVTT" 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?("")).to be_false + expect(Invidious::SubtitleCache.valid_vtt?("{\"error\": \"not found\"}")).to be_false + expect(Invidious::SubtitleCache.valid_vtt?("")).to be_false + end + + it "stores and retrieves entries" do + cache = Invidious::SubtitleCache.new + vtt = "WEBVTT\n\n00:00:00.000 --> 00:00:01.000\nTest cue\n" + + entry = cache.put("video:123|label:en", vtt, "text/vtt; charset=utf-8") + expect(entry).not_to be_nil + expect(cache.size).to eq(1) + + retrieved = cache.get("video:123|label:en") + expect(retrieved).not_to be_nil + expect(retrieved.try &.body).to eq(vtt) + end + + it "isolates different languages and labels" do + cache = Invidious::SubtitleCache.new + vtt_en = "WEBVTT\n\n00:00:00.000 --> 00:00:01.000\nEnglish\n" + vtt_zh = "WEBVTT\n\n00:00:00.000 --> 00:00:01.000\nChinese\n" + + cache.put("video:123|label:English|lang:en|tlang:", vtt_en, "text/vtt") + cache.put("video:123|label:Chinese|lang:zh|tlang:", vtt_zh, "text/vtt") + + expect(cache.get("video:123|label:English|lang:en|tlang:").try &.body).to eq(vtt_en) + expect(cache.get("video:123|label:Chinese|lang:zh|tlang:").try &.body).to eq(vtt_zh) + end + + it "returns miss on first fetch and hit on subsequent fetch" do + cache = Invidious::SubtitleCache.new + fetch_count = 0 + vtt = "WEBVTT\n\n00:00:00.000 --> 00:00:01.000\nHello" + + entry1, status1 = cache.get_or_fetch("video:abc") do + fetch_count += 1 + {200, "text/vtt", vtt} + end + + expect(status1).to eq("miss") + expect(entry1.try &.body).to eq(vtt) + expect(fetch_count).to eq(1) + + entry2, status2 = cache.get_or_fetch("video:abc") do + fetch_count += 1 + {200, "text/vtt", vtt} + end + + expect(status2).to eq("hit") + expect(entry2.try &.body).to eq(vtt) + expect(fetch_count).to eq(1) + end + + it "coalesces concurrent fetches for the same key" do + cache = Invidious::SubtitleCache.new + fetch_count = 0 + vtt = "WEBVTT\n\n00:00:00.000 --> 00:00:01.000\nCoalesced" + done_ch = ::Channel(String).new(5) + + 5.times do + spawn do + entry, status = cache.get_or_fetch("video:concurrent") do + fetch_count += 1 + sleep 0.05.seconds + {200, "text/vtt", vtt} + end + done_ch.send(status) + end + end + + results = Array(String).new + 5.times { results << done_ch.receive } + + expect(fetch_count).to eq(1) + expect(results.count("miss")).to eq(1) + expect(results.count("hit")).to eq(4) + end + + it "evicts oldest entries when max_entries is exceeded (LRU)" do + cache = Invidious::SubtitleCache.new(max_entries: 2) + vtt = "WEBVTT\n\n00:00:00.000 --> 00:00:01.000\nLine" + + cache.put("k1", vtt, "text/vtt") + cache.put("k2", vtt, "text/vtt") + expect(cache.size).to eq(2) + + # Access k1 to make it most recently used + cache.get("k1") + + # Adding k3 should evict k2 (the oldest) + cache.put("k3", vtt, "text/vtt") + expect(cache.size).to eq(2) + expect(cache.get("k1")).not_to be_nil + expect(cache.get("k3")).not_to be_nil + expect(cache.get("k2")).to be_nil + end + + it "evicts oldest entries when max_bytes is exceeded" do + cache = Invidious::SubtitleCache.new(max_entries: 10, max_bytes: 60) + vtt1 = "WEBVTT\n\n12345678901234567890" # ~30 bytes + vtt2 = "WEBVTT\n\nABCDEFGHIJABCDEFGHIJ" # ~30 bytes + vtt3 = "WEBVTT\n\nXYZXYZXYZXYZXYZXYZ" # ~27 bytes + + cache.put("k1", vtt1, "text/vtt") + cache.put("k2", vtt2, "text/vtt") + expect(cache.size).to eq(2) + + # Adding k3 pushes total bytes over 60, evicts k1 + cache.put("k3", vtt3, "text/vtt") + expect(cache.get("k1")).to be_nil + expect(cache.get("k2")).not_to be_nil + expect(cache.get("k3")).not_to be_nil + end + + 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\"}"} + end + + expect(status).to eq("bypass") + expect(entry).to be_nil + expect(cache.size).to eq(0) + end + end +end diff --git a/src/invidious/helpers/subtitle_cache.cr b/src/invidious/helpers/subtitle_cache.cr new file mode 100644 index 000000000..ecd167662 --- /dev/null +++ b/src/invidious/helpers/subtitle_cache.cr @@ -0,0 +1,167 @@ +module Invidious + class SubtitleCache + struct Entry + getter body : String + getter content_type : String + getter created_at : Time + getter size : Int32 + + def initialize(@body : String, @content_type : String, @created_at : Time = Time.utc) + @size = @body.bytesize + end + + def expired?(ttl : Time::Span) : Bool + Time.utc - @created_at > ttl + end + end + + DEFAULT_MAX_ENTRIES = 256 + DEFAULT_MAX_BYTES = 64 * 1024 * 1024 # 64 MiB + DEFAULT_TTL = 6.hours # 21600 seconds + MAX_ENTRY_BYTES = 2 * 1024 * 1024 # 2 MiB + + getter max_entries : Int32 + getter max_bytes : Int32 + getter ttl : Time::Span + getter total_bytes : Int32 + + def initialize(@max_entries : Int32 = DEFAULT_MAX_ENTRIES, + @max_bytes : Int32 = DEFAULT_MAX_BYTES, + @ttl : Time::Span = DEFAULT_TTL) + @entries = Hash(String, Entry).new + @total_bytes = 0 + @mutex = Mutex.new + @in_flight = Hash(String, Array(::Channel(Entry?))).new + end + + def self.valid_vtt?(body : String) : Bool + return false if body.empty? + trimmed = body.lstrip("\uFEFF \t\r\n") + trimmed.starts_with?("WEBVTT") + end + + def size : Int32 + @mutex.synchronize { @entries.size } + end + + def clear : Nil + @mutex.synchronize do + @entries.clear + @total_bytes = 0 + end + end + + def get(key : String) : Entry? + @mutex.synchronize do + get_internal(key) + end + end + + private def get_internal(key : String) : Entry? + if entry = @entries[key]? + if entry.expired?(@ttl) + @entries.delete(key) + @total_bytes -= entry.size + nil + else + # Move to end (most recently used in Crystal Hash) + @entries.delete(key) + @entries[key] = entry + entry + end + else + nil + end + 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? + if old = @entries.delete(key) + @total_bytes -= old.size + end + + body_size = body.bytesize + + # Evict LRU entries if capacity exceeded + while (@entries.size >= @max_entries || @total_bytes + body_size > @max_bytes) && !@entries.empty? + oldest_key = @entries.first_key + if removed = @entries.delete(oldest_key) + @total_bytes -= removed.size + end + end + + entry = Entry.new(body, content_type, Time.utc) + @entries[key] = entry + @total_bytes += entry.size + entry + end + + def get_or_fetch(key : String, &fetch_block : -> Tuple(Int32, String, String)?) : Tuple(Entry?, String) + wait_ch : ::Channel(Entry?)? = nil + + @mutex.synchronize do + if entry = get_internal(key) + return {entry, "hit"} + end + + if waiting_list = @in_flight[key]? + ch = ::Channel(Entry?).new(1) + waiting_list << ch + wait_ch = ch + else + @in_flight[key] = Array(::Channel(Entry?)).new + end + end + + if ch = wait_ch + entry = ch.receive + if entry + return {entry, "hit"} + else + return {nil, "bypass"} + end + end + + # Primary fetcher for this key + fetch_result = begin + fetch_block.call + rescue + nil + end + + cached_entry : Entry? = nil + waiting_channels = Array(::Channel(Entry?)).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 + end + + if channels = @in_flight.delete(key) + waiting_channels = channels + end + end + + waiting_channels.each do |w_ch| + w_ch.send(cached_entry) + end + + if cached_entry + {cached_entry, "miss"} + else + {nil, "bypass"} + end + end + end +end diff --git a/src/invidious/helpers/utils.cr b/src/invidious/helpers/utils.cr index 7a262f849..013df0866 100644 --- a/src/invidious/helpers/utils.cr +++ b/src/invidious/helpers/utils.cr @@ -404,3 +404,40 @@ def invidious_companion_encrypt(data) encrypted_data = encrypt_ecb_without_salt("#{timestamp}|#{data}", CONFIG.invidious_companion_key) return Base64.urlsafe_encode(encrypted_data) end + +def decrypt_ecb_without_salt(data : Bytes | IO, key : String) : IO::Memory + cipher = OpenSSL::Cipher.new("aes-128-ecb") + cipher.decrypt + cipher.key = key + + io = IO::Memory.new + io.write(cipher.update(data)) + io.write(cipher.final) + io.rewind + + return io +end + +def invidious_companion_verify_check(check : String, expected_video_id : String, max_age_seconds : Int64 = 21600_i64) : Bool + return false if check.empty? || CONFIG.invidious_companion_key.size != 16 + + raw_bytes = Base64.decode(check) rescue nil + return false unless raw_bytes + + decrypted_io = decrypt_ecb_without_salt(raw_bytes, CONFIG.invidious_companion_key) rescue nil + return false unless decrypted_io + + plain = decrypted_io.to_s + parts = plain.split('|', 2) + return false unless parts.size == 2 + + timestamp = parts[0].to_i64? + return false unless timestamp + + video_id = parts[1] + return false unless video_id == expected_video_id + + now = Time.utc.to_unix + diff = (now - timestamp).abs + return diff <= max_age_seconds +end diff --git a/src/invidious/routes/companion.cr b/src/invidious/routes/companion.cr index 811393aba..d359a73ef 100644 --- a/src/invidious/routes/companion.cr +++ b/src/invidious/routes/companion.cr @@ -1,4 +1,6 @@ module Invidious::Routes::Companion + SUBTITLE_CACHE = Invidious::SubtitleCache.new + # GET /companion def self.get_companion(env) url = env.request.path @@ -6,6 +8,12 @@ module Invidious::Routes::Companion url += "?#{env.request.query}" end + path = env.request.path.rstrip('/') + if match = path.match(%r{^/(?:companion/)?api/v1/captions/([^/?#]+)}) + video_id = match[1] + return self.handle_caption_request(env, url, video_id) + end + begin COMPANION_POOL.client do |wrapper| wrapper.client.get(url, env.request.headers) do |resp| @@ -49,6 +57,67 @@ module Invidious::Routes::Companion end end + private def self.handle_caption_request(env, url, video_id) + label = env.params.query["label"]? || "" + lang = env.params.query["lang"]? || "" + tlang = env.params.query["tlang"]? || "" + check = env.params.query["check"]? || "" + + # Verify check token before accessing cache + valid_check = invidious_companion_verify_check(check, video_id) + + unless valid_check + # Bypass cache if token is invalid or expired + begin + COMPANION_POOL.client do |wrapper| + wrapper.client.get(url, env.request.headers) do |resp| + env.response.headers["X-Invidious-Subtitle-Cache"] = "bypass" + return self.proxy_companion(env, resp) + end + end + rescue ex + end + return + end + + cache_key = "video:#{video_id}|label:#{label}|lang:#{lang}|tlang:#{tlang}" + + entry, cache_status = SUBTITLE_CACHE.get_or_fetch(cache_key) do + fetch_from_companion(url, env.request.headers) + end + + if entry + env.response.status_code = 200 + 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.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 + end + end + end + + private def self.fetch_from_companion(url : String, headers : HTTP::Headers) : Tuple(Int32, String, String)? + 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} + end + end + rescue ex + nil + end + private def self.proxy_companion(env, response) env.response.status_code = response.status_code response.headers.each do |key, value|