From 70070a1f6a159cc30748fc0227529585f02177bf Mon Sep 17 00:00:00 2001 From: un-lock-able Date: Fri, 23 Oct 2026 16:35:40 +0800 Subject: [PATCH 1/2] refactor(auth): add token role and merge executor-facing endpoints --- netmito/src/api/admin.rs | 6 +- netmito/src/api/groups.rs | 56 +++++++---- netmito/src/api/mod.rs | 10 +- netmito/src/api/suites.rs | 4 +- netmito/src/api/tasks.rs | 56 ++++++++--- netmito/src/api/users.rs | 6 +- netmito/src/api/workers.rs | 91 ++---------------- netmito/src/schema/task.rs | 2 + netmito/src/service/auth/mod.rs | 146 ++++++++++++++-------------- netmito/src/service/auth/token.rs | 34 +++++++ netmito/src/service/s3.rs | 81 +++++----------- netmito/src/service/worker/mod.rs | 11 +++ netmito/src/worker.rs | 19 ++-- openapi.yaml | 152 ++++++++---------------------- 14 files changed, 296 insertions(+), 378 deletions(-) diff --git a/netmito/src/api/admin.rs b/netmito/src/api/admin.rs index 182512b2..ab333aa2 100644 --- a/netmito/src/api/admin.rs +++ b/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, )) .with_state(st) } diff --git a/netmito/src/api/groups.rs b/netmito/src/api/groups.rs index a93adeb8..9000c1fe 100644 --- a/netmito/src/api/groups.rs +++ b/netmito/src/api/groups.rs @@ -12,12 +12,25 @@ 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 { - 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)) @@ -25,10 +38,6 @@ pub fn groups_router(st: InfraPool) -> Router { "/{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 { ) .route_layer(middleware::from_fn_with_state( st.clone(), - user_auth_middleware, - )) + auth_middleware::, + )); + 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, + Extension(p): Extension, State(pool): State, Path((group_name, key)): Path<(String, String)>, ) -> Result, 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)) } diff --git a/netmito/src/api/mod.rs b/netmito/src/api/mod.rs index 9269fd75..cbb6bf1b 100644 --- a/netmito/src/api/mod.rs +++ b/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::, )), ) .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::, )), ) .route( "/revoke", post(users::revoke).layer(middleware::from_fn_with_state( st.clone(), - user_auth_middleware, + auth_middleware::, )), ) .route( "/redis", get(query_redis_connection_info).layer(middleware::from_fn_with_state( st.clone(), - user_auth_middleware, + auth_middleware::, )), ) .nest("/users", users::users_router(st.clone())) diff --git a/netmito/src/api/suites.rs b/netmito/src/api/suites.rs index 56cdea7a..149cbb68 100644 --- a/netmito/src/api/suites.rs +++ b/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 { .route("/{uuid}/agents/override", post(override_agents_for_suite)) .route_layer(middleware::from_fn_with_state( st.clone(), - user_auth_middleware, + auth_middleware::, )) .with_state(st) } diff --git a/netmito/src/api/tasks.rs b/netmito/src/api/tasks.rs index e834b445..a00d72d9 100644 --- a/netmito/src/api/tasks.rs +++ b/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 { - 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::, + )); + 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 { .route("/submit", post(batch_submit_tasks)) .route_layer(middleware::from_fn_with_state( st.clone(), - user_auth_middleware, - )) + auth_middleware::, + )); + 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, + Extension(_): Extension, State(pool): State, Path(uuid): Path, ) -> Result, ApiError> { @@ -147,6 +164,21 @@ pub async fn query_task( Ok(Json(task)) } +pub async fn report_task( + Extension(w): Extension, + State(pool): State, + Json(ReportTaskReq { id, op }): Json, +) -> Result, 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, State(pool): State, @@ -222,7 +254,7 @@ pub async fn upload_artifact( } pub async fn download_artifact( - Extension(_): Extension, + Extension(_): Extension, State(pool): State, Path((uuid, content_type)): Path<(Uuid, ArtifactContentType)>, ) -> Result, ApiError> { diff --git a/netmito/src/api/users.rs b/netmito/src/api/users.rs index 864bf391..75db2b07 100644 --- a/netmito/src/api/users.rs +++ b/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 { .route("/groups", get(query_groups)) .route_layer(middleware::from_fn_with_state( st.clone(), - user_auth_middleware, + auth_middleware::, )) .with_state(st) } @@ -52,7 +52,7 @@ pub async fn user_login( Ok(Json(UserLoginResp { token })) } -pub async fn user_auth(Extension(u): Extension) -> String { +pub async fn user_auth(Extension(u): Extension) -> String { u.username } diff --git a/netmito/src/api/workers.rs b/netmito/src/api/workers.rs index 571fb9f6..8cdaf493 100644 --- a/netmito/src/api/workers.rs +++ b/netmito/src/api/workers.rs @@ -9,14 +9,15 @@ use uuid::Uuid; use crate::{ config::{InfraPool, REDIS_CONNECTION_INFO}, - entity::content::ArtifactContentType, error::{ApiError, ApiResult, Error}, schema::*, service::{ self, auth::{ - token::generate_worker_token, user_auth_middleware, worker_auth_middleware, AuthUser, - AuthWorker, + auth_middleware, + roles::{USER, WORKER}, + token::generate_worker_token, + AuthUser, AuthWorker, }, }, }; @@ -25,19 +26,10 @@ pub fn workers_router(st: InfraPool) -> Router { let worker_auth_router = Router::new() .route("/", delete(worker_unregister)) .route("/heartbeat", post(heartbeat)) - .route("/tasks", post(worker_report_task).get(worker_fetch_task)) - .route("/tasks/{uuid}", get(worker_query_task)) - .route( - "/tasks/{uuid}/artifacts/{content_type}", - get(download_artifact), - ) - .route( - "/tasks/{uuid}/attachments/{*key}", - get(worker_download_attachment), - ) + .route("/tasks", get(worker_fetch_task)) .route_layer(middleware::from_fn_with_state( st.clone(), - worker_auth_middleware, + auth_middleware::, )) .with_state(st.clone()); let user_auth_router = Router::new() @@ -58,7 +50,7 @@ pub fn workers_router(st: InfraPool) -> Router { .route("/shutdown/list", post(user_shutdown_workers_by_uuids)) .route_layer(middleware::from_fn_with_state( st.clone(), - user_auth_middleware, + auth_middleware::, )) .with_state(st.clone()); Router::new() @@ -145,75 +137,6 @@ pub async fn worker_fetch_task( } } -pub async fn worker_report_task( - Extension(w): Extension, - State(pool): State, - Json(ReportTaskReq { id, op }): Json, -) -> ApiResult> { - match service::worker::report_task(w.id, w.uuid, id, op, &pool).await { - Ok(url) => Ok(Json(ReportTaskResp { url })), - Err(Error::ApiError(e)) => Err(e), - Err(e) => { - tracing::error!("{}", e); - Err(ApiError::InternalServerError) - } - } -} - -pub async fn download_artifact( - Extension(_): Extension, - State(pool): State, - Path((uuid, content_type)): Path<(Uuid, ArtifactContentType)>, -) -> Result, ApiError> { - let artifact = service::s3::download_artifact_by_uuid(&pool, uuid, content_type) - .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(artifact)) -} - -pub async fn worker_download_attachment( - Extension(_): Extension, - State(pool): State, - Path((uuid, key)): Path<(Uuid, String)>, -) -> Result, ApiError> { - let attachment = service::s3::worker_download_attachment(&pool, uuid, 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)) -} - -pub async fn worker_query_task( - Extension(_): Extension, - State(pool): State, - Path(uuid): Path, -) -> Result, ApiError> { - let task = service::task::get_task_by_uuid(&pool, uuid) - .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(task)) -} - pub async fn user_query_worker( Extension(_): Extension, State(pool): State, diff --git a/netmito/src/schema/task.rs b/netmito/src/schema/task.rs index 6c6ab5c4..56a73cff 100644 --- a/netmito/src/schema/task.rs +++ b/netmito/src/schema/task.rs @@ -15,6 +15,8 @@ pub struct WorkerTaskResp { pub id: i64, pub uuid: Uuid, pub upstream_task_uuid: Option, + /// The task's group, where its attachment resources are fetched from. + pub group_name: String, pub spec: ExecSpec, pub exec_options: Option, } diff --git a/netmito/src/service/auth/mod.rs b/netmito/src/service/auth/mod.rs index 20985ad9..2c10a3cc 100644 --- a/netmito/src/service/auth/mod.rs +++ b/netmito/src/service/auth/mod.rs @@ -22,21 +22,35 @@ use crate::{ error::{ApiError, AuthError}, schema::{UserChangePasswordReq, UserLoginReq}, }; -use token::{generate_token, verify_token}; +use token::{generate_token, verify_token, TokenClaims, TokenRole}; // TODO: should check if the structure and logic is high cohesive and low coupling enough, and // refactor if necessary. // TODO: check if we allow non-expiry lifetime -#[derive(Debug, Clone)] -pub struct AuthUser { - pub id: i64, +/// Role bits for [`auth_middleware`]'s `ROLES`. +pub mod roles { + use super::token::TokenRole; + + pub const USER: u8 = 1 << 0; + pub const WORKER: u8 = 1 << 1; + pub const AGENT: u8 = 1 << 2; + + pub(super) fn bit_mask_of(role: TokenRole) -> u8 { + match role { + TokenRole::User => USER, + TokenRole::Worker => WORKER, + TokenRole::Agent => AGENT, + } + } } #[derive(Debug, Clone)] -pub struct AuthUserWithName { +pub struct AuthUser { pub id: i64, pub username: String, + /// Mitosis system admin (`users.admin`) + pub admin: bool, } #[derive(Debug, Clone)] @@ -50,6 +64,13 @@ pub struct AuthWorker { pub uuid: Uuid, } +/// The caller of a route open to more than one role. +#[derive(Debug, Clone)] +pub enum AuthIdentity { + User(AuthUser), + Worker(AuthWorker), +} + fn generate_auth_signature() -> i64 { StdRng::from_os_rng().next_u32() as i64 + 1 } @@ -301,40 +322,70 @@ pub async fn admin_change_password( Ok(()) } -pub async fn user_auth_middleware( +/// Accept the token roles set in `ROLES` (a union of [`roles`] bits) and attach the +/// caller to the request: always as [`AuthIdentity`], and also as the role's own type +/// (`AuthUser`, `AuthWorker`) for handlers that serve a single role. +pub async fn auth_middleware( State(pool): State, TypedHeader(Authorization(bearer)): TypedHeader>, mut req: Request, next: Next, ) -> Result { - let auth_user = user_auth(&pool.db, &bearer).await?; - req.extensions_mut().insert(AuthUser { id: auth_user.id }); + let claims = verify_token(bearer.token()).map_err(|_| AuthError::InvalidToken)?; + if ROLES & roles::bit_mask_of(claims.role) == 0 { + return Err(AuthError::PermissionDenied.into()); + } + let identity = match claims.role { + TokenRole::User => { + let user = user_auth(&pool.db, &claims).await?; + let auth_user = AuthUser { + id: user.id, + username: user.username, + admin: user.admin, + }; + req.extensions_mut().insert(auth_user.clone()); + AuthIdentity::User(auth_user) + } + TokenRole::Worker => { + let auth_worker = worker_auth(&pool.db, &claims).await?; + req.extensions_mut().insert(auth_worker.clone()); + AuthIdentity::Worker(auth_worker) + } + // No agent tokens are issued yet. + TokenRole::Agent => return Err(AuthError::InvalidToken.into()), + }; + req.extensions_mut().insert(identity); Ok(next.run(req).await) } -pub async fn user_auth_with_name_middleware( - State(pool): State, - TypedHeader(Authorization(bearer)): TypedHeader>, +/// Admit only system admins. Must run inside `auth_middleware::`, which attaches +/// the `AuthUser` this reads. +pub async fn admin_middleware( mut req: Request, next: Next, ) -> Result { - let auth_user = user_auth(&pool.db, &bearer).await?; - req.extensions_mut().insert(AuthUserWithName { - id: auth_user.id, - username: auth_user.username, - }); + let user = req + .extensions() + .get::() + .ok_or(ApiError::InternalServerError)?; + if !user.admin { + return Err(AuthError::PermissionDenied.into()); + } + let admin_user = AuthAdminUser { id: user.id }; + req.extensions_mut().insert(admin_user); Ok(next.run(req).await) } -async fn user_auth(db: &DatabaseConnection, bearer: &Bearer) -> Result { - let token = bearer.token(); - let claims = verify_token(token).map_err(|_| AuthError::InvalidToken)?; +async fn user_auth( + db: &DatabaseConnection, + claims: &TokenClaims<'_>, +) -> Result { if claims.exp.is_none() { return Err(AuthError::InvalidToken); } let user = User::Entity::find() - .filter(User::Column::Username.eq(claims.sub)) + .filter(User::Column::Username.eq(claims.sub.as_ref())) .one(db) .await .map_err(|_| AuthError::WrongCredentials)? @@ -349,57 +400,10 @@ async fn user_auth(db: &DatabaseConnection, bearer: &Bearer) -> Result, - TypedHeader(Authorization(bearer)): TypedHeader>, - mut req: Request, - next: Next, -) -> Result { - let admin_user = admin_auth(&pool.db, &bearer).await?; - req.extensions_mut().insert(admin_user); - Ok(next.run(req).await) -} - -async fn admin_auth(db: &DatabaseConnection, bearer: &Bearer) -> Result { - let token = bearer.token(); - let claims = verify_token(token).map_err(|_| AuthError::InvalidToken)?; - if claims.exp.is_none() { - return Err(AuthError::InvalidToken); - } - - let user = User::Entity::find() - .filter(User::Column::Username.eq(claims.sub)) - .one(db) - .await - .map_err(|_| AuthError::WrongCredentials)? - .ok_or(AuthError::WrongCredentials)?; - if user.admin { - if user.state != UserState::Active { - Err(AuthError::PermissionDenied) - } else if user.auth_signature != Some(claims.sign) { - Err(AuthError::WrongCredentials) - } else { - Ok(AuthAdminUser { id: user.id }) - } - } else { - Err(AuthError::PermissionDenied) - } -} - -pub async fn worker_auth_middleware( - State(pool): State, - TypedHeader(Authorization(bearer)): TypedHeader>, - mut req: Request, - next: Next, -) -> Result { - let auth_worker = worker_auth(&pool.db, &bearer).await?; - req.extensions_mut().insert(auth_worker); - Ok(next.run(req).await) -} - -async fn worker_auth(db: &DatabaseConnection, bearer: &Bearer) -> Result { - let token = bearer.token(); - let claims = verify_token(token).map_err(|_| AuthError::InvalidToken)?; +async fn worker_auth( + db: &DatabaseConnection, + claims: &TokenClaims<'_>, +) -> Result { let uuid = Uuid::parse_str(&claims.sub).map_err(|_| AuthError::InvalidToken)?; let worker = Worker::Entity::find() diff --git a/netmito/src/service/auth/token.rs b/netmito/src/service/auth/token.rs index 730fe290..3b8d8a31 100644 --- a/netmito/src/service/auth/token.rs +++ b/netmito/src/service/auth/token.rs @@ -7,6 +7,36 @@ use time::OffsetDateTime; use crate::error::{ApiError, DecodeTokenError}; +/// Which kind of caller a token was issued to. Serialized as its number; a token +/// without a role does not decode. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(try_from = "u8", into = "u8")] +#[repr(u8)] +pub enum TokenRole { + User = 1, + Worker = 2, + Agent = 3, +} + +impl From for u8 { + fn from(role: TokenRole) -> Self { + role as u8 + } +} + +impl TryFrom for TokenRole { + type Error = String; + + fn try_from(value: u8) -> Result { + match value { + 1 => Ok(Self::User), + 2 => Ok(Self::Worker), + 3 => Ok(Self::Agent), + _ => Err(format!("invalid token role {value}")), + } + } +} + #[derive(Debug, Serialize, Deserialize)] pub struct TokenClaims<'a> { /// username @@ -20,6 +50,7 @@ pub struct TokenClaims<'a> { pub exp: Option, /// random number pub sign: i64, + pub role: TokenRole, } pub fn generate_token(username: T, sign: i64) -> crate::error::Result @@ -35,6 +66,7 @@ where sub: Cow::from(username.as_ref()), exp: Some(OffsetDateTime::now_utc() + token_ttl.token_expires_in), sign, + role: TokenRole::User, }; let encoding_key = crate::config::ENCODING_KEY @@ -69,6 +101,7 @@ where sub: Cow::from(username.as_ref()), exp, sign, + role: TokenRole::Worker, }; let encoding_key = crate::config::ENCODING_KEY @@ -134,6 +167,7 @@ MCowBQYDK2VwAyEAa+75592cq25dhxP5a9zZMvLn+7yXyq6cUj2xmjB2PW4=\n\ sub: Cow::from("test_user"), exp, sign: 42, + role: TokenRole::User, }; encode_token(&claims, &test_encoding_key()).expect("token should encode") } diff --git a/netmito/src/service/s3.rs b/netmito/src/service/s3.rs index cdab0e14..e56b9513 100644 --- a/netmito/src/service/s3.rs +++ b/netmito/src/service/s3.rs @@ -8,7 +8,7 @@ use reqwest::{header::CONTENT_LENGTH, Response}; use sea_orm::{ prelude::*, sea_query::{Expr, Query}, - FromQueryResult, Set, TransactionTrait, + FromQueryResult, QuerySelect, Set, TransactionTrait, }; use tokio::io::AsyncWriteExt; use tokio_util::io::ReaderStream; @@ -28,7 +28,10 @@ use crate::{ error::AuthError, }; use crate::{ - entity::{content::AttachmentContentType, state::GroupState}, + entity::{ + content::AttachmentContentType, + state::{GroupState, TaskState}, + }, error::Error, schema::{AttachmentQueryInfo, AttachmentsQueryReq, TasksQueryReq}, }; @@ -550,65 +553,31 @@ pub async fn user_upload_artifact( group_upload_artifact(pool, task, content_type, content_length).await } -#[derive(FromQueryResult)] -struct GroupInfo { - id: i64, - group_name: String, -} - -pub async fn worker_download_attachment( +/// A worker may read a group's attachments only while it runs one of that group's tasks. +pub async fn worker_get_attachment( pool: &InfraPool, - uuid: Uuid, + worker_uuid: Uuid, + group_name: String, key: String, ) -> Result { - let builder = pool.db.get_database_backend(); - let active_group_name_stmt = Query::select() - .column((Group::Entity, Group::Column::Id)) - .column((Group::Entity, Group::Column::GroupName)) - .from(Group::Entity) - .join( - sea_orm::JoinType::Join, - ActiveTask::Entity, - Expr::col((ActiveTask::Entity, ActiveTask::Column::GroupId)) - .eq(Expr::col((Group::Entity, Group::Column::Id))), - ) - .and_where(Expr::col((ActiveTask::Entity, ActiveTask::Column::Uuid)).eq(uuid)) - .limit(1) - .to_owned(); - let GroupInfo { - id: group_id, - group_name, - } = match GroupInfo::find_by_statement(builder.build(&active_group_name_stmt)) + let group_id = Group::Entity::find() + .filter(Group::Column::GroupName.eq(group_name.clone())) .one(&pool.db) .await? - { - Some(g) => g, - None => { - let archived_group_name_stmt = Query::select() - .column((Group::Entity, Group::Column::GroupName)) - .from(Group::Entity) - .join( - sea_orm::JoinType::Join, - ArchivedTask::Entity, - Expr::col((ArchivedTask::Entity, ArchivedTask::Column::GroupId)) - .eq(Expr::col((Group::Entity, Group::Column::Id))), - ) - .and_where(Expr::col((ArchivedTask::Entity, ArchivedTask::Column::Uuid)).eq(uuid)) - .limit(1) - .to_owned(); - match GroupInfo::find_by_statement(builder.build(&archived_group_name_stmt)) - .one(&pool.db) - .await? - { - Some(g) => g, - None => { - return Err( - crate::error::ApiError::NotFound(format!("Task with uuid {uuid}")).into(), - ) - } - } - } - }; + .ok_or(crate::error::ApiError::NotFound(format!( + "Attachment of group {group_name} and key {key}" + )))? + .id; + ActiveTask::Entity::find() + .select_only() + .column(ActiveTask::Column::Id) + .filter(ActiveTask::Column::GroupId.eq(group_id)) + .filter(ActiveTask::Column::RunnerUuid.eq(worker_uuid)) + .filter(ActiveTask::Column::State.eq(TaskState::Running)) + .into_tuple::() + .one(&pool.db) + .await? + .ok_or(AuthError::PermissionDenied)?; let attachment = Attachment::Entity::find() .filter(Attachment::Column::GroupId.eq(group_id)) .filter(Attachment::Column::Key.eq(key.clone())) diff --git a/netmito/src/service/worker/mod.rs b/netmito/src/service/worker/mod.rs index 4c80f629..48298aa6 100644 --- a/netmito/src/service/worker/mod.rs +++ b/netmito/src/service/worker/mod.rs @@ -553,10 +553,21 @@ pub async fn fetch_task( let spec: ExecSpec = serde_json::from_value(task.spec)?; let exec_options: Option = task.exec_options.map(serde_json::from_value).transpose()?; + let group_name = Group::Entity::find_by_id(task.group_id) + .select_only() + .column(Group::Column::GroupName) + .into_tuple::() + .one(&pool.db) + .await? + .ok_or(Error::Custom(format!( + "group of task {} not found", + task.uuid + )))?; let task_resp = WorkerTaskResp { id: task.id, uuid: task.uuid, upstream_task_uuid: task.upstream_task_uuid, + group_name, spec, exec_options, }; diff --git a/netmito/src/worker.rs b/netmito/src/worker.rs index e70cabbc..301b7447 100644 --- a/netmito/src/worker.rs +++ b/netmito/src/worker.rs @@ -137,7 +137,7 @@ impl TaskExecutor { } } self.task_url - .set_path(format!("workers/tasks/{uuid}").as_str()); + .set_path(format!("tasks/{uuid}").as_str()); let resp = self .task_client .get(self.task_url.as_str()) @@ -176,8 +176,7 @@ impl TaskExecutor { break; } } - self.task_url - .set_path(format!("workers/tasks/{uuid}").as_str()); + self.task_url.set_path(format!("tasks/{uuid}").as_str()); let resp = self .task_client .get(self.task_url.as_str()) @@ -617,14 +616,14 @@ async fn execute_task( let content_serde_val = serde_json::to_value(content_type)?; let content_serde_str = content_serde_val.as_str().unwrap_or("result"); task_executor.task_url.set_path(&format!( - "workers/tasks/{uuid}/artifacts/{content_serde_str}" + "tasks/{uuid}/download/artifacts/{content_serde_str}" )); } RemoteResource::Attachment { key } => { - let uuid = task.uuid; - task_executor - .task_url - .set_path(&format!("workers/tasks/{uuid}/attachments/{key}")); + task_executor.task_url.set_path(&format!( + "groups/{}/download/attachments/{key}", + task.group_name + )); } }; let resp = loop { @@ -1273,7 +1272,7 @@ async fn process_task_result( content_length, }, }; - task_executor.task_url.set_path("workers/tasks"); + task_executor.task_url.set_path("tasks/report"); let resp = loop { let resp = task_executor .task_client @@ -1480,7 +1479,7 @@ async fn report_task( task_executor: &mut TaskExecutor, req: ReportTaskReq, ) -> crate::error::Result<()> { - task_executor.task_url.set_path("workers/tasks"); + task_executor.task_url.set_path("tasks/report"); loop { let resp = task_executor .task_client diff --git a/openapi.yaml b/openapi.yaml index 8f218fb2..f73ba3e7 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -7,7 +7,9 @@ info: ## Authentication Most endpoints require JWT bearer token authentication. Get a token by calling `/login` with - valid credentials. + valid credentials. Each token carries a `role` claim (1 = user, 2 = worker, 3 = agent), and + each endpoint accepts only the roles listed in its `security`; a token without a role is + rejected. ## API Structure - `/auth` - User authentication status @@ -492,12 +494,15 @@ paths: /groups/{group_name}/download/attachments/{key}: get: summary: Download attachment - description: Get download URL for an attachment + description: | + Get download URL for an attachment. A user must be a member of the group; a + worker must be running a task of the group. operationId: downloadAttachment tags: - Groups security: - bearerAuth: [] + - workerAuth: [] parameters: - name: group_name in: path @@ -1043,118 +1048,6 @@ paths: "401": $ref: "#/components/responses/Unauthorized" - post: - summary: Report task status - description: Worker reports task execution status - operationId: reportTask - tags: - - Workers - security: - - workerAuth: [] - requestBody: - required: true - content: - application/json: - schema: - $ref: "#/components/schemas/ReportTaskReq" - responses: - "200": - description: Status reported successfully - content: - application/json: - schema: - $ref: "#/components/schemas/ReportTaskResp" - "401": - $ref: "#/components/responses/Unauthorized" - - /workers/tasks/{uuid}: - get: - summary: Get task details for worker - description: Worker queries task details - operationId: getTaskForWorker - tags: - - Workers - security: - - workerAuth: [] - parameters: - - name: uuid - in: path - required: true - schema: - type: string - format: uuid - responses: - "200": - description: Task details - content: - application/json: - schema: - $ref: "#/components/schemas/WorkerTaskResp" - "401": - $ref: "#/components/responses/Unauthorized" - "404": - $ref: "#/components/responses/NotFound" - - /workers/tasks/{uuid}/artifacts/{content_type}: - get: - summary: Download task artifact for worker - description: Worker downloads task artifacts - operationId: downloadTaskArtifactForWorker - tags: - - Workers - security: - - workerAuth: [] - parameters: - - name: uuid - in: path - required: true - schema: - type: string - format: uuid - - name: content_type - in: path - required: true - schema: - $ref: "#/components/schemas/ArtifactContentType" - responses: - "302": - description: Redirect to download URL - "401": - $ref: "#/components/responses/Unauthorized" - "404": - $ref: "#/components/responses/NotFound" - - /workers/tasks/{uuid}/attachments/{key}: - get: - summary: Download attachment for worker - description: Worker downloads group attachments needed for task - operationId: downloadAttachmentForWorker - tags: - - Workers - security: - - workerAuth: [] - parameters: - - name: uuid - in: path - required: true - schema: - type: string - format: uuid - - name: key - in: path - required: true - schema: - type: string - style: simple - explode: false - responses: - "302": - description: Redirect to download URL - "401": - $ref: "#/components/responses/Unauthorized" - "404": - $ref: "#/components/responses/NotFound" - /tasks: post: summary: Submit task @@ -1407,6 +1300,31 @@ paths: "403": $ref: "#/components/responses/Forbidden" + /tasks/report: + post: + summary: Report task status + description: Worker reports task execution status + operationId: reportTask + tags: + - Tasks + security: + - workerAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/ReportTaskReq" + responses: + "200": + description: Status reported successfully + content: + application/json: + schema: + $ref: "#/components/schemas/ReportTaskResp" + "401": + $ref: "#/components/responses/Unauthorized" + /tasks/{uuid}: get: summary: Get task information @@ -1416,6 +1334,7 @@ paths: - Tasks security: - bearerAuth: [] + - workerAuth: [] parameters: - name: uuid in: path @@ -1605,6 +1524,7 @@ paths: - Tasks security: - bearerAuth: [] + - workerAuth: [] parameters: - name: uuid in: path @@ -3643,6 +3563,7 @@ components: required: - id - uuid + - group_name - timeout - spec properties: @@ -3654,6 +3575,9 @@ components: type: string format: uuid description: Task UUID + group_name: + type: string + description: The task's group, where its attachment resources are fetched from timeout: type: string description: Task timeout duration From 51d39ae2bb0ebdc09777abccadc6033be0c45526 Mon Sep 17 00:00:00 2001 From: un-lock-able Date: Sat, 24 Oct 2026 22:57:34 +0800 Subject: [PATCH 2/2] refactor(auth): derive role bits from TokenRole discriminants --- netmito/src/service/auth/mod.rs | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/netmito/src/service/auth/mod.rs b/netmito/src/service/auth/mod.rs index 2c10a3cc..e73e6586 100644 --- a/netmito/src/service/auth/mod.rs +++ b/netmito/src/service/auth/mod.rs @@ -28,20 +28,18 @@ use token::{generate_token, verify_token, TokenClaims, TokenRole}; // refactor if necessary. // TODO: check if we allow non-expiry lifetime -/// Role bits for [`auth_middleware`]'s `ROLES`. +/// Role bits for [`auth_middleware`]'s `ROLES`, so one endpoint can accept several roles by using +/// something like `USER | WORKER`. These are bit flags, not [`TokenRole`] discriminants; +/// each bit is derived from its discriminant (which starts at 1). pub mod roles { use super::token::TokenRole; - pub const USER: u8 = 1 << 0; - pub const WORKER: u8 = 1 << 1; - pub const AGENT: u8 = 1 << 2; + pub const USER: u8 = 1 << (TokenRole::User as u8 - 1); + pub const WORKER: u8 = 1 << (TokenRole::Worker as u8 - 1); + pub const AGENT: u8 = 1 << (TokenRole::Agent as u8 - 1); pub(super) fn bit_mask_of(role: TokenRole) -> u8 { - match role { - TokenRole::User => USER, - TokenRole::Worker => WORKER, - TokenRole::Agent => AGENT, - } + 1 << (role as u8 - 1) } }