Extract MCP machine access module
This commit is contained in:
@@ -0,0 +1,140 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::http::{HeaderMap, StatusCode, header::AUTHORIZATION};
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use crank_core::{OperationSecurityLevel, PlatformApiKeyScope};
|
||||
use sha2::{Digest, Sha256};
|
||||
use time::OffsetDateTime;
|
||||
|
||||
use crate::{
|
||||
app::{AgentRoutePath, AppState},
|
||||
auth::VerifiedMachineCredential,
|
||||
};
|
||||
|
||||
pub(super) async fn require_machine_access(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
headers: &HeaderMap,
|
||||
required_scope: PlatformApiKeyScope,
|
||||
) -> Result<VerifiedMachineCredential, StatusCode> {
|
||||
let secret = bearer_token(headers).ok_or(StatusCode::UNAUTHORIZED)?;
|
||||
let credential = resolve_machine_credential(state, path, secret).await?;
|
||||
|
||||
if !allows_scope(&credential.scopes, required_scope) {
|
||||
return Err(StatusCode::FORBIDDEN);
|
||||
}
|
||||
|
||||
Ok(credential)
|
||||
}
|
||||
|
||||
pub(super) fn bearer_token(headers: &HeaderMap) -> Option<&str> {
|
||||
let value = headers.get(AUTHORIZATION)?.to_str().ok()?;
|
||||
let (scheme, token) = value.split_once(' ')?;
|
||||
if !scheme.eq_ignore_ascii_case("Bearer") || token.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(token)
|
||||
}
|
||||
|
||||
pub(super) fn credential_allows_security_level(
|
||||
credential: &VerifiedMachineCredential,
|
||||
required_level: OperationSecurityLevel,
|
||||
) -> bool {
|
||||
security_level_rank(credential.max_security_level) >= security_level_rank(required_level)
|
||||
}
|
||||
|
||||
pub(super) fn serialize_security_level(level: OperationSecurityLevel) -> &'static str {
|
||||
match level {
|
||||
OperationSecurityLevel::Standard => "standard",
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn serialize_machine_access_mode(mode: crank_core::MachineAccessMode) -> &'static str {
|
||||
match mode {
|
||||
crank_core::MachineAccessMode::StaticAgentKey => "static_agent_key",
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn hash_access_secret(secret: &str) -> String {
|
||||
let digest = Sha256::digest(secret.as_bytes());
|
||||
URL_SAFE_NO_PAD.encode(digest)
|
||||
}
|
||||
|
||||
async fn resolve_machine_credential(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
token: &str,
|
||||
) -> Result<VerifiedMachineCredential, StatusCode> {
|
||||
if let Some(credential) = verify_static_agent_key(state, path, token).await? {
|
||||
return Ok(credential);
|
||||
}
|
||||
|
||||
state
|
||||
.credential_verifier
|
||||
.verify_bearer_token(&path.workspace_slug, &path.agent_slug, token)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
|
||||
.ok_or(StatusCode::UNAUTHORIZED)
|
||||
}
|
||||
|
||||
async fn verify_static_agent_key(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
secret: &str,
|
||||
) -> Result<Option<VerifiedMachineCredential>, StatusCode> {
|
||||
let secret_hash = hash_access_secret(secret);
|
||||
let Some(api_key) = state
|
||||
.registry
|
||||
.get_platform_api_key_by_secret_for_agent_slug(
|
||||
&path.workspace_slug,
|
||||
&path.agent_slug,
|
||||
&secret_hash,
|
||||
)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let used_at = OffsetDateTime::now_utc();
|
||||
state
|
||||
.registry
|
||||
.touch_platform_api_key(&api_key.api_key.workspace_id, &api_key.api_key.id, &used_at)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
|
||||
Ok(Some(VerifiedMachineCredential {
|
||||
machine_access_mode: crank_core::MachineAccessMode::StaticAgentKey,
|
||||
max_security_level: crank_core::OperationSecurityLevel::Standard,
|
||||
scopes: api_key.api_key.scopes,
|
||||
}))
|
||||
}
|
||||
|
||||
fn allows_scope(scopes: &[PlatformApiKeyScope], required_scope: PlatformApiKeyScope) -> bool {
|
||||
match required_scope {
|
||||
PlatformApiKeyScope::Read => scopes.iter().any(|scope| {
|
||||
matches!(
|
||||
scope,
|
||||
PlatformApiKeyScope::Read
|
||||
| PlatformApiKeyScope::Write
|
||||
| PlatformApiKeyScope::Deploy
|
||||
)
|
||||
}),
|
||||
PlatformApiKeyScope::Write => scopes.iter().any(|scope| {
|
||||
matches!(
|
||||
scope,
|
||||
PlatformApiKeyScope::Write | PlatformApiKeyScope::Deploy
|
||||
)
|
||||
}),
|
||||
PlatformApiKeyScope::Deploy => scopes
|
||||
.iter()
|
||||
.any(|scope| matches!(scope, PlatformApiKeyScope::Deploy)),
|
||||
}
|
||||
}
|
||||
|
||||
fn security_level_rank(level: OperationSecurityLevel) -> u8 {
|
||||
match level {
|
||||
OperationSecurityLevel::Standard => 0,
|
||||
}
|
||||
}
|
||||
@@ -8,17 +8,13 @@ use std::{
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::{Path, State},
|
||||
http::{
|
||||
HeaderMap, HeaderValue, StatusCode,
|
||||
header::{AUTHORIZATION, RETRY_AFTER},
|
||||
},
|
||||
http::{HeaderMap, HeaderValue, StatusCode, header::RETRY_AFTER},
|
||||
response::{IntoResponse, Response, sse::Event},
|
||||
routing::get,
|
||||
};
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use crank_core::{
|
||||
AuthProfile, CoordinationStateStore, InvocationLevel, InvocationLog, InvocationLogId,
|
||||
InvocationSource, InvocationStatus, OperationSecurityLevel, PlatformApiKeyScope, SecretId,
|
||||
InvocationSource, InvocationStatus, PlatformApiKeyScope, SecretId,
|
||||
};
|
||||
use crank_registry::{CreateInvocationLogRequest, PostgresRegistry, PublishedAgentTool};
|
||||
use crank_runtime::{
|
||||
@@ -28,11 +24,14 @@ use crank_runtime::{
|
||||
use futures_util::stream;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Value, json};
|
||||
use sha2::{Digest, Sha256};
|
||||
use time::OffsetDateTime;
|
||||
use tracing::info;
|
||||
|
||||
use crate::{
|
||||
access::{
|
||||
bearer_token, credential_allows_security_level, hash_access_secret, require_machine_access,
|
||||
serialize_machine_access_mode, serialize_security_level,
|
||||
},
|
||||
auth::{SharedMachineCredentialVerifier, VerifiedMachineCredential},
|
||||
catalog::PublishedToolCatalog,
|
||||
jsonrpc::{
|
||||
@@ -56,14 +55,14 @@ use crate::{
|
||||
const TRANSPORT_SESSION_TTL_MS: u64 = 86_400_000;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
registry: PostgresRegistry,
|
||||
pub(super) struct AppState {
|
||||
pub(super) registry: PostgresRegistry,
|
||||
catalog: PublishedToolCatalog,
|
||||
runtime: RuntimeExecutor,
|
||||
api_rate_limiter: RequestRateLimiter,
|
||||
secret_crypto: SecretCrypto,
|
||||
sessions: SharedSessionStore,
|
||||
credential_verifier: SharedMachineCredentialVerifier,
|
||||
pub(super) credential_verifier: SharedMachineCredentialVerifier,
|
||||
allowed_origins: AllowedOrigins,
|
||||
}
|
||||
|
||||
@@ -92,9 +91,9 @@ struct ToolCallExecution {
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
struct AgentRoutePath {
|
||||
workspace_slug: String,
|
||||
agent_slug: String,
|
||||
pub(super) struct AgentRoutePath {
|
||||
pub(super) workspace_slug: String,
|
||||
pub(super) agent_slug: String,
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
@@ -811,82 +810,6 @@ async fn require_initialized_session(
|
||||
Ok(session)
|
||||
}
|
||||
|
||||
async fn require_machine_access(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
headers: &HeaderMap,
|
||||
required_scope: PlatformApiKeyScope,
|
||||
) -> Result<VerifiedMachineCredential, StatusCode> {
|
||||
let secret = bearer_token(headers).ok_or(StatusCode::UNAUTHORIZED)?;
|
||||
let credential = resolve_machine_credential(state, path, secret).await?;
|
||||
|
||||
if !allows_scope(&credential.scopes, required_scope) {
|
||||
return Err(StatusCode::FORBIDDEN);
|
||||
}
|
||||
|
||||
Ok(credential)
|
||||
}
|
||||
|
||||
async fn resolve_machine_credential(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
token: &str,
|
||||
) -> Result<VerifiedMachineCredential, StatusCode> {
|
||||
if let Some(credential) = verify_static_agent_key(state, path, token).await? {
|
||||
return Ok(credential);
|
||||
}
|
||||
|
||||
state
|
||||
.credential_verifier
|
||||
.verify_bearer_token(&path.workspace_slug, &path.agent_slug, token)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
|
||||
.ok_or(StatusCode::UNAUTHORIZED)
|
||||
}
|
||||
|
||||
async fn verify_static_agent_key(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
secret: &str,
|
||||
) -> Result<Option<VerifiedMachineCredential>, StatusCode> {
|
||||
let secret_hash = hash_access_secret(secret);
|
||||
let Some(api_key) = state
|
||||
.registry
|
||||
.get_platform_api_key_by_secret_for_agent_slug(
|
||||
&path.workspace_slug,
|
||||
&path.agent_slug,
|
||||
&secret_hash,
|
||||
)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let used_at = OffsetDateTime::now_utc();
|
||||
state
|
||||
.registry
|
||||
.touch_platform_api_key(&api_key.api_key.workspace_id, &api_key.api_key.id, &used_at)
|
||||
.await
|
||||
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
|
||||
Ok(Some(VerifiedMachineCredential {
|
||||
machine_access_mode: crank_core::MachineAccessMode::StaticAgentKey,
|
||||
max_security_level: crank_core::OperationSecurityLevel::Standard,
|
||||
scopes: api_key.api_key.scopes,
|
||||
}))
|
||||
}
|
||||
|
||||
fn bearer_token(headers: &HeaderMap) -> Option<&str> {
|
||||
let value = headers.get(AUTHORIZATION)?.to_str().ok()?;
|
||||
let (scheme, token) = value.split_once(' ')?;
|
||||
if !scheme.eq_ignore_ascii_case("Bearer") || token.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(token)
|
||||
}
|
||||
|
||||
async fn enforce_post_rate_limit(
|
||||
state: &Arc<AppState>,
|
||||
path: &AgentRoutePath,
|
||||
@@ -961,53 +884,6 @@ fn rate_limited_status_response(rejection: RateLimitRejection) -> Response {
|
||||
response
|
||||
}
|
||||
|
||||
fn allows_scope(scopes: &[PlatformApiKeyScope], required_scope: PlatformApiKeyScope) -> bool {
|
||||
match required_scope {
|
||||
PlatformApiKeyScope::Read => scopes.iter().any(|scope| {
|
||||
matches!(
|
||||
scope,
|
||||
PlatformApiKeyScope::Read
|
||||
| PlatformApiKeyScope::Write
|
||||
| PlatformApiKeyScope::Deploy
|
||||
)
|
||||
}),
|
||||
PlatformApiKeyScope::Write => scopes.iter().any(|scope| {
|
||||
matches!(
|
||||
scope,
|
||||
PlatformApiKeyScope::Write | PlatformApiKeyScope::Deploy
|
||||
)
|
||||
}),
|
||||
PlatformApiKeyScope::Deploy => scopes
|
||||
.iter()
|
||||
.any(|scope| matches!(scope, PlatformApiKeyScope::Deploy)),
|
||||
}
|
||||
}
|
||||
|
||||
fn credential_allows_security_level(
|
||||
credential: &VerifiedMachineCredential,
|
||||
required_level: OperationSecurityLevel,
|
||||
) -> bool {
|
||||
security_level_rank(credential.max_security_level) >= security_level_rank(required_level)
|
||||
}
|
||||
|
||||
fn security_level_rank(level: OperationSecurityLevel) -> u8 {
|
||||
match level {
|
||||
OperationSecurityLevel::Standard => 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn serialize_security_level(level: OperationSecurityLevel) -> &'static str {
|
||||
match level {
|
||||
OperationSecurityLevel::Standard => "standard",
|
||||
}
|
||||
}
|
||||
|
||||
fn serialize_machine_access_mode(mode: crank_core::MachineAccessMode) -> &'static str {
|
||||
match mode {
|
||||
crank_core::MachineAccessMode::StaticAgentKey => "static_agent_key",
|
||||
}
|
||||
}
|
||||
|
||||
fn internal_jsonrpc_error(message: &Value, error: impl std::fmt::Display) -> Response {
|
||||
transport_response(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
@@ -1154,11 +1030,6 @@ fn add_millis(timestamp: OffsetDateTime, millis: u64) -> OffsetDateTime {
|
||||
timestamp + delta
|
||||
}
|
||||
|
||||
fn hash_access_secret(secret: &str) -> String {
|
||||
let digest = Sha256::digest(secret.as_bytes());
|
||||
URL_SAFE_NO_PAD.encode(digest)
|
||||
}
|
||||
|
||||
fn resolve_generated_tool(
|
||||
tools: &[PublishedAgentTool],
|
||||
tool_name: &str,
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
mod access;
|
||||
mod app;
|
||||
pub mod auth;
|
||||
pub mod catalog;
|
||||
|
||||
Reference in New Issue
Block a user