beeta/openresty/lua/thumbnail.lua

local _M = {}

local THUMB_DIR = "/tmp/thumbs"
local GENERATION_INTERVAL = 600  -- 10 minutes in seconds
local NEW_STREAM_DELAY = 60  -- 1 minute delay for new streams
local NEW_STREAM_CHECK_INTERVAL = 10  -- Check for new streams every 10 seconds
local FFMPEG_TIMEOUT = 15  -- seconds (needs more time for animated capture)
local ANIMATION_DURATION = 3  -- seconds of video to capture
local ANIMATION_FPS = 8  -- frames per second in output

-- Track known live streams (in-memory, resets on nginx restart)
local known_streams = {}
-- Track streams with pending thumbnail generation
local pending_streams = {}

-- SECURITY FIX: Shell escape function to prevent command injection
local function shell_escape(str)
    -- Escape single quotes by ending the quoted string, adding escaped quote, and starting new quoted string
    return "'" .. string.gsub(str, "'", "'\\''") .. "'"
end

local function file_exists(path)
    local f = io.open(path, "r")
    if f then
        f:close()
        return true
    end
    return false
end

local function file_size(path)
    local f = io.open(path, "rb")
    if f then
        local size = f:seek("end")
        f:close()
        return size or 0
    end
    return 0
end

-- Lookup stream_key from realm name via internal API
local function get_stream_key_for_realm(realm_name)
    local http = require "resty.http"
    local httpc = http.new()

    local res, err = httpc:request_uri("http://drogon-backend:8080/internal/realm-stream-key/" .. ngx.escape_uri(realm_name), {
        method = "GET",
        headers = { ["Content-Type"] = "application/json" }
    })

    if not res then
        ngx.log(ngx.ERR, "Failed to lookup stream key for realm ", realm_name, ": ", err)
        return nil
    end

    if res.status ~= 200 then
        return nil
    end

    local cjson = require "cjson"
    local ok, data = pcall(cjson.decode, res.body)
    if not ok or not data.streamKey then
        return nil
    end

    return data.streamKey
end

-- Get all live realms from API
local function get_live_realms()
    local http = require "resty.http"
    local httpc = http.new()

    local res, err = httpc:request_uri("http://drogon-backend:8080/api/realms/live", {
        method = "GET",
        headers = { ["Content-Type"] = "application/json" }
    })

    if not res then
        ngx.log(ngx.ERR, "Failed to get live realms: ", err)
        return {}
    end

    if res.status ~= 200 then
        ngx.log(ngx.ERR, "Failed to get live realms: status ", res.status)
        return {}
    end

    local cjson = require "cjson"
    local ok, realms = pcall(cjson.decode, res.body)
    if not ok then
        ngx.log(ngx.ERR, "Failed to parse live realms response")
        return {}
    end

    return realms
end

-- Generate thumbnail for a single stream
local function generate_thumbnail(stream_key)
    -- Build LLHLS URL - internal docker network
    local llhls_url = "http://ovenmediaengine:8080/app/" .. stream_key .. "/llhls.m3u8"
    local thumb_path = THUMB_DIR .. "/" .. stream_key .. ".webp"
    local log_file = THUMB_DIR .. "/" .. stream_key .. ".log"

    -- First, try to generate animated WebP
    local cmd = string.format(
        "timeout %d ffmpeg -y -i %s -t %d -vf 'fps=%d,scale=320:-1:flags=lanczos' -c:v libwebp -lossless 0 -compression_level 3 -q:v 70 -loop 0 -preset default -an %s > %s 2>&1",
        FFMPEG_TIMEOUT,
        shell_escape(llhls_url),
        ANIMATION_DURATION,
        ANIMATION_FPS,
        shell_escape(thumb_path),
        shell_escape(log_file)
    )

    ngx.log(ngx.INFO, "Generating thumbnail for stream: ", stream_key)

    local exit_code = os.execute(cmd)

    -- Check if file was created and has content
    if file_exists(thumb_path) and file_size(thumb_path) > 100 then
        ngx.log(ngx.INFO, "Thumbnail generated successfully: ", thumb_path)
        os.remove(log_file)
        return true
    end

    -- If animated webp failed, try static webp (single frame)
    ngx.log(ngx.WARN, "Animated webp failed (exit: ", tostring(exit_code), "), trying static")

    local static_cmd = string.format(
        "timeout %d ffmpeg -y -i %s -vframes 1 -vf 'scale=320:-1:flags=lanczos' -c:v libwebp -q:v 75 %s > %s 2>&1",
        FFMPEG_TIMEOUT,
        shell_escape(llhls_url),
        shell_escape(thumb_path),
        shell_escape(log_file)
    )

    exit_code = os.execute(static_cmd)

    if file_exists(thumb_path) and file_size(thumb_path) > 100 then
        ngx.log(ngx.INFO, "Static thumbnail generated: ", thumb_path)
        os.remove(log_file)
        return true
    end

    ngx.log(ngx.ERR, "Thumbnail generation failed for stream: ", stream_key)
    return false
end

-- Generate thumbnail for a specific new stream (delayed callback)
local function generate_new_stream_thumbnail(premature, realm_name, stream_key)
    if premature then
        return
    end

    -- Clear pending flag
    pending_streams[realm_name] = nil

    ngx.log(ngx.INFO, "Generating initial thumbnail for new stream: ", realm_name)

    -- Ensure thumb directory exists
    os.execute("mkdir -p " .. THUMB_DIR)

    -- Generate the thumbnail
    if generate_thumbnail(stream_key) then
        ngx.log(ngx.INFO, "Initial thumbnail generated for new stream: ", realm_name)
    else
        ngx.log(ngx.WARN, "Failed to generate initial thumbnail for new stream: ", realm_name)
    end
end

-- Check for new streams and schedule thumbnail generation
local function check_for_new_streams(premature)
    if premature then
        return
    end

    local realms = get_live_realms()
    local current_streams = {}

    for _, realm in ipairs(realms) do
        if realm.name then
            current_streams[realm.name] = true

            -- Check if this is a new stream we haven't seen
            if not known_streams[realm.name] and not pending_streams[realm.name] then
                ngx.log(ngx.INFO, "New stream detected: ", realm.name, " - scheduling thumbnail in ", NEW_STREAM_DELAY, " seconds")

                local stream_key = get_stream_key_for_realm(realm.name)
                if stream_key then
                    -- Mark as pending to avoid duplicate scheduling
                    pending_streams[realm.name] = true

                    -- Schedule thumbnail generation after delay
                    local ok, err = ngx.timer.at(NEW_STREAM_DELAY, generate_new_stream_thumbnail, realm.name, stream_key)
                    if not ok then
                        ngx.log(ngx.ERR, "Failed to schedule new stream thumbnail: ", err)
                        pending_streams[realm.name] = nil
                    end
                end
            end

            -- Mark stream as known
            known_streams[realm.name] = true
        end
    end

    -- Clean up streams that are no longer live
    for name, _ in pairs(known_streams) do
        if not current_streams[name] then
            known_streams[name] = nil
            pending_streams[name] = nil
        end
    end

    -- Schedule next check
    local ok, err = ngx.timer.at(NEW_STREAM_CHECK_INTERVAL, check_for_new_streams)
    if not ok then
        ngx.log(ngx.ERR, "Failed to schedule next new stream check: ", err)
    end
end

-- Background job: Generate thumbnails for all live streams
local function generate_all_thumbnails(premature)
    if premature then
        return
    end

    ngx.log(ngx.INFO, "Starting thumbnail generation for all live streams")

    -- Ensure thumb directory exists
    os.execute("mkdir -p " .. THUMB_DIR)

    local realms = get_live_realms()
    local count = 0

    for _, realm in ipairs(realms) do
        if realm.name then
            local stream_key = get_stream_key_for_realm(realm.name)
            if stream_key then
                generate_thumbnail(stream_key)
                count = count + 1
            end
        end
    end

    ngx.log(ngx.INFO, "Thumbnail generation complete. Generated ", count, " thumbnails")

    -- Schedule next run
    local ok, err = ngx.timer.at(GENERATION_INTERVAL, generate_all_thumbnails)
    if not ok then
        ngx.log(ngx.ERR, "Failed to schedule next thumbnail generation: ", err)
    end
end

-- Initialize background thumbnail generation (call from init_worker_by_lua_block)
function _M.init_worker()
    -- Only run on worker 0 to avoid duplicate work
    if ngx.worker.id() ~= 0 then
        return
    end

    ngx.log(ngx.INFO, "Initializing thumbnail generator on worker 0")

    -- Ensure thumb directory exists
    os.execute("mkdir -p " .. THUMB_DIR)

    -- Start first full generation after 10 seconds (let services start up)
    local ok, err = ngx.timer.at(10, generate_all_thumbnails)
    if not ok then
        ngx.log(ngx.ERR, "Failed to start thumbnail generator: ", err)
    end

    -- Start new stream detection after 15 seconds
    ok, err = ngx.timer.at(15, check_for_new_streams)
    if not ok then
        ngx.log(ngx.ERR, "Failed to start new stream detector: ", err)
    end
end

-- Serve existing thumbnail (no generation on request)
function _M.serve()
    -- Get realm name from nginx variable
    local realm_name = ngx.var.realm_name
    if not realm_name or realm_name == "" then
        ngx.status = 400
        ngx.header["Content-Type"] = "text/plain"
        ngx.say("Missing realm name")
        return ngx.exit(400)
    end

    -- Sanitize realm name (alphanumeric, dash only)
    if not realm_name:match("^[%w%-]+$") then
        ngx.status = 400
        ngx.header["Content-Type"] = "text/plain"
        ngx.say("Invalid realm name")
        return ngx.exit(400)
    end

    -- Lookup stream key from realm name
    local stream_key = get_stream_key_for_realm(realm_name)
    if not stream_key then
        ngx.status = 404
        ngx.header["Content-Type"] = "text/plain"
        ngx.say("Realm not found or not live")
        return ngx.exit(404)
    end

    local thumb_path = THUMB_DIR .. "/" .. stream_key .. ".webp"

    -- Check if thumbnail exists
    if not file_exists(thumb_path) or file_size(thumb_path) < 100 then
        ngx.status = 404
        ngx.header["Content-Type"] = "text/plain"
        ngx.say("Thumbnail not available")
        return ngx.exit(404)
    end

    -- Serve the thumbnail file
    ngx.header["Content-Type"] = "image/webp"
    ngx.header["Cache-Control"] = "public, max-age=60"  -- Cache for 1 minute on client
    ngx.header["Access-Control-Allow-Origin"] = "*"

    local f = io.open(thumb_path, "rb")
    if not f then
        ngx.status = 500
        ngx.header["Content-Type"] = "text/plain"
        ngx.say("Failed to read thumbnail")
        return ngx.exit(500)
    end

    local content = f:read("*a")
    f:close()

    ngx.print(content)
end

return _M