Add token role and merge endpoints

netmito/src/api/admin.rs

@@ -15,7 +15,7 @@ use crate::{
     schema::*,
     service::{
         self,
-        auth::{admin_auth_middleware, AuthAdminUser},
+        auth::{admin_middleware, auth_middleware, roles::USER, AuthAdminUser},
         name_validator,
     },
 };
@@ -50,9 +50,11 @@ pub fn admin_router(st: InfraPool, cancel_token: CancellationToken) -> Router<In
             "/shutdown",
             post(admin_shutdown_coordinator).with_state(cancel_token),
         )
+        // The last layer added runs first: auth attaches the user, then the admin check reads it.
+        .route_layer(middleware::from_fn(admin_middleware))
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
-            admin_auth_middleware,
+            auth_middleware::<USER>,
         ))
         .with_state(st)
 }

netmito/src/api/groups.rs

@@ -12,23 +12,32 @@ use crate::{
     schema::*,
     service::{
         self,
-        auth::{user_auth_middleware, AuthUser},
+        auth::{
+            auth_middleware,
+            roles::{USER, WORKER},
+            AuthIdentity, AuthUser,
+        },
     },
 };
 
 pub fn groups_router(st: InfraPool) -> Router<InfraPool> {
-    Router::new()
+    let shared_router = Router::new()
+        .route(
+            "/{group_name}/download/attachments/{*key}",
+            get(download_attachment),
+        )
+        .route_layer(middleware::from_fn_with_state(
+            st.clone(),
+            auth_middleware::<{ USER | WORKER }>,
+        ));
+    let user_router = Router::new()
         .route("/", post(create_group))
         .route("/{group_name}", get(get_group))
         .route("/{group_name}/users/remove", post(remove_user_group))
         .route(
             "/{group_name}/users",
             put(update_user_group).delete(remove_user_group_params),
         )
-        .route(
-            "/{group_name}/download/attachments/{*key}",
-            get(download_attachment),
-        )
         .route(
             "/{group_name}/attachments/{*key}",
             get(query_single_attachment).delete(delete_attachment),
@@ -53,8 +62,11 @@ pub fn groups_router(st: InfraPool) -> Router<InfraPool> {
         )
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
-            user_auth_middleware,
-        ))
+            auth_middleware::<USER>,
+        ));
+    Router::new()
+        .merge(shared_router)
+        .merge(user_router)
         .with_state(st)
 }
 
@@ -172,20 +184,26 @@ pub async fn upload_attachment(
 }
 
 pub async fn download_attachment(
-    Extension(u): Extension<AuthUser>,
+    Extension(p): Extension<AuthIdentity>,
     State(pool): State<InfraPool>,
     Path((group_name, key)): Path<(String, String)>,
 ) -> Result<Json<RemoteResourceDownloadResp>, ApiError> {
-    let attachment = service::s3::user_get_attachment(&pool, u.id, group_name, key)
-        .await
-        .map_err(|e| match e {
-            crate::error::Error::AuthError(err) => ApiError::AuthError(err),
-            crate::error::Error::ApiError(e) => e,
-            _ => {
-                tracing::error!("{}", e);
-                ApiError::InternalServerError
-            }
-        })?;
+    let attachment = match p {
+        AuthIdentity::User(u) => {
+            service::s3::user_get_attachment(&pool, u.id, group_name, key).await
+        }
+        AuthIdentity::Worker(w) => {
+            service::s3::worker_get_attachment(&pool, w.uuid, group_name, key).await
+        }
+    }
+    .map_err(|e| match e {
+        crate::error::Error::AuthError(err) => ApiError::AuthError(err),
+        crate::error::Error::ApiError(e) => e,
+        _ => {
+            tracing::error!("{}", e);
+            ApiError::InternalServerError
+        }
+    })?;
     Ok(Json(attachment))
 }
 

netmito/src/api/mod.rs

@@ -32,7 +32,7 @@ use crate::{
     config::{InfraPool, REDIS_CONNECTION_INFO},
     error::ApiError,
     schema::RedisConnectionInfo,
-    service::auth::{user_auth_middleware, user_auth_with_name_middleware, AuthUser},
+    service::auth::{auth_middleware, roles::USER, AuthUser},
 };
 
 pub fn router(st: InfraPool, cancel_token: CancellationToken) -> Router {
@@ -42,7 +42,7 @@ pub fn router(st: InfraPool, cancel_token: CancellationToken) -> Router {
             "/auth",
             get(users::user_auth).layer(middleware::from_fn_with_state(
                 st.clone(),
-                user_auth_with_name_middleware,
+                auth_middleware::<USER>,
             )),
         )
         .route(
@@ -54,21 +54,21 @@ pub fn router(st: InfraPool, cancel_token: CancellationToken) -> Router {
             "/refresh",
             post(users::refresh_token).layer(middleware::from_fn_with_state(
                 st.clone(),
-                user_auth_middleware,
+                auth_middleware::<USER>,
             )),
         )
         .route(
             "/revoke",
             post(users::revoke).layer(middleware::from_fn_with_state(
                 st.clone(),
-                user_auth_middleware,
+                auth_middleware::<USER>,
             )),
         )
         .route(
             "/redis",
             get(query_redis_connection_info).layer(middleware::from_fn_with_state(
                 st.clone(),
-                user_auth_middleware,
+                auth_middleware::<USER>,
             )),
         )
         .nest("/users", users::users_router(st.clone()))

netmito/src/api/suites.rs

@@ -16,7 +16,7 @@ use crate::{
     },
     service::{
         self,
-        auth::{user_auth_middleware, AuthUser},
+        auth::{auth_middleware, roles::USER, AuthUser},
     },
 };
 
@@ -29,7 +29,7 @@ pub fn suites_router(st: InfraPool) -> Router<InfraPool> {
         .route("/{uuid}/agents/override", post(override_agents_for_suite))
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
-            user_auth_middleware,
+            auth_middleware::<USER>,
         ))
         .with_state(st)
 }

netmito/src/api/tasks.rs

@@ -13,22 +13,35 @@ use crate::{
     schema::*,
     service::{
         self,
-        auth::{user_auth_middleware, AuthUser},
+        auth::{
+            auth_middleware,
+            roles::{USER, WORKER},
+            AuthIdentity, AuthUser, AuthWorker,
+        },
     },
 };
 
 pub fn tasks_router(st: InfraPool) -> Router<InfraPool> {
-    Router::new()
-        .route("/", post(submit_task))
-        .route(
-            "/{uuid}",
-            get(query_task).put(change_task).delete(cancel_task),
-        )
-        .route("/{uuid}/labels", put(change_task_labels))
+    let shared_router = Router::new()
+        .route("/{uuid}", get(query_task))
         .route(
             "/{uuid}/download/artifacts/{content_type}",
             get(download_artifact),
         )
+        .route_layer(middleware::from_fn_with_state(
+            st.clone(),
+            auth_middleware::<{ USER | WORKER }>,
+        ));
+    let worker_router = Router::new()
+        .route("/report", post(report_task))
+        .route_layer(middleware::from_fn_with_state(
+            st.clone(),
+            auth_middleware::<WORKER>,
+        ));
+    let user_router = Router::new()
+        .route("/", post(submit_task))
+        .route("/{uuid}", put(change_task).delete(cancel_task))
+        .route("/{uuid}/labels", put(change_task_labels))
         .route("/{uuid}/artifacts/{content_type}", delete(delete_artifact))
         .route("/{uuid}/artifacts", post(upload_artifact))
         .route("/query", post(query_tasks))
@@ -50,8 +63,12 @@ pub fn tasks_router(st: InfraPool) -> Router<InfraPool> {
         .route("/submit", post(batch_submit_tasks))
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
-            user_auth_middleware,
-        ))
+            auth_middleware::<USER>,
+        ));
+    Router::new()
+        .merge(shared_router)
+        .merge(worker_router)
+        .merge(user_router)
         .with_state(st)
 }
 
@@ -130,7 +147,7 @@ pub async fn cancel_task(
 }
 
 pub async fn query_task(
-    Extension(_): Extension<AuthUser>,
+    Extension(_): Extension<AuthIdentity>,
     State(pool): State<InfraPool>,
     Path(uuid): Path<Uuid>,
 ) -> Result<Json<TaskQueryResp>, ApiError> {
@@ -147,6 +164,21 @@ pub async fn query_task(
     Ok(Json(task))
 }
 
+pub async fn report_task(
+    Extension(w): Extension<AuthWorker>,
+    State(pool): State<InfraPool>,
+    Json(ReportTaskReq { id, op }): Json<ReportTaskReq>,
+) -> Result<Json<ReportTaskResp>, ApiError> {
+    match service::worker::report_task(w.id, w.uuid, id, op, &pool).await {
+        Ok(url) => Ok(Json(ReportTaskResp { url })),
+        Err(crate::error::Error::ApiError(e)) => Err(e),
+        Err(e) => {
+            tracing::error!("{}", e);
+            Err(ApiError::InternalServerError)
+        }
+    }
+}
+
 pub async fn query_tasks(
     Extension(u): Extension<AuthUser>,
     State(pool): State<InfraPool>,
@@ -222,7 +254,7 @@ pub async fn upload_artifact(
 }
 
 pub async fn download_artifact(
-    Extension(_): Extension<AuthUser>,
+    Extension(_): Extension<AuthIdentity>,
     State(pool): State<InfraPool>,
     Path((uuid, content_type)): Path<(Uuid, ArtifactContentType)>,
 ) -> Result<Json<RemoteResourceDownloadResp>, ApiError> {

netmito/src/api/users.rs

@@ -13,7 +13,7 @@ use crate::{
     schema::*,
     service::{
         self,
-        auth::{user_auth_middleware, AuthUser, AuthUserWithName},
+        auth::{auth_middleware, roles::USER, AuthUser},
     },
 };
 
@@ -23,7 +23,7 @@ pub fn users_router(st: InfraPool) -> Router<InfraPool> {
         .route("/groups", get(query_groups))
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
-            user_auth_middleware,
+            auth_middleware::<USER>,
         ))
         .with_state(st)
 }
@@ -52,7 +52,7 @@ pub async fn user_login(
     Ok(Json(UserLoginResp { token }))
 }
 
-pub async fn user_auth(Extension(u): Extension<AuthUserWithName>) -> String {
+pub async fn user_auth(Extension(u): Extension<AuthUser>) -> String {
     u.username
 }