VeritableValor/invidious
1
0
Fork 0
mirror of https://github.com/iv-org/invidious.git synced 2026-09-06 17:12:45 -05:00
Code Issues Packages Projects Releases Wiki Activity

perf(companion): cache valid caption responses

Browse Source
This commit is contained in:
XZH 2026-08-28 07:41:29 -07:00
parent 34173a3d84
commit ca2108e669
4 changed files with 444 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

171
spec/invidious/subtitle_cache_spec.cr Normal file
View File

@ -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?("<!DOCTYPE html><html></html>")).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

167
src/invidious/helpers/subtitle_cache.cr Normal file
View File

@ -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

37
src/invidious/helpers/utils.cr
View File

@ -404,3 +404,40 @@ def invidious_companion_encrypt(data)
encrypted_data = encrypt_ecb_without_salt("#{timestamp}|#{data}", CONFIG.invidious_companion_key) encrypted_data = encrypt_ecb_without_salt("#{timestamp}|#{data}", CONFIG.invidious_companion_key)
return Base64.urlsafe_encode(encrypted_data) return Base64.urlsafe_encode(encrypted_data)
end 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

69
src/invidious/routes/companion.cr
View File

@ -1,4 +1,6 @@
module Invidious::Routes::Companion module Invidious::Routes::Companion
SUBTITLE_CACHE = Invidious::SubtitleCache.new
# GET /companion # GET /companion
def self.get_companion(env) def self.get_companion(env)
url = env.request.path url = env.request.path
@ -6,6 +8,12 @@ module Invidious::Routes::Companion
url += "?#{env.request.query}" url += "?#{env.request.query}"
end 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 begin
COMPANION_POOL.client do |wrapper| COMPANION_POOL.client do |wrapper|
wrapper.client.get(url, env.request.headers) do |resp| wrapper.client.get(url, env.request.headers) do |resp|
@ -49,6 +57,67 @@ module Invidious::Routes::Companion
end end
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) private def self.proxy_companion(env, response)
env.response.status_code = response.status_code env.response.status_code = response.status_code
response.headers.each do |key, value| response.headers.each do |key, value|