feat: implement password protection

src/backend/Cargo.toml

@@ -18,6 +18,8 @@ nanoserde = "0.1.29"
 pty-process = "0.1.1"
 heim = { git = "https://github.com/heim-rs/heim", features = ["cpu", "disk", "host", "memory", "net", "process"] }
 infer = { version = "0.5.0", default-features = false }
+jwts = "0.2.3"
+sha2 = "0.9.8"
 
 [profile.release]
 lto = "fat"

src/backend/src/config.rs

@@ -6,6 +6,10 @@ pub struct Config {
     pub tls: bool,
     pub cert: String,
     pub key: String,
+
+    pub pass: bool,
+    pub hash: String,
+    pub secret: String,
 }
 
 pub fn config() -> Config {
@@ -32,7 +36,6 @@ pub fn config() -> Config {
     let port: u16 = cfg.get("port").unwrap_or(&Toml::Num(8088.0)).num() as u16;
 
     let tls = cfg.get("tls").unwrap_or(&Toml::Bool(false));
-
     let cert = cfg
         .get("cert")
         .unwrap_or(&Toml::Str(String::new()))
@@ -44,10 +47,25 @@ pub fn config() -> Config {
         .str()
         .to_string();
 
+    let pass = cfg.get("pass").unwrap_or(&Toml::Bool(false));
+    let hash = cfg
+        .get("hash")
+        .unwrap_or(&Toml::Str(String::new()))
+        .str()
+        .to_string();
+    let secret = cfg
+        .get("secret")
+        .unwrap_or(&Toml::Str(String::new()))
+        .str()
+        .to_string();
+
     Config {
         port,
         tls: tls == &Toml::Bool(true),
         cert,
         key,
+        pass: pass == &Toml::Bool(true),
+        hash,
+        secret,
     }
 }

src/backend/src/main.rs

@@ -1,5 +1,6 @@
 #![warn(clippy::pedantic)]
 #![warn(clippy::cognitive_complexity)]
+use sha2::{Digest, Sha512};
 use simple_logger::SimpleLogger;
 use warp::Filter;
 
@@ -9,6 +10,11 @@ mod systemdata;
 mod terminal;
 mod types;
 
+lazy_static::lazy_static! {
+    static ref CONFIG: config::Config = config::config();
+}
+
+#[allow(clippy::too_many_lines)]
 fn main() {
     tokio::runtime::Builder::new_multi_thread()
         .worker_threads(num_cpus::get().max(2)) // We have to use num_cpus because heim is async, and the runtime hasn't been started yet. Minimum of 2 threads.
@@ -18,8 +24,6 @@ fn main() {
         .block_on(async {
             const DIR: include_dir::Dir = include_dir::include_dir!("dist");
 
-            let cfg = config::config();
-
             SimpleLogger::new()
                 .with_level(log::LevelFilter::Info)
                 .env()
@@ -51,6 +55,44 @@ fn main() {
                     )
                 });
 
+            let login_route = warp::path("login")
+                .and(warp::post())
+                .and(warp::body::bytes())
+                .map(|pass| {
+                    if CONFIG.pass {
+                        let mut hasher = Sha512::new();
+                        hasher.update(pass);
+                        let shasum = format!("{:x?}", hasher.finalize())
+                            .split(&['[', ']', ',', ' '][..])
+                            .collect::<String>();
+                        if shasum == CONFIG.hash {
+                            let mut claims = jwts::Claims::new();
+                            claims.exp = Some(
+                                (std::time::SystemTime::now()
+                                    + std::time::Duration::from_secs(3600))
+                                .duration_since(std::time::UNIX_EPOCH)
+                                .unwrap()
+                                .as_secs(),
+                            );
+                            let mut token = jwts::jws::Token::with_payload(claims);
+                            let key =
+                                jwts::jws::Key::new(&CONFIG.secret, jwts::jws::Algorithm::HS256);
+                            return warp::reply::with_status(
+                                token.sign(&key).unwrap(),
+                                warp::http::StatusCode::OK,
+                            );
+                        }
+                        return warp::reply::with_status(
+                            "Unauthorized".to_string(),
+                            warp::http::StatusCode::UNAUTHORIZED,
+                        );
+                    }
+                    warp::reply::with_status(
+                        "No login needed".to_string(),
+                        warp::http::StatusCode::OK,
+                    )
+                });
+
             let terminal_route = warp::path!("ws" / "term")
                 .and(warp::ws())
                 .map(|ws: warp::ws::Ws| ws.on_upgrade(terminal::term_handler));
@@ -65,6 +107,7 @@ fn main() {
 
             let page_routes = favicon_route
                 .or(assets_route)
+                .or(login_route)
                 .or(main_route)
                 .with(warp::compression::gzip());
 
@@ -83,15 +126,15 @@ fn main() {
                     );
                 }));
 
-            if cfg.tls {
+            if CONFIG.tls {
                 warp::serve(routes)
                     .tls()
-                    .cert_path(cfg.cert)
-                    .key_path(cfg.key)
-                    .run(([0, 0, 0, 0], cfg.port))
+                    .cert_path(&CONFIG.cert)
+                    .key_path(&CONFIG.key)
+                    .run(([0, 0, 0, 0], CONFIG.port))
                     .await;
             } else {
-                warp::serve(routes).run(([0, 0, 0, 0], cfg.port)).await;
+                warp::serve(routes).run(([0, 0, 0, 0], CONFIG.port)).await;
             }
         });
 }

src/backend/src/sockets.rs

@@ -11,7 +11,7 @@ use tokio::sync::{
 use tokio::time::sleep;
 use warp::ws::Message;
 
-use crate::{systemdata, types};
+use crate::{systemdata, types, CONFIG};
 
 async fn main_handler(
     socket_ptr: Arc<Mutex<SplitSink<warp::ws::WebSocket, warp::ws::Message>>>,
@@ -292,6 +292,52 @@ pub async fn socket_handler(socket: warp::ws::WebSocket) {
                 break;
             }
             req = DeJson::deserialize_json(data.to_str().unwrap()).unwrap();
+            if CONFIG.pass {
+                let key = jwts::jws::Key::new(&CONFIG.secret, jwts::jws::Algorithm::HS256);
+                let verified: jwts::jws::Token<jwts::Claims>;
+                if let Ok(token) = jwts::jws::Token::verify_with_key(&req.token, &key) {
+                    verified = token;
+                } else {
+                    log::error!("Couldn't verify token");
+                    if !first_message {
+                        data_send.send(None).await.unwrap();
+                    }
+                    data_send
+                        .send(Some(types::Request {
+                            page: "/login".to_string(),
+                            token: String::new(),
+                            cmd: String::new(),
+                            args: Vec::new(),
+                        }))
+                        .await
+                        .unwrap();
+                    continue;
+                };
+                let config = jwts::ValidationConfig {
+                    iat_validation: false,
+                    nbf_validation: false,
+                    exp_validation: true,
+                    expected_iss: None,
+                    expected_sub: None,
+                    expected_aud: None,
+                    expected_jti: None,
+                };
+                if verified.validate_claims(&config).is_err() {
+                    if !first_message {
+                        data_send.send(None).await.unwrap();
+                    }
+                    data_send
+                        .send(Some(types::Request {
+                            page: "/login".to_string(),
+                            token: String::new(),
+                            cmd: String::new(),
+                            args: Vec::new(),
+                        }))
+                        .await
+                        .unwrap();
+                    continue;
+                }
+            }
             if req.cmd.is_empty() {
                 if first_message {
                     first_message = false;
@@ -333,6 +379,14 @@ pub async fn socket_handler(socket: warp::ws::WebSocket) {
             "/browser" => {
                 browser_handler(Arc::clone(&socket_ptr), &mut data_recv).await;
             }
+            "/login" => {
+                // Internal poll, see other thread
+                let _send = (*socket_ptr.lock().await)
+                    .send(Message::text(SerJson::serialize_json(&types::TokenError {
+                        error: true,
+                    })))
+                    .await;
+            }
             _ => {}
         }
     }

src/backend/src/systemdata.rs

@@ -422,7 +422,10 @@ pub fn services() -> Vec<types::ServiceData> {
 pub fn global() -> types::GlobalData {
     let update =
         fs::read_to_string("/run/dietpi/.update_available").unwrap_or_else(|_| String::new());
-    types::GlobalData { update }
+    types::GlobalData {
+        update,
+        login: crate::CONFIG.pass,
+    }
 }
 
 pub fn browser_dir(path: &std::path::Path) -> Vec<types::BrowserDirData> {

src/backend/src/types.rs

@@ -30,6 +30,8 @@ pub struct Request {
     pub cmd: String,
     #[nserde(default)]
     pub args: Vec<String>,
+    #[nserde(default)]
+    pub token: String,
 }
 
 #[derive(SerJson)]
@@ -91,6 +93,7 @@ pub struct ServiceList {
 #[derive(SerJson)]
 pub struct GlobalData {
     pub update: String,
+    pub login: bool,
 }
 
 #[derive(SerJson, Debug)]
@@ -112,3 +115,8 @@ pub struct BrowserFileData {
 pub struct BrowserList {
     pub contents: Vec<BrowserDirData>,
 }
+
+#[derive(SerJson)]
+pub struct TokenError {
+    pub error: bool,
+}

src/frontend/index.html

@@ -3,6 +3,8 @@
 
 <head>
   <meta charset="UTF-8" />
+  <!--<meta http-equiv="Content-Security-Policy"
+    content="default-src 'self'; font-src 'self'; img-src 'self'; script-src 'self'; style-src 'unsafe-inline' 'self';" />-->
   <link rel="icon" href="/favicon.png" />
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
   <title>DietPi Dashboard</title>