PagesServer/src/main.rs

use std::io::Read;
use std::path::{Path, PathBuf};

use actix_files::{Files, NamedFile};
use actix_multipart::Multipart;
use actix_web::dev::{ServiceRequest, ServiceResponse, fn_service};
use actix_web::error::{ErrorBadRequest, ErrorInternalServerError, ErrorUnauthorized};
use actix_web::middleware::Logger;
use actix_web::{App, Error, HttpRequest, HttpResponse, HttpServer, web};
use bytes::BufMut;
use clap::Parser;
use futures_util::TryStreamExt;

/// Simple pages server
#[derive(Parser, Debug, Clone)]
#[clap(author, version, about, long_about = None)]
struct Args {
    /// The address and port this server will listen on
    #[clap(short, long, env, default_value = "0.0.0.0:8000")]
    listen_address: String,

    /// The place wheres files to serve are hosted
    #[clap(short, long, env)]
    files_path: String,

    /// Update token
    #[clap(short, long, env)]
    update_token: String,

    /// The IP addresses allowed to perform updates
    #[clap(short, long, env, default_value = "*")]
    allowed_ips_for_update: String,

    /// Index file name
    #[clap(short, long, env, default_value = "index.html")]
    index_file: String,

    /// Handle for not found files. By default, a basic message is returned
    #[clap(short, long, env, default_value = "")]
    not_found_file: String,

    /// Optional proxy IP
    #[clap(short, long, env)]
    proxy_ip: Option<String>,

    /// Specify whether HTML extensions can be bypassed
    #[clap(short, long, env)]
    can_bypass_html_ext: bool,
}

impl Args {
    pub fn storage_path(&self) -> &Path {
        Path::new(&self.files_path)
    }

    pub fn default_handler_path(&self) -> Option<PathBuf> {
        match self.not_found_file.is_empty() {
            true => None,
            false => Some(self.storage_path().join(&self.not_found_file)),
        }
    }
}

struct NewFile {
    path: PathBuf,
    bytes: Vec<u8>,
}

// Check if two ips matches
pub fn match_ip(pattern: &str, ip: &str) -> bool {
    if pattern.eq(ip) {
        return true;
    }

    if pattern.ends_with('*') && ip.starts_with(&pattern.replace('*', "")) {
        return true;
    }

    false
}

/// Get the remote IP address
fn get_remote_ip(req: &HttpRequest, args: &Args) -> String {
    let mut ip = req.peer_addr().unwrap().ip().to_string();

    // We check if the request comes from a trusted reverse proxy
    if let Some(proxy) = args.proxy_ip.as_ref() {
        if match_ip(proxy, &ip) {
            if let Some(header) = req.headers().get("X-Forwarded-For") {
                let header: Vec<String> = header
                    .to_str()
                    .unwrap()
                    .split(',')
                    .map(|f| f.to_string())
                    .collect();

                if !header.is_empty() {
                    ip = header[0].to_string();
                }
            }
        }
    }

    ip
}

/// Replace all the files of the website
async fn replace_files(
    args: web::Data<Args>,
    req: HttpRequest,
    mut payload: Multipart,
) -> Result<HttpResponse, Error> {
    // Validate remote IP
    let remote_ip = get_remote_ip(&req, &args);
    if !match_ip(&args.allowed_ips_for_update, &remote_ip) {
        log::warn!(
            "Block unauthorized attempt to perform site update from {}",
            remote_ip
        );
        return Err(ErrorUnauthorized("You are not allowed to perform updates!"));
    }

    // Check token
    let token = match req.headers().get("Token") {
        None => {
            return Err(ErrorUnauthorized("Token required!"));
        }
        Some(t) => t
            .to_str()
            .map_err(|_| ErrorInternalServerError("Failed to parse token!"))?,
    };

    if !token.eq(&args.update_token) || args.update_token.is_empty() {
        return Err(ErrorUnauthorized("Invalid update token!"));
    }

    // Get base folder to keep from tar-file
    let base_uri = match req.headers().get("BaseURI") {
        None => "/",
        Some(t) => t
            .to_str()
            .map_err(|_| ErrorInternalServerError("Failed to parse base URI to keep!"))?,
    };

    let mut new_files = Vec::new();

    // iterate over multipart stream
    if let Some(mut field) = payload.try_next().await? {
        let mut b = bytes::BytesMut::new();

        // Field in turn is stream of *Bytes* object
        while let Some(chunk) = field.try_next().await? {
            b.put(chunk);
        }

        let mut archive = tar::Archive::new(b.as_ref());
        for entry in archive
            .entries()
            .map_err(|_| ErrorInternalServerError("Failed to parse TAR archive!"))?
        {
            let mut file = entry?;
            let inner_path = file.header().path()?;
            let inner_path_str = inner_path.to_string_lossy();

            if !inner_path_str.starts_with(base_uri) {
                continue;
            }

            // Remove base URI before extraction
            let mut inner_path_str = &inner_path_str[base_uri.len()..];
            if inner_path_str.starts_with('/') {
                inner_path_str = &inner_path_str[1..];
            }
            let inner_path = Path::new(inner_path_str);

            if inner_path.is_dir() || inner_path_str.ends_with('/') || inner_path_str.is_empty() {
                continue;
            }

            // Read file to buffer
            let dest_file = args.storage_path().join(inner_path);
            let mut buf = Vec::with_capacity(file.size() as usize);
            file.read_to_end(&mut buf)?;

            new_files.push(NewFile {
                path: dest_file,
                bytes: buf,
            });
        }
    }

    // Check if at least one file was sent
    if new_files.is_empty() {
        return Err(ErrorBadRequest("No file to extract!"));
    }

    // Delete all current files in storage
    for entry in std::fs::read_dir(args.storage_path())? {
        let entry = entry?;
        if entry.path().is_dir() {
            std::fs::remove_dir_all(entry.path())?;
        } else {
            std::fs::remove_file(entry.path())?;
        }
    }

    // Apply new files
    for file in new_files {
        std::fs::create_dir_all(file.path.parent().unwrap())?;
        if let Err(e) = std::fs::write(&file.path, file.bytes) {
            log::error!("Failed to write new file {:?} : {:?}", file.path, e);
            return Err(ErrorInternalServerError("Failed to write a file!"));
        }
    }

    Ok(HttpResponse::Ok().into())
}

async fn default_files_handler(req: ServiceRequest) -> Result<ServiceResponse, Error> {
    let (req, _) = req.into_parts();
    let args: &web::Data<Args> = req.app_data().unwrap();

    // Search for alternate paths
    if args.can_bypass_html_ext
        && !req.path().ends_with(".html")
        && !req.path().ends_with('/')
        && !req.path().is_empty()
    {
        let alt_file = args
            .storage_path()
            .join(format!("{}.html", &req.path()[1..]));

        if alt_file.exists() {
            let file = NamedFile::open_async(alt_file).await?;
            let res = file.into_response(&req);
            return Ok(ServiceResponse::new(req, res));
        }
    }

    // Default handler
    if let Some(h) = args.default_handler_path() {
        let file = NamedFile::open_async(h).await?;
        let res = file.into_response(&req);
        Ok(ServiceResponse::new(req, res))
    }
    // Dummy response
    else {
        Ok(ServiceResponse::new(
            req,
            HttpResponse::NotFound().body("404 Not found"),
        ))
    }
}

#[actix_web::main]
async fn main() -> std::io::Result<()> {
    let args: Args = Args::parse();
    env_logger::init_from_env(env_logger::Env::new().default_filter_or("info"));

    log::info!("starting HTTP server at {}", args.listen_address);
    let listen_address = args.listen_address.to_string();

    if !args.storage_path().exists() {
        panic!("Specified files path does not exists!");
    }

    HttpServer::new(move || {
        App::new()
            // Update service
            .service(web::resource("/_mgmt/replace_files").route(web::post().to(replace_files)))
            // Serve a tree of static files at the web root and specify the index file.
            // Note that the root path should always be defined as the last item. The paths are
            // resolved in the order they are defined. If this would be placed before the `/images`
            // path then the service for the static images would never be reached.
            .service(
                Files::new("/", &args.files_path)
                    .index_file(&args.index_file)
                    .default_handler(fn_service(default_files_handler)),
            )
            // Enable the logger.
            .wrap(Logger::default())
            .app_data(web::Data::new(args.clone()))
    })
    .bind(listen_address)?
    .run()
    .await
}

#[cfg(test)]
mod test {
    use crate::Args;
    #[test]
    fn verify_cli() {
        use clap::CommandFactory;
        Args::command().debug_assert()
    }
}