Can update files

src/main.rs

@@ -1,10 +1,15 @@
+use std::io::Read;
 use std::path::{Path, PathBuf};
 
 use actix_files::{Files, NamedFile};
-use actix_web::{App, HttpResponse, HttpServer};
+use actix_multipart::Multipart;
+use actix_web::{App, Error, HttpRequest, HttpResponse, HttpServer, web};
 use actix_web::dev::{fn_service, ServiceRequest, ServiceResponse};
+use actix_web::error::{ErrorInternalServerError, ErrorUnauthorized};
 use actix_web::middleware::Logger;
+use bytes::BufMut;
 use clap::Parser;
+use futures_util::TryStreamExt;
 
 /// Simple pages server
 #[derive(Parser, Debug, Clone)]
@@ -50,6 +55,82 @@ impl Args {
 
 static mut ARGS: Option<Args> = None;
 
+struct NewFile {
+    path: PathBuf,
+    bytes: Vec<u8>,
+}
+
+/// Replace all the files of the website
+async fn replace_files(args: web::Data<Args>, req: HttpRequest, mut payload: Multipart) -> Result<HttpResponse, Error> {
+
+    // 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!"));
+    }
+
+    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()?;
+
+            if inner_path.is_absolute() {
+                return Err(ErrorInternalServerError("Cowardly refusing to unpack absolute file!"));
+            }
+
+            if inner_path.is_dir() || inner_path.to_string_lossy().ends_with("/") {
+                continue;
+            }
+
+            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,
+            });
+        }
+    }
+
+    // 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())?;
+        std::fs::write(file.path, file.bytes)?;
+    }
+
+    Ok(HttpResponse::Ok().into())
+}
+
 #[actix_web::main]
 async fn main() -> std::io::Result<()> {
     let args: Args = Args::parse();
@@ -67,6 +148,10 @@ async fn main() -> std::io::Result<()> {
     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`
@@ -90,6 +175,7 @@ async fn main() -> std::io::Result<()> {
 
             // Enable the logger.
             .wrap(Logger::default())
+            .app_data(web::Data::new(args.clone()))
     })
         .bind(listen_address)?
         .run()