transotf/src/main.rs

#![allow(dead_code)]

use crate::av::init;
use crate::transcoder_manager::{TranscoderInstance, TranscoderManager};
use askama::Template;
use async_std::fs::File;
use async_std::sync::{Arc, RwLock};
use ffmpeg_sys_next::{av_make_error_string, av_malloc, av_version_info};
use http_types::Error as HttpError;
use lazy_static::lazy_static;
use serde::Serialize;
use std::collections::HashMap;
use std::ffi::{CStr, CString};
use std::ops::Deref;
use std::option::Option::Some;
use std::os::raw::{c_char, c_int};
use std::str::FromStr;
use tide::{Body, Request, Response, StatusCode};
use uuid::Uuid;

pub mod av;
pub mod transcoder;
pub mod transcoder_manager;
pub mod utils;

const OUTPUT_PATH: &str = "/tmp/transotf";

fn av_err2str(error_code: c_int) -> String {
    unsafe {
        let str: *mut c_char = av_malloc(1024).cast();
        av_make_error_string(str, 1024, error_code);
        CString::from_raw(str).to_str().unwrap().to_string()
    }
}

lazy_static! {
    static ref HOME_PAGE: String = {
        format!(
            "{}: {}\nffmpeg: {}",
            env!("CARGO_PKG_NAME"),
            env!("CARGO_PKG_VERSION"),
            unsafe { CStr::from_ptr(av_version_info()).to_str().unwrap_or("n/a") },
        )
    };
    static ref FAVICON: Vec<u8> = include_bytes!("../resources/favicon.ico").to_vec();
}

#[derive(Default, Clone)]
struct State {
    manager: Arc<RwLock<TranscoderManager>>,
}

#[derive(Serialize)]
struct CreateDTO {
    id: String,
    manifest: String,
    _player: String,
}

fn build_url(req: &Request<State>, url: String) -> String {
    let host = req.header("host").and_then(|x| x.get(0)).unwrap().clone();
    format!("http://{}/{}", host, url.trim_start_matches("/"))
}

async fn create_transcode(req: Request<State>) -> Result<Response, HttpError> {
    let params = req.query::<HashMap<String, String>>()?;
    let target = params.get(&"target".to_string());
    let target = if let Some(target) = target {
        target
    } else {
        return Err(HttpError::from_str(
            StatusCode::BadRequest,
            "no target given",
        ));
    };

    {
        File::open(target).await.map_err(|err| {
            HttpError::from_str(
                StatusCode::BadRequest,
                format!("Failed to open file: {}", err),
            )
        })?;
    }

    let id = req
        .state()
        .manager
        .write()
        .await
        .start(target.clone(), None, None)
        .map_err(|err| {
            HttpError::from_str(
                StatusCode::InternalServerError,
                format!("Failed to start transcoder: {}", err),
            )
        })?;

    let id = id.to_string();
    let dto = CreateDTO {
        manifest: build_url(&req, format!("/session/{}/manifest.mpd", id)),
        _player: build_url(&req, format!("/session/{}/player", id)),
        id,
    };

    let mut resp = Response::new(200);
    resp.set_body(Body::from_json(&dto)?);
    resp.insert_header("Content-Type", "application/json");

    Ok(resp)
}

async fn get_instance(req: &Request<State>) -> Result<Arc<TranscoderInstance>, HttpError> {
    let id = req.param("id")?;
    let id = Uuid::from_str(&id)?;
    let manager = req.state().manager.read().await;
    if let Some(instance) = manager.get(id) {
        Ok(instance)
    } else {
        Err(HttpError::from_str(
            StatusCode::NotFound,
            format!("Can't find session with id {}", id),
        ))
    }
}

#[derive(Template)]
#[template(path = "manifest.xml")]
struct MPDManifest {
    id: Uuid,
    has_audio: bool,
    duration: f64,
    is_vlc: bool,
    audio_time_base_den: i32,
    video_time_base_den: i32,
}

mod filters {
    pub fn iso_duration(secs: &f64) -> ::askama::Result<String> {
        let minutes = (secs / 60.0).floor();
        let secs = ((secs % 60.0) * 100.0).floor() / 100.0;
        let hours = (minutes / 60.0).floor();
        let minutes = minutes % 60.0;
        return Ok(format!("PT{}H{}M{}S", hours, minutes, secs));
    }
}

async fn get_manifest(req: Request<State>) -> Result<Response, HttpError> {
    let transcoder_instance = get_instance(&req).await?;
    let id = transcoder_instance.id();
    let status = transcoder_instance.status().await;

    let mut resp = Response::new(200);
    resp.insert_header("Content-Type", "application/dash+xml");

    let mpd = MPDManifest {
        id,
        has_audio: status.audio.is_some(),
        duration: status.duration_secs,
        is_vlc: req
            .header("User-Agent")
            .map(|x| x.last().to_string())
            .map_or(false, |x| x.starts_with("VLC")),
        audio_time_base_den: status.audio.map_or(0, |audio| audio.time_scale),
        video_time_base_den: status.video.time_scale,
    };

    resp.set_body(mpd.render().unwrap());

    Ok(resp)
}

async fn get_init(req: Request<State>) -> Result<Response, HttpError> {
    let transcoder_instance = get_instance(&req).await?;
    let type_ = req.param("type").unwrap();
    let id = transcoder_instance.id();

    if (type_ != "audio" && type_ != "video") || !transcoder_instance.wait_for_init().await {
        return Ok(Response::new(StatusCode::NotFound));
    }

    let mut ok = Response::new(StatusCode::Ok);
    ok.insert_header("Content-Type", "video/mp4");
    ok.set_body(Body::from_file(format!("{}/{}/{}-init.mp4", OUTPUT_PATH, id, type_)).await?);

    Ok(ok)
}

async fn get_segment(req: Request<State>) -> Result<Response, HttpError> {
    let transcoder_instance = get_instance(&req).await?;
    let type_ = req.param("type").unwrap();
    let segment = req.param("nr").unwrap();
    let id = transcoder_instance.id();

    if !segment.ends_with(".m4s")
        || (type_ != "audio" && type_ != "video")
        || !transcoder_instance.wait_for_init().await
    {
        return Ok(Response::new(StatusCode::NotFound));
    }

    let segment = &segment[0..segment.len() - 4];
    let segment = if let Ok(segment) = u32::from_str(segment) {
        segment
    } else {
        return Ok(Response::new(StatusCode::NotFound));
    };

    let status = transcoder_instance.status().await;
    if (segment as i64 - 5) > status.current_segment as i64 || segment < status.current_segment {
        if !status.segments.contains(&segment) {
            transcoder_instance.seek(segment).await;
        }
    }

    if !transcoder_instance.wait_for_segment(segment).await {
        return Ok(Response::new(StatusCode::NotFound));
    }

    let mut ok = Response::new(StatusCode::Ok);
    ok.insert_header("Content-Type", "video/mp4");
    ok.set_body(
        Body::from_file(format!(
            "{}/{}/{}-segment-{:0>5}.m4s",
            OUTPUT_PATH, id, type_, segment
        ))
        .await?,
    );

    Ok(ok)
}

async fn get_player(req: Request<State>) -> Result<Response, HttpError> {
    get_instance(&req).await?;
    let mut resp = Response::new(StatusCode::Ok);
    resp.insert_header("Content-Type", "text/html");
    resp.set_body(include_str!("../resources/player.html"));

    Ok(resp)
}

async fn get_favicon(_: Request<State>) -> Result<Response, HttpError> {
    let mut resp = Response::new(StatusCode::Ok);
    resp.insert_header("Content-Type", "image/x-icon");
    resp.set_body(Vec::deref(&FAVICON));

    Ok(resp)
}

async fn get_status(req: Request<State>) -> Result<Response, HttpError> {
    let instance = get_instance(&req).await?;
    let mut resp = Response::new(StatusCode::Ok);
    resp.insert_header("Content-Type", "application/json");
    resp.set_body(Body::from_json(&instance.status().await)?);

    Ok(resp)
}

#[async_std::main]
async fn main() -> std::io::Result<()> {
    init();

    let mut app = tide::with_state(State::default());
    app.at("/").get(|_| async { Ok(HOME_PAGE.clone()) });
    app.at("/favicon.ico").get(get_favicon);
    app.at("/transcode").get(create_transcode);
    app.at("/session/:id/status.json").get(get_status);
    app.at("/session/:id/player").get(get_player);
    app.at("/session/:id/manifest").get(get_manifest);
    app.at("/session/:id/manifest.xml").get(get_manifest);
    app.at("/session/:id/manifest.mpd").get(get_manifest);
    app.at("/session/:id/:type/init.mp4").get(get_init);
    app.at("/session/:id/:type/:nr").get(get_segment);
    let listen_fut = app.listen("0:8000");
    println!("Listening on 0:8000");
    listen_fut.await
}