microsandbox-server 0.2.6

`microsandbox-server` implements the sandbox server responsible for orchestrating sandboxes.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373

//! Middleware components for the microsandbox server.
//!
//! This module handles:
//! - Request/response middleware
//! - Authentication and authorization
//! - Request tracing and logging
//! - Error handling
//!
//! The module provides:
//! - Middleware components for common operations
//! - Authentication middleware for API security
//! - Logging and tracing middleware

use axum::{
    body::{to_bytes, Body},
    extract::State,
    http::{HeaderMap, Request, StatusCode, Uri},
    middleware::Next,
    response::IntoResponse,
};
use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation};
use serde_json::Value;

use crate::{
    config::PROXY_AUTH_HEADER,
    error::{AuthenticationError, ServerError, ValidationError},
    management::API_KEY_PREFIX,
    state::AppState,
    Claims,
};

//--------------------------------------------------------------------------------------------------
// Middleware Functions
//--------------------------------------------------------------------------------------------------

/// Proxy middleware for forwarding requests to a target service
pub async fn proxy_middleware(
    State(_state): State<AppState>,
    req: Request<Body>,
    next: Next,
) -> impl IntoResponse {
    // Default to passing the request to the next handler
    // This middleware can be extended to implement actual proxying logic
    next.run(req).await
}

/// Convert a URI to a proxied URI targeting a sandbox
pub fn proxy_uri(original_uri: Uri, namespace: &str, sandbox_name: &str) -> Uri {
    // In a real implementation, you would:
    // 1. Look up the sandbox's address from a registry or state
    // 2. Construct a new URI that points to the sandbox
    // 3. Return the new URI for proxying

    // For demonstration purposes, we'll construct a simple URI
    // In production, you would get this from a sandbox registry
    let target_host = format!("sandbox-{}.{}.internal", sandbox_name, namespace);

    let uri_string = if let Some(path_and_query) = original_uri.path_and_query() {
        format!("http://{}:{}{}", target_host, 8080, path_and_query)
    } else {
        format!("http://{}:{}/", target_host, 8080)
    };

    // Try to parse the string into a URI
    // In case of errors, fallback to a default URI
    uri_string
        .parse()
        .unwrap_or_else(|_| "http://localhost:8080/".parse().unwrap())
}

/// Log incoming requests
pub async fn logging_middleware(
    req: Request<Body>,
    next: Next,
) -> Result<impl IntoResponse, (StatusCode, String)> {
    let method = req.method().clone();
    let uri = req.uri().clone();

    // Log the request
    tracing::info!("Request: {} {}", method, uri);

    // Process the request
    let response = next.run(req).await;

    // Log the response
    tracing::info!("Response: {} {}: {}", method, uri, response.status());

    Ok(response)
}

/// Authentication middleware for verifying API keys and namespace access
pub async fn auth_middleware(
    State(state): State<AppState>,
    req: Request<Body>,
    next: Next,
) -> Result<impl IntoResponse, ServerError> {
    // Skip auth in dev mode if configured
    if *state.get_config().get_dev_mode() {
        return Ok(next.run(req).await);
    }

    // Extract API key from authorization header
    let api_key = extract_api_key_from_headers(req.headers())?;

    // Validate the token and get its claims
    let claims = validate_token(&api_key, &state)?;

    // If token has wildcard namespace access, we can skip further namespace validation
    if claims.namespace == "*" {
        return Ok(next.run(req).await);
    }

    // For namespace-specific tokens, we need to ensure the token has access to the requested namespace
    // We need to read the request body to extract the namespace
    let (parts, body) = req.into_parts();

    // Use axum's to_bytes to buffer the body
    let bytes = to_bytes(body, usize::MAX)
        .await
        .map_err(|e| ServerError::InternalError(format!("Failed to read request body: {}", e)))?;

    // Parse the JSON-RPC request and extract the namespace
    let namespace_from_request = extract_namespace_from_json_rpc(&bytes)?;

    // Validate that the token has access to the requested namespace
    if claims.namespace != namespace_from_request {
        return Err(ServerError::AuthorizationError(
            crate::error::AuthorizationError::AccessDenied(format!(
                "Token does not have access to namespace '{}'",
                namespace_from_request
            )),
        ));
    }

    // Reconstruct the request with the original body
    let body = Body::from(bytes);
    let req = Request::from_parts(parts, body);

    // If everything is valid, continue with the request
    Ok(next.run(req).await)
}

/// Smart authentication middleware for MCP requests that handles protocol vs tool methods differently
/// Protocol methods (initialize, tools/list, prompts/list, prompts/get) don't require namespace validation
/// Tool methods (tools/call) require namespace validation
pub async fn mcp_smart_auth_middleware(
    State(state): State<AppState>,
    req: Request<Body>,
    next: Next,
) -> Result<impl IntoResponse, ServerError> {
    // Skip auth in dev mode if configured
    if *state.get_config().get_dev_mode() {
        return Ok(next.run(req).await);
    }

    // Extract API key from authorization header
    let api_key = extract_api_key_from_headers(req.headers())?;

    // Validate the token and get its claims
    let claims = validate_token(&api_key, &state)?;

    // If token has wildcard namespace access, we can skip further namespace validation
    if claims.namespace == "*" {
        return Ok(next.run(req).await);
    }

    // For namespace-specific tokens, we need to check if this is a tool execution method
    // that requires namespace validation
    let (parts, body) = req.into_parts();

    // Use axum's to_bytes to buffer the body
    let bytes = to_bytes(body, usize::MAX)
        .await
        .map_err(|e| ServerError::InternalError(format!("Failed to read request body: {}", e)))?;

    // Parse the JSON to check the method
    let json_value: serde_json::Value = serde_json::from_slice(&bytes).map_err(|e| {
        ServerError::ValidationError(crate::error::ValidationError::InvalidInput(format!(
            "Invalid JSON-RPC request: {}",
            e
        )))
    })?;

    let method = json_value
        .get("method")
        .and_then(serde_json::Value::as_str)
        .unwrap_or("unknown");

    // Check if this is a tool execution method that requires namespace validation
    let requires_namespace_validation = matches!(method, "tools/call");

    if requires_namespace_validation {
        // Extract namespace from params for tool execution methods
        let namespace_from_request = extract_namespace_from_json_rpc(&bytes)?;

        // Validate that the token has access to the requested namespace
        if claims.namespace != namespace_from_request {
            return Err(ServerError::AuthorizationError(
                crate::error::AuthorizationError::AccessDenied(format!(
                    "Token does not have access to namespace '{}'",
                    namespace_from_request
                )),
            ));
        }
    }

    // Reconstruct the request with the original body
    let body = Body::from(bytes);
    let req = Request::from_parts(parts, body);

    // If everything is valid, continue with the request
    Ok(next.run(req).await)
}

//--------------------------------------------------------------------------------------------------
// Helper Functions
//--------------------------------------------------------------------------------------------------

/// Extract the namespace from a JSON-RPC request body
fn extract_namespace_from_json_rpc(bytes: &[u8]) -> Result<String, ServerError> {
    // Parse the request body as JSON
    let json_value: Value = serde_json::from_slice(bytes).map_err(|e| {
        ServerError::ValidationError(ValidationError::InvalidInput(format!(
            "Invalid JSON-RPC request: {}",
            e
        )))
    })?;

    // Extract the method for logging purposes
    let method = json_value
        .get("method")
        .and_then(Value::as_str)
        .unwrap_or("unknown");

    // Extract params object
    let params = json_value.get("params").ok_or_else(|| {
        ServerError::ValidationError(ValidationError::InvalidInput(
            "Missing 'params' field in JSON-RPC request".to_string(),
        ))
    })?;

    // Extract namespace from params for any method
    params
        .get("namespace")
        .and_then(Value::as_str)
        .map(String::from)
        .ok_or_else(|| {
            ServerError::ValidationError(ValidationError::InvalidInput(format!(
                "Missing or invalid 'namespace' in params for method '{}'",
                method
            )))
        })
}

/// Extract API key from request headers
fn extract_api_key_from_headers(headers: &HeaderMap) -> Result<String, ServerError> {
    // First check the Proxy-Authorization header
    if let Some(auth_header) = headers.get(PROXY_AUTH_HEADER) {
        let auth_value = auth_header.to_str().map_err(|_| {
            ServerError::Authentication(AuthenticationError::InvalidCredentials(
                "Invalid authorization header format".to_string(),
            ))
        })?;

        // Check if it has the Bearer prefix
        if let Some(token) = auth_value.strip_prefix("Bearer ") {
            return Ok(token.to_string());
        }

        // Or if it's just the raw token
        return Ok(auth_value.to_string());
    }

    // Then check standard Authorization header
    if let Some(auth_header) = headers.get("Authorization") {
        let auth_value = auth_header.to_str().map_err(|_| {
            ServerError::Authentication(AuthenticationError::InvalidCredentials(
                "Invalid authorization header format".to_string(),
            ))
        })?;

        // Check if it has the Bearer prefix
        if let Some(token) = auth_value.strip_prefix("Bearer ") {
            return Ok(token.to_string());
        }

        // Or if it's just the raw token
        return Ok(auth_value.to_string());
    }

    Err(ServerError::Authentication(
        AuthenticationError::InvalidCredentials("Missing authorization header".to_string()),
    ))
}

/// Convert a custom API key back to a standard JWT format
fn convert_api_key_to_jwt(api_key: &str) -> Result<String, ServerError> {
    // Check if the API key has the expected prefix
    if !api_key.starts_with(API_KEY_PREFIX) {
        return Err(ServerError::Authentication(
            AuthenticationError::InvalidCredentials(
                "Invalid API key format: missing prefix".to_string(),
            ),
        ));
    }

    // Remove the prefix and return the original JWT
    Ok(api_key[API_KEY_PREFIX.len()..].to_string())
}

/// Get the server key from the AppState config
fn get_server_key(state: &AppState) -> Result<String, ServerError> {
    // Get the key from the config - we already know we're not in dev mode
    // by the time this function is called
    match state.get_config().get_key() {
        Some(key) => Ok(key.clone()),
        None => Err(ServerError::Authentication(
            AuthenticationError::InvalidCredentials(
                "Server key not found in configuration".to_string(),
            ),
        )),
    }
}

/// Validate the token
fn validate_token(api_key: &str, state: &AppState) -> Result<Claims, ServerError> {
    // Convert API key back to JWT format
    let jwt = convert_api_key_to_jwt(api_key)?;

    // Get server key for validation
    let server_key = get_server_key(state)?;

    // Decode and validate the JWT
    let token_data = decode::<Claims>(
        &jwt,
        &DecodingKey::from_secret(server_key.as_bytes()),
        &Validation::new(Algorithm::HS256),
    )
    .map_err(|e| {
        let error_message = match e.kind() {
            jsonwebtoken::errors::ErrorKind::ExpiredSignature => "Token expired".to_string(),
            jsonwebtoken::errors::ErrorKind::InvalidSignature => {
                "Invalid token signature".to_string()
            }
            _ => format!("Token validation error: {}", e),
        };
        ServerError::Authentication(AuthenticationError::InvalidToken(error_message))
    })?;

    Ok(token_data.claims)
}

/// Validate the token and check namespace access
pub fn validate_token_and_namespace(
    api_key: &str,
    requested_namespace: &str,
    state: &AppState,
) -> Result<Claims, ServerError> {
    // Validate token
    let claims = validate_token(api_key, state)?;

    // Check if the token's namespace matches the requested namespace
    if claims.namespace != requested_namespace && claims.namespace != "*" {
        return Err(ServerError::Authentication(
            AuthenticationError::InvalidCredentials(format!(
                "Token does not have access to namespace '{}'",
                requested_namespace
            )),
        ));
    }

    Ok(claims)
}