From 6dd6e46aab0361bbb3f7005a27b628b3960514ec Mon Sep 17 00:00:00 2001
From: catvayor <catvayor@katvayor.net>
Date: Mon, 17 Jun 2024 09:11:01 +0200
Subject: [PATCH] cleanner admin auth

---
 src/admin.rs  | 105 ++++++++++++++++++++++++++------------------------
 src/global.rs |   1 -
 src/main.rs   |   3 --
 3 files changed, 54 insertions(+), 55 deletions(-)

diff --git a/src/admin.rs b/src/admin.rs
index ce3ce4e..cfdb6a9 100644
--- a/src/admin.rs
+++ b/src/admin.rs
@@ -1,4 +1,6 @@
 use rocket::{
+    http::Status,
+    request::{self, FromRequest, Outcome, Request},
     response::stream::{Event, EventStream},
     serde::json::Json,
     tokio::{
@@ -11,74 +13,75 @@ use rocket_dyn_templates::{context, Template};
 
 use crate::global::*;
 
-#[get("/?<tok>&<dbg>")]
-fn admin_page(
-    tok: Option<AdminKey>,
-    dbg: Option<bool>,
-    admin_key: &State<AdminKey>,
-) -> Option<Template> {
-    if tok == Some(admin_key.to_string()) {
-        Some(Template::render(
-            "admin",
-            context! { tok: tok.unwrap(), dbg: dbg.unwrap_or(false) },
-        ))
-    } else {
-        None
+struct AdminAuth(String);
+
+#[rocket::async_trait]
+impl<'r> FromRequest<'r> for AdminAuth {
+    type Error = ();
+
+    async fn from_request(req: &'r Request<'_>) -> request::Outcome<Self, Self::Error> {
+        if let Some(query) = req.uri().query() {
+            for (key, value) in query.segments() {
+                if key == "tok" {
+                    return if value == req.rocket().state::<Config>().unwrap().admin_token {
+                        Outcome::Success(AdminAuth(value.to_string()))
+                    } else {
+                        //498 Token expired/invalid
+                        Outcome::Error((Status::new(498), ()))
+                    };
+                }
+            }
+        }
+        Outcome::Error((Status::Unauthorized, ()))
     }
 }
 
-#[patch("/<id>?<tok>", data = "<nstate>")]
+#[get("/?<dbg>")]
+fn admin_page(auth: AdminAuth, dbg: Option<bool>) -> Template {
+    Template::render("admin", context! { tok: auth.0, dbg: dbg.unwrap_or(false) })
+}
+
+#[patch("/<id>", data = "<nstate>")]
 fn admin_set_state(
-    tok: Option<AdminKey>,
+    _auth: AdminAuth,
     id: &str,
-    admin_key: &State<AdminKey>,
     nstate: Json<TrackedState>,
     tracking: &State<Tracking>,
     evt_queue: &State<TrackingEventQueue>,
     admin_queue: &State<AdminEventQueue>,
-) -> Option<()> {
-    if tok == Some(admin_key.to_string()) {
-        let tracked = tracking.get(&id.to_string()).unwrap();
-        tracked.write().unwrap().state = nstate.into_inner();
-        state_update(&tracked.read().unwrap(), &evt_queue, &admin_queue);
-        Some(())
-    } else {
-        None
-    }
+) {
+    let tracked = tracking.get(&id.to_string()).unwrap(); //TODO: clean error
+    tracked.write().unwrap().state = nstate.into_inner();
+    state_update(&tracked.read().unwrap(), &evt_queue, &admin_queue);
 }
 
-#[get("/events?<tok>")]
+#[get("/events")]
 fn admin_events<'a>(
-    tok: Option<AdminKey>,
-    admin_key: &State<AdminKey>,
+    _auth: AdminAuth,
     admin_queue: &'a State<AdminEventQueue>,
     tracking: &State<Tracking>,
     config: &State<Config>,
     mut shutdown: Shutdown,
-) -> Option<EventStream![Event + 'a]> {
-    if tok == Some(admin_key.to_string()) {
-        let full_info: Vec<AdminTrackedInfo> = tracking
-            .iter()
-            .map(|(_, tracked)| admin_view(&tracked.read().unwrap()))
-            .collect();
-        let timeout = Duration::from_millis(config.event_timeout);
-        Some(EventStream! {
-            yield Event::json(&full_info).event("full_update");
-            let mut interval = time::interval(timeout);
-            loop {
-                select!{
-                    _ = interval.tick() =>{
-                        for evt in evts_to_send(admin_queue, timeout){
-                            //println!("{:?}", evt);
-                            yield evt;
-                        }
-                    },
-                    _ = &mut shutdown => break
-                }
+) -> EventStream![Event + 'a] {
+    let full_info: Vec<AdminTrackedInfo> = tracking
+        .iter()
+        .map(|(_, tracked)| admin_view(&tracked.read().unwrap()))
+        .collect();
+    let timeout = Duration::from_millis(config.event_timeout);
+    EventStream! {
+        yield Event::json(&full_info).event("full_update");
+        let mut interval = time::interval(timeout);
+        loop {
+            select!{
+                _ = interval.tick() =>{
+                    for evt in evts_to_send(admin_queue, timeout){
+                        //println!("{:?}", evt);
+                        yield evt;
+                    }
+                },
+                _ = &mut shutdown => break
             }
-        })
-    } else {
-        None
+        }
     }
 }
 
diff --git a/src/global.rs b/src/global.rs
index d5b1209..a9725fa 100644
--- a/src/global.rs
+++ b/src/global.rs
@@ -171,7 +171,6 @@ impl From<QueuedEvent> for Event {
 pub type Tracking = Arc<HashMap<String, RwLock<Tracked>>>;
 pub type TrackingEventQueue = Arc<HashMap<String, RwLock<VecDeque<QueuedEvent>>>>;
 pub type AdminEventQueue = Arc<RwLock<VecDeque<QueuedEvent>>>;
-pub type AdminKey = String;
 
 #[derive(Serialize)]
 #[serde(crate = "rocket::serde")]
diff --git a/src/main.rs b/src/main.rs
index d5ad1ee..fea3101 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -98,15 +98,12 @@ async fn rocket() -> _ {
             .collect(),
     );
     let admin_evt_queue: AdminEventQueue = Arc::new(RwLock::new(VecDeque::new()));
-    let key: AdminKey = config.admin_token.clone();
-    println!("Admin token: {}", key);
     let mut rocket = rocket
         .attach(Template::fairing())
         .attach(AdHoc::config::<Config>())
         .manage(tracking.clone())
         .manage(evt_queue.clone())
         .manage(admin_evt_queue.clone())
-        .manage(key)
         .mount("/", routes![index])
         .mount("/track", track::routes())
         .mount("/admin", admin::routes());