webserver/src/main.rs

mod icon;
mod index;
// mod lua;

use std::{collections::HashMap, fs, io::ErrorKind, path::PathBuf, str::FromStr, sync::LazyLock};

use index::render_markdown;
use mlua::{
    Lua, Number, Table,
    Value::{self as LuaValue},
};

use axum::{
    extract::Request,
    handler::HandlerWithoutStateExt,
    http::{HeaderName, HeaderValue, StatusCode},
    response::{Html, IntoResponse},
    routing::get,
};
use serde::Serialize;
use time::OffsetDateTime;
use tokio::time::Instant;
use tower_http::services::ServeDir;

pub type AError = Box<dyn std::error::Error>;
pub type AResult<T> = Result<T, AError>;

pub static CWD: LazyLock<PathBuf> = LazyLock::new(|| std::env::current_dir().unwrap());

pub static LUA_CPATH: LazyLock<String> = LazyLock::new(|| {
    String::from_utf8_lossy(
        &std::process::Command::new("luarocks")
            .arg("path")
            .arg("--lr-cpath")
            .arg("--full")
            .output()
            .map_err::<std::vec::Vec<u8>, _>(|_| vec![])
            .map(|r| r.stdout)
            .unwrap_or_default(),
    )
    .trim()
    .to_string()
});

pub static LUA_PATH: LazyLock<String> = LazyLock::new(|| {
    String::from_utf8_lossy(
        &std::process::Command::new("luarocks")
            .arg("path")
            .arg("--lr-path")
            .arg("--full")
            .output()
            .map_err::<std::vec::Vec<u8>, _>(|_| vec![])
            .map(|r| r.stdout)
            .unwrap_or_default(),
    )
    .trim()
    .to_string()
});

pub static BASE_DIR: LazyLock<PathBuf> = LazyLock::new(|| {
    PathBuf::from_str(&std::env::var("BASE_DIR").unwrap_or(String::from("./public"))).unwrap()
});

#[derive(Debug, Serialize)]
struct FileInfo {
    filename: String,
    size: usize,
    format: String,
    icon: Option<Vec<u8>>,
    modified: OffsetDateTime,
    description: Option<String>,
}

#[tokio::main]
async fn main() -> AResult<()> {
    let listener = tokio::net::TcpListener::bind("[::]:3000").await.unwrap();
    println!("Server running on port 3000");
    Ok(axum::serve(listener, file_handler.into_make_service()).await?)
}

async fn file_handler(request: Request) -> impl IntoResponse {
    std::env::set_current_dir(&*CWD).unwrap();

    let base_dir = BASE_DIR.canonicalize().map_err(|e| {
        eprintln!("{:#?}", e);
        (StatusCode::NOT_FOUND, "404: Not Found".to_string())
    })?;

    let mut request_uri_path = request.uri().path().to_owned();

    if request_uri_path.ends_with("/") {
        let mut a = base_dir.clone();
        a.push(format!(".{}/{}", &request_uri_path, "init.lua"));
        if let Ok(b) = a.canonicalize() {
            if b.exists() {
                request_uri_path += "init.lua";
            }
        }
    }

    if request_uri_path.to_lowercase().ends_with(".md") {
        let mut path = base_dir.clone();

        let uri = urlencoding::decode(&request_uri_path[1..]).map_err(|e| {
            println!("{:?}", e);
            (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
        })?;

        path.push(&*uri);
        let full_path = fs::canonicalize(&path).map_err(|e| {
            println!("{:?}", e);
            (StatusCode::NOT_FOUND, "404: Not Found".to_string())
        })?;

        let full_base_path = fs::canonicalize(&base_dir).map_err(|e| {
            println!("{:?}", e);
            (StatusCode::NOT_FOUND, "404: Not Found".to_string())
        })?;

        if !full_path.starts_with(&full_base_path) {
            return Err((StatusCode::BAD_REQUEST, "400: Bad Request".to_string()));
        }

        let script = fs::read_to_string(&full_path).map_err(|e| {
            eprintln!("Lua Read Error: {:?}", e);
            (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
        })?;

        let filename = full_path.file_name().unwrap_or_default().to_string_lossy();

        if let Ok(md) = markdown::to_html_with_options(&script, &markdown::Options::gfm()) {
            return Ok(Html(render_markdown(&filename, &md).unwrap()).into_response());
        }
    }

    if !request_uri_path.ends_with(".lua") {
        return ServeDir::new(&base_dir)
            .fallback(get(handler))
            .try_call(request)
            .await
            .map(|e| e.into_response())
            .map_err(|_| (StatusCode::NOT_FOUND, "404: Not Found".to_string()));
    };

    let lua = unsafe { Lua::unsafe_new() };

    // Fix load paths
    lua.load(format!(
        "package.path = \"{}\"\npackage.cpath = \"{}\"",
        &*LUA_PATH, &*LUA_CPATH
    ))
    .exec()
    .map_err(|e| {
        eprintln!("Lua Error: {:#?}", e);
        render_lua_error(e)
    })?;

    let globals = lua.globals();

    let request_table = lua.create_table().map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        render_lua_error(e)
    })?;

    let headers_table = lua
        .create_table_from(
            request
                .headers()
                .into_iter()
                .map(|(k, v)| (k.to_string(), v.to_str().unwrap())),
        )
        .map_err(|e| {
            eprintln!("Lua Error: {:?}", e);
            render_lua_error(e)
        })?;

    let query = request.uri().query().unwrap_or("");

    let query: HashMap<String, String> = serde_qs::from_str(query)
        .map_err(|e| {
            println!("{:#?}", e);
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                "500: Unable to parse query string".to_string(),
            )
        })
        .unwrap_or(HashMap::new());

    let query_table = lua.create_table_from(query).map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        render_lua_error(e)
    })?;

    request_table.set("headers", headers_table).map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        render_lua_error(e)
    })?;

    request_table
        .set("uri", request.uri().to_string())
        .map_err(|e| {
            eprintln!("Lua Error: {:?}", e);
            render_lua_error(e)
        })?;

    request_table
        .set("method", request.method().to_string())
        .map_err(|e| {
            eprintln!("Lua Error: {:?}", e);
            render_lua_error(e)
        })?;

    request_table.set("query", query_table).map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        render_lua_error(e)
    })?;

    // Inject functions to change headers and such
    globals.set("request", request_table).map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        render_lua_error(e)
    })?;

    let response_table = lua.create_table().map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        render_lua_error(e)
    })?;

    response_table
        .set(
            "headers",
            lua.create_table().map_err(|e| {
                eprintln!("Lua Error: {:?}", e);
                render_lua_error(e)
            })?,
        )
        .map_err(|e| {
            eprintln!("Lua Error: {:?}", e);
            render_lua_error(e)
        })?;

    globals.set("response", response_table).map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        render_lua_error(e)
    })?;

    let mut path = base_dir.clone();

    let uri = urlencoding::decode(&request_uri_path[1..]).map_err(|e| {
        println!("{:?}", e);
        (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
    })?;

    path.push(&*uri);
    let full_path = fs::canonicalize(&path).map_err(|e| {
        println!("{:?}", e);
        (StatusCode::NOT_FOUND, "404: Not Found".to_string())
    })?;

    let full_base_path = fs::canonicalize(base_dir).map_err(|e| {
        println!("{:?}", e);
        (StatusCode::NOT_FOUND, "404: Not Found".to_string())
    })?;

    if !full_path.starts_with(&full_base_path) {
        return Err((StatusCode::BAD_REQUEST, "400: Bad Request".to_string()));
    }

    let script = fs::read_to_string(full_path).map_err(|e| {
        eprintln!("Lua Read Error: {:?}", e);
        (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
    })?;

    // ------------------------------------------------------------------------

    // TODO: FIX THIS, THIS IS DISGUSTING
    std::env::set_current_dir(full_base_path).unwrap();

    // Before executing user-code, make sure they can't screw with `global`
    lua.load(
        r#"
    setmetatable(_G, {
        __newindex = function (_, n)
            error("attempt to write to undeclared variable "..n, 2)
        end,
        __index = function (_, n)
            error("attempt to read undeclared variable "..n, 2)
        end,
    })"#,
    )
    .exec()
    .map_err(|e| {
        std::env::set_current_dir(&*CWD).unwrap();
        eprintln!("Lua Error: {:?}", e);
        render_lua_error(e)
    })?;

    let script = lua.load(script).eval::<LuaValue>().map_err(|e| {
        std::env::set_current_dir(&*CWD).unwrap();
        eprintln!("Lua Error: {:?}", e);
        render_lua_error(e)
    })?;
    std::env::set_current_dir(&*CWD).unwrap();

    let result = lua.load("return response").eval::<Table>().map_err(|e| {
        eprintln!("Lua Error: {:?}", e);
        render_lua_error(e)
    })?;

    // ------------------------------------------------------------------------

    let script = if let LuaValue::String(script) = script {
        script.to_string_lossy().to_string()
    } else {
        serde_json::to_string(&script).map_err(|e| {
            eprintln!("Serde Error: {:?}", e);
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Serde error: {}", e),
            )
        })?
    };

    let mut response = Html(script).into_response();

    if let Ok(status_code) = result.get::<Number>("code") {
        if let Ok(code) = StatusCode::from_u16(status_code as u16) {
            *response.status_mut() = code;
        }
    }

    if let Ok(headers) = result.get::<Table>("headers") {
        let pairs = headers.pairs::<String, String>();

        for pair in pairs {
            let (k, v) = pair.map_err(|e| {
                eprintln!("Lua Error: {:?}", e);
                render_lua_error(e)
            })?;
            response.headers_mut().insert(
                HeaderName::from_str(&k).map_err(|e| {
                    eprintln!("Response Error: {:?}", e);
                    (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
                })?,
                HeaderValue::from_str(&v).map_err(|e| {
                    eprintln!("Response Error: {:?}", e);
                    (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
                })?,
            );
        }
    };

    Ok(response)
}

#[allow(unused_variables)]
fn render_lua_error(e: mlua::Error) -> (StatusCode, String) {
    match e {
        mlua::Error::SyntaxError {
            message,
            incomplete_input,
        } => {
            let mut message = message.split("]:");
            message.next();
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                message.next().unwrap_or_default().to_string(),
            )
        }
        mlua::Error::RuntimeError(message) => (StatusCode::INTERNAL_SERVER_ERROR, message),
        mlua::Error::MemoryError(message) => (StatusCode::INTERNAL_SERVER_ERROR, message),
        mlua::Error::SafetyError(message) => (StatusCode::INTERNAL_SERVER_ERROR, message),
        mlua::Error::MemoryControlNotAvailable => (
            StatusCode::INTERNAL_SERVER_ERROR,
            "MemoryControlNotAvailable".to_string(),
        ),
        mlua::Error::RecursiveMutCallback => (
            StatusCode::INTERNAL_SERVER_ERROR,
            "RecursiveMutCallback".to_string(),
        ),
        mlua::Error::CallbackDestructed => (
            StatusCode::INTERNAL_SERVER_ERROR,
            "CallbackDestructed".to_string(),
        ),
        mlua::Error::StackError => (StatusCode::INTERNAL_SERVER_ERROR, "StackError".to_string()),
        mlua::Error::BindError => (StatusCode::INTERNAL_SERVER_ERROR, "BindError".to_string()),
        mlua::Error::BadArgument {
            to,
            pos,
            name,
            cause,
        } => (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!(
                "Bad Argument:\n\tto: {}\n\tpos: {}\n\tname: {}\n\tcause: {:?}",
                to.unwrap_or_default(),
                pos,
                name.unwrap_or_default(),
                cause
            ),
        ),
        mlua::Error::ToLuaConversionError { from, to, message } => (
            StatusCode::INTERNAL_SERVER_ERROR,
            message.unwrap_or_default(),
        ),
        mlua::Error::FromLuaConversionError { from, to, message } => (
            StatusCode::INTERNAL_SERVER_ERROR,
            message.unwrap_or_default(),
        ),
        mlua::Error::CoroutineUnresumable => (StatusCode::INTERNAL_SERVER_ERROR, String::new()),
        mlua::Error::UserDataTypeMismatch => (StatusCode::INTERNAL_SERVER_ERROR, String::new()),
        mlua::Error::UserDataDestructed => (StatusCode::INTERNAL_SERVER_ERROR, String::new()),
        mlua::Error::UserDataBorrowError => (StatusCode::INTERNAL_SERVER_ERROR, String::new()),
        mlua::Error::UserDataBorrowMutError => (StatusCode::INTERNAL_SERVER_ERROR, String::new()),
        mlua::Error::MetaMethodRestricted(message) => (StatusCode::INTERNAL_SERVER_ERROR, message),
        mlua::Error::MetaMethodTypeError {
            method,
            type_name,
            message,
        } => (
            StatusCode::INTERNAL_SERVER_ERROR,
            message.unwrap_or_default(),
        ),
        mlua::Error::MismatchedRegistryKey => (StatusCode::INTERNAL_SERVER_ERROR, String::new()),
        mlua::Error::CallbackError { traceback, cause } => (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("{}\n\n{:#?}", traceback, cause),
        ),
        mlua::Error::PreviouslyResumedPanic => (StatusCode::INTERNAL_SERVER_ERROR, String::new()),
        mlua::Error::SerializeError(message) => (StatusCode::INTERNAL_SERVER_ERROR, message),
        mlua::Error::DeserializeError(message) => (StatusCode::INTERNAL_SERVER_ERROR, message),
        mlua::Error::ExternalError(error) => {
            (StatusCode::INTERNAL_SERVER_ERROR, format!("{:#?}", error))
        }
        mlua::Error::WithContext { context, cause } => (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("{}\n\n{:#?}", context, cause),
        ),
        _ => (
            StatusCode::INTERNAL_SERVER_ERROR,
            String::from("Unknown Error"),
        ),
    }
}

async fn handler(request: Request) -> Result<Html<String>, (StatusCode, String)> {
    let start = Instant::now();
    let mut path = BASE_DIR.clone();

    let uri = urlencoding::decode(&request.uri().path()[1..]).map_err(|e| {
        println!("{:?}", e);
        (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
    })?;

    path.push(&*uri);
    let full_path = path.canonicalize().map_err(|e| {
        println!("{:?}", e);
        (StatusCode::NOT_FOUND, "404: Not Found".to_string())
    })?;

    let full_base_path = BASE_DIR.canonicalize().map_err(|e| {
        println!("{:?}", e);
        (StatusCode::NOT_FOUND, "404: Not Found".to_string())
    })?;

    if !full_path.starts_with(full_base_path) {
        return Err((StatusCode::BAD_REQUEST, "400: Bad Request".to_string()));
    }

    let body = match crate::index::render_index(&path, start).await {
        Ok(body) => Ok(body),
        Err(e) => {
            let ioerror: Option<&std::io::Error> = e.downcast_ref();
            if let Some(ioerror) = ioerror {
                if ioerror.kind() == ErrorKind::NotFound {
                    Err((StatusCode::NOT_FOUND, "404: Not Found".to_string()))
                } else {
                    println!("{:?}", ioerror);
                    Err((StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))
                }
            } else {
                println!("{:?}", e);
                Err((StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))
            }
        }
    };

    Ok(Html(body?))
}