mirror of
https://github.com/iv-org/invidious.git
synced 2026-09-06 00:52:45 -05:00
perf(companion): cache valid caption responses
This commit is contained in:
XZH
parent
34173a3d84
commit
ca2108e669
171
spec/invidious/subtitle_cache_spec.cr
Normal file
171
spec/invidious/subtitle_cache_spec.cr
Normal 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
167
src/invidious/helpers/subtitle_cache.cr
Normal 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
|
||||
@ -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
|
||||
|
||||
@ -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|
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user