Skip to main content

statespace_server/
server.rs

1//! HTTP server and Axum router.
2
3use crate::content::{ContentResolver, LocalContentResolver};
4use crate::error::ErrorExt;
5use axum::{
6    Json, Router,
7    extract::{Path, Query, State, rejection::JsonRejection},
8    http::{StatusCode, header},
9    response::{IntoResponse, Response},
10    routing::get,
11};
12use statespace_tool_runtime::{
13    ActionRequest, ActionResponse, BuiltinTool, ErrorResponse, ExecutionLimits, SandboxEnv,
14    SuccessResponse, ToolExecutor, eval, expand_command_for_execution, parse_frontmatter,
15    validate_command_with_specs, validate_env_map,
16};
17use std::collections::HashMap;
18use std::path::PathBuf;
19use std::sync::Arc;
20use tokio::fs;
21use tower_http::trace::TraceLayer;
22use tracing::Span;
23
24#[derive(Clone)]
25pub struct ServerConfig {
26    pub content_root: PathBuf,
27    pub host: String,
28    pub port: u16,
29    pub limits: ExecutionLimits,
30    pub env: HashMap<String, String>,
31    pub sandbox_env: SandboxEnv,
32}
33
34impl std::fmt::Debug for ServerConfig {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        f.debug_struct("ServerConfig")
37            .field("content_root", &self.content_root)
38            .field("host", &self.host)
39            .field("port", &self.port)
40            .field("limits", &self.limits)
41            .field("env_keys", &self.env.len())
42            .field("sandbox_path", &self.sandbox_env.path())
43            .finish()
44    }
45}
46
47impl ServerConfig {
48    #[must_use]
49    pub fn new(content_root: PathBuf) -> Self {
50        Self {
51            content_root,
52            host: "127.0.0.1".to_string(),
53            port: 8000,
54            limits: ExecutionLimits::default(),
55            env: HashMap::new(),
56            sandbox_env: SandboxEnv::default(),
57        }
58    }
59
60    #[must_use]
61    pub fn with_host(mut self, host: impl Into<String>) -> Self {
62        self.host = host.into();
63        self
64    }
65
66    #[must_use]
67    pub const fn with_port(mut self, port: u16) -> Self {
68        self.port = port;
69        self
70    }
71
72    #[must_use]
73    pub fn with_limits(mut self, limits: ExecutionLimits) -> Self {
74        self.limits = limits;
75        self
76    }
77
78    #[must_use]
79    pub fn with_env(mut self, env: HashMap<String, String>) -> Self {
80        self.env = env;
81        self
82    }
83
84    #[must_use]
85    pub fn with_sandbox_env(mut self, sandbox_env: SandboxEnv) -> Self {
86        self.sandbox_env = sandbox_env;
87        self
88    }
89
90    #[must_use]
91    pub fn socket_addr(&self) -> String {
92        format!("{}:{}", self.host, self.port)
93    }
94
95    #[must_use]
96    pub fn base_url(&self) -> String {
97        format!("http://{}:{}", self.host, self.port)
98    }
99}
100
101#[derive(Clone)]
102pub struct ServerState {
103    pub content_resolver: Arc<dyn ContentResolver>,
104    pub limits: ExecutionLimits,
105    pub content_root: PathBuf,
106    pub env: Arc<HashMap<String, String>>,
107    pub sandbox_env: Arc<SandboxEnv>,
108}
109
110impl std::fmt::Debug for ServerState {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        f.debug_struct("ServerState")
113            .field("limits", &self.limits)
114            .field("content_root", &self.content_root)
115            .field("env_keys", &self.env.len())
116            .field("sandbox_path", &self.sandbox_env.path())
117            .finish_non_exhaustive()
118    }
119}
120
121impl ServerState {
122    /// # Errors
123    ///
124    /// Returns an error if the content root path cannot be canonicalized.
125    pub fn from_config(config: &ServerConfig) -> crate::error::Result<Self> {
126        Ok(Self {
127            content_resolver: Arc::new(LocalContentResolver::new(&config.content_root)?),
128            limits: config.limits.clone(),
129            content_root: config.content_root.clone(),
130            env: Arc::new(config.env.clone()),
131            sandbox_env: Arc::new(config.sandbox_env.clone()),
132        })
133    }
134}
135
136/// # Errors
137///
138/// Returns an error if the content root path cannot be canonicalized.
139pub fn build_router(config: &ServerConfig) -> crate::error::Result<Router> {
140    let state = ServerState::from_config(config)?;
141
142    let trace_layer = TraceLayer::new_for_http()
143        .make_span_with(|request: &axum::http::Request<_>| {
144            tracing::info_span!(
145                "",
146                method = %request.method(),
147                path = %request.uri().path(),
148            )
149        })
150        .on_response(
151            |response: &axum::http::Response<_>, latency: std::time::Duration, _span: &Span| {
152                let status = response.status();
153                let code = status.as_u16();
154                let reason = status.canonical_reason().unwrap_or("");
155                let ms = latency.as_secs_f64() * 1000.0;
156
157                if code < 400 {
158                    tracing::info!("{code} {reason} {ms:.1}ms");
159                } else {
160                    tracing::error!("{code} {reason} {ms:.1}ms");
161                }
162            },
163        );
164
165    let router = Router::new()
166        .route("/", get(index_handler).post(action_handler_root))
167        .route("/favicon.svg", get(favicon_handler))
168        .route("/favicon.ico", get(favicon_handler))
169        .route("/{*path}", get(file_handler).post(action_handler))
170        .layer(trace_layer);
171
172    Ok(router.with_state(state))
173}
174async fn index_handler(
175    Query(query_env): Query<HashMap<String, String>>,
176    State(state): State<ServerState>,
177) -> Response {
178    serve_page("", &query_env, &state).await
179}
180
181async fn favicon_handler(State(state): State<ServerState>) -> Response {
182    match fs::read_to_string(state.content_root.join("favicon.svg")).await {
183        Ok(content) => (
184            StatusCode::OK,
185            [(header::CONTENT_TYPE, "image/svg+xml")],
186            content,
187        )
188            .into_response(),
189        Err(_) => StatusCode::NOT_FOUND.into_response(),
190    }
191}
192
193async fn file_handler(
194    Path(path): Path<String>,
195    Query(query_env): Query<HashMap<String, String>>,
196    State(state): State<ServerState>,
197) -> Response {
198    serve_page(&path, &query_env, &state).await
199}
200
201fn content_type_for_path(path: &std::path::Path) -> &'static str {
202    match path
203        .extension()
204        .and_then(|e| e.to_str())
205        .map(str::to_ascii_lowercase)
206        .as_deref()
207    {
208        Some("md") => "text/markdown; charset=utf-8",
209        Some("json") => "application/json; charset=utf-8",
210        Some("yaml" | "yml") => "text/yaml; charset=utf-8",
211        Some("csv") => "text/csv; charset=utf-8",
212        Some("html" | "htm") => "text/html; charset=utf-8",
213        _ => "text/plain; charset=utf-8",
214    }
215}
216
217async fn serve_page(
218    path: &str,
219    query_env: &HashMap<String, String>,
220    state: &ServerState,
221) -> Response {
222    if let Err(e) = validate_env_map(query_env) {
223        return json_error(StatusCode::BAD_REQUEST, &e.to_string());
224    }
225
226    let file_path = match state.content_resolver.resolve_path(path).await {
227        Ok(p) => p,
228        Err(e) => return json_error(e.status_code(), &e.user_message()),
229    };
230
231    let Ok(content) = fs::read_to_string(&file_path).await else {
232        return json_error(StatusCode::INTERNAL_SERVER_ERROR, "Internal server error");
233    };
234
235    let content_type = content_type_for_path(&file_path);
236    let working_dir = file_path.parent().unwrap_or(&state.content_root);
237    let has_eval = !eval::parse_eval_blocks(&content).is_empty();
238    let merged_env = eval::merge_eval_env(state.env.as_ref(), query_env);
239    let rendered = eval::process_eval_blocks_with_sandbox(
240        &content,
241        working_dir,
242        &merged_env,
243        &state.sandbox_env,
244        &state.limits,
245    )
246    .await;
247
248    if has_eval {
249        (
250            [
251                (header::CONTENT_TYPE, content_type),
252                (header::CACHE_CONTROL, "no-store"),
253                (header::X_CONTENT_TYPE_OPTIONS, "nosniff"),
254            ],
255            rendered,
256        )
257            .into_response()
258    } else {
259        (
260            [
261                (header::CONTENT_TYPE, content_type),
262                (header::X_CONTENT_TYPE_OPTIONS, "nosniff"),
263            ],
264            rendered,
265        )
266            .into_response()
267    }
268}
269
270async fn action_handler_root(
271    State(state): State<ServerState>,
272    body: Result<Json<ActionRequest>, JsonRejection>,
273) -> Response {
274    match body {
275        Ok(Json(request)) => execute_action("", &state, request).await,
276        Err(e) => json_error(e.status(), &e.body_text()),
277    }
278}
279
280async fn action_handler(
281    Path(path): Path<String>,
282    State(state): State<ServerState>,
283    body: Result<Json<ActionRequest>, JsonRejection>,
284) -> Response {
285    match body {
286        Ok(Json(request)) => execute_action(&path, &state, request).await,
287        Err(e) => json_error(e.status(), &e.body_text()),
288    }
289}
290
291fn runtime_error_response(e: &statespace_tool_runtime::Error) -> Response {
292    json_error(e.status_code(), &e.user_message())
293}
294
295async fn execute_action(path: &str, state: &ServerState, request: ActionRequest) -> Response {
296    if let Err(msg) = request.validate() {
297        return json_error(StatusCode::BAD_REQUEST, &msg);
298    }
299
300    let file_path = match state.content_resolver.resolve_path(path).await {
301        Ok(p) => p,
302        Err(e) => return runtime_error_response(&e),
303    };
304
305    if file_path.extension().and_then(|e| e.to_str()) != Some("md") {
306        return json_error(
307            StatusCode::BAD_REQUEST,
308            "POST is only supported on Markdown pages",
309        );
310    }
311
312    let content = match state.content_resolver.resolve(path).await {
313        Ok(c) => c,
314        Err(e) => return runtime_error_response(&e),
315    };
316
317    let frontmatter = match parse_frontmatter(&content) {
318        Ok(fm) => fm,
319        Err(e) => return runtime_error_response(&e),
320    };
321
322    let merged_env = eval::merge_eval_env(state.env.as_ref(), &request.env);
323    let expanded_command =
324        expand_command_for_execution(&request.command, &frontmatter.specs, &merged_env);
325
326    if let Err(e) = validate_command_with_specs(&frontmatter.specs, &request.command) {
327        return runtime_error_response(&e);
328    }
329
330    let tool = match BuiltinTool::from_command(&expanded_command) {
331        Ok(t) => t,
332        Err(e) => return runtime_error_response(&e),
333    };
334
335    let working_dir = file_path.parent().unwrap_or(&file_path);
336    let executor = ToolExecutor::new(working_dir.to_path_buf(), state.limits.clone())
337        .with_sandbox_env((*state.sandbox_env).clone())
338        .with_user_env(merged_env);
339
340    match executor.execute(&tool).await {
341        Ok(output) => {
342            let data = ActionResponse {
343                stdout: output.stdout(),
344                stderr: output.stderr().to_string(),
345                returncode: output.exit_code(),
346            };
347            let response = SuccessResponse::ok(data);
348            (StatusCode::OK, Json(response)).into_response()
349        }
350        Err(e) => runtime_error_response(&e),
351    }
352}
353
354fn json_error(status: StatusCode, message: &str) -> Response {
355    let response = ErrorResponse::new(message);
356    (status, Json(response)).into_response()
357}
358
359#[cfg(test)]
360#[allow(clippy::unwrap_used)]
361mod tests {
362    use super::*;
363    use axum::body;
364    use std::collections::HashMap;
365
366    async fn response_text(response: Response) -> String {
367        let bytes = body::to_bytes(response.into_body(), usize::MAX)
368            .await
369            .unwrap();
370        String::from_utf8_lossy(&bytes).to_string()
371    }
372
373    #[tokio::test]
374    async fn eval_pages_set_no_store_cache_control() {
375        let dir = tempfile::tempdir().unwrap();
376        std::fs::write(
377            dir.path().join("README.md"),
378            "```component\necho hello\n```\n",
379        )
380        .unwrap();
381
382        let config = ServerConfig::new(dir.path().to_path_buf());
383        let state = ServerState::from_config(&config).unwrap();
384
385        let response = serve_page("README.md", &HashMap::new(), &state).await;
386        assert_eq!(response.status(), StatusCode::OK);
387
388        let cache_control = response
389            .headers()
390            .get(header::CACHE_CONTROL)
391            .and_then(|v| v.to_str().ok());
392        assert_eq!(cache_control, Some("no-store"));
393    }
394
395    #[tokio::test]
396    async fn query_params_injected_into_component_env() {
397        let dir = tempfile::tempdir().unwrap();
398        std::fs::write(
399            dir.path().join("README.md"),
400            "```component\nprintf '%s/%s' \"$USER_ID\" \"$PAGE\"\n```\n",
401        )
402        .unwrap();
403
404        let config = ServerConfig::new(dir.path().to_path_buf());
405        let state = ServerState::from_config(&config).unwrap();
406        let query = HashMap::from([
407            ("USER_ID".to_string(), "42".to_string()),
408            ("PAGE".to_string(), "stats".to_string()),
409        ]);
410
411        let response = serve_page("README.md", &query, &state).await;
412        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
413            .await
414            .unwrap();
415        assert_eq!(String::from_utf8_lossy(&body).trim_end(), "42/stats");
416    }
417
418    #[tokio::test]
419    async fn configured_env_overrides_query_params() {
420        let dir = tempfile::tempdir().unwrap();
421        std::fs::write(
422            dir.path().join("README.md"),
423            "```component\necho \"$USER_ID\"\n```\n",
424        )
425        .unwrap();
426
427        let config = ServerConfig::new(dir.path().to_path_buf()).with_env(HashMap::from([(
428            "USER_ID".to_string(),
429            "trusted".to_string(),
430        )]));
431        let state = ServerState::from_config(&config).unwrap();
432        let query = HashMap::from([("USER_ID".to_string(), "untrusted".to_string())]);
433
434        let response = serve_page("README.md", &query, &state).await;
435        let body = axum::body::to_bytes(response.into_body(), usize::MAX)
436            .await
437            .unwrap();
438        assert_eq!(String::from_utf8_lossy(&body).trim_end(), "trusted");
439    }
440
441    #[tokio::test]
442    async fn invalid_query_key_returns_bad_request() {
443        let dir = tempfile::tempdir().unwrap();
444        std::fs::write(dir.path().join("README.md"), "ok\n").unwrap();
445
446        let config = ServerConfig::new(dir.path().to_path_buf());
447        let state = ServerState::from_config(&config).unwrap();
448        let query = HashMap::from([("A=B".to_string(), "1".to_string())]);
449
450        let response = serve_page("README.md", &query, &state).await;
451        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
452    }
453
454    #[tokio::test]
455    async fn invalid_query_value_returns_bad_request() {
456        let dir = tempfile::tempdir().unwrap();
457        std::fs::write(dir.path().join("README.md"), "ok\n").unwrap();
458
459        let config = ServerConfig::new(dir.path().to_path_buf());
460        let state = ServerState::from_config(&config).unwrap();
461        let query = HashMap::from([("USER_ID".to_string(), "abc\0def".to_string())]);
462
463        let response = serve_page("README.md", &query, &state).await;
464        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
465    }
466
467    #[tokio::test]
468    async fn action_expands_trusted_literal_env_segments() -> anyhow::Result<()> {
469        let dir = tempfile::tempdir()?;
470        std::fs::write(
471            dir.path().join("README.md"),
472            "---\ntools:\n  - [echo, $DATABASE_URL]\n---\n",
473        )?;
474
475        let config = ServerConfig::new(dir.path().to_path_buf())
476            .with_env(HashMap::from([(
477                "DATABASE_URL".to_string(),
478                "postgresql://gateway:gateway@localhost:5432/gateway_dev".to_string(),
479            )]))
480            .with_sandbox_env(SandboxEnv::from_host_process());
481        let state = ServerState::from_config(&config)?;
482
483        let request = ActionRequest {
484            command: vec!["echo".to_string(), "$DATABASE_URL".to_string()],
485
486            env: HashMap::new(),
487        };
488
489        let response = execute_action("README.md", &state, request).await;
490        assert_eq!(response.status(), StatusCode::OK);
491
492        let body = response_text(response).await;
493        let json: serde_json::Value = serde_json::from_str(&body)?;
494        let data = json
495            .get("data")
496            .ok_or_else(|| anyhow::anyhow!("missing data"))?;
497        let stdout = data
498            .get("stdout")
499            .and_then(|v| v.as_str())
500            .ok_or_else(|| anyhow::anyhow!("missing stdout"))?;
501        let stderr = data
502            .get("stderr")
503            .and_then(|v| v.as_str())
504            .ok_or_else(|| anyhow::anyhow!("missing stderr"))?;
505        let returncode = data
506            .get("returncode")
507            .and_then(serde_json::Value::as_i64)
508            .ok_or_else(|| anyhow::anyhow!("missing returncode"))?;
509        assert!(stdout.contains("postgresql://gateway:gateway@localhost:5432/gateway_dev"));
510        assert_eq!(stderr, "");
511        assert_eq!(returncode, 0);
512        Ok(())
513    }
514
515    #[tokio::test]
516    async fn action_does_not_expand_placeholders_into_trusted_env() -> anyhow::Result<()> {
517        let dir = tempfile::tempdir()?;
518        std::fs::write(
519            dir.path().join("README.md"),
520            "---\ntools:\n  - [echo, { }]\n---\n",
521        )?;
522
523        let config = ServerConfig::new(dir.path().to_path_buf())
524            .with_env(HashMap::from([(
525                "DATABASE_URL".to_string(),
526                "postgresql://gateway:gateway@localhost:5432/gateway_dev".to_string(),
527            )]))
528            .with_sandbox_env(SandboxEnv::from_host_process());
529        let state = ServerState::from_config(&config)?;
530
531        let request = ActionRequest {
532            command: vec!["echo".to_string(), "$DATABASE_URL".to_string()],
533
534            env: HashMap::new(),
535        };
536
537        let response = execute_action("README.md", &state, request).await;
538        assert_eq!(response.status(), StatusCode::OK);
539
540        let body = response_text(response).await;
541        let json: serde_json::Value = serde_json::from_str(&body)?;
542        let data = json
543            .get("data")
544            .ok_or_else(|| anyhow::anyhow!("missing data"))?;
545        let stdout = data
546            .get("stdout")
547            .and_then(|v| v.as_str())
548            .ok_or_else(|| anyhow::anyhow!("missing stdout"))?;
549        let stderr = data
550            .get("stderr")
551            .and_then(|v| v.as_str())
552            .ok_or_else(|| anyhow::anyhow!("missing stderr"))?;
553        let returncode = data
554            .get("returncode")
555            .and_then(serde_json::Value::as_i64)
556            .ok_or_else(|| anyhow::anyhow!("missing returncode"))?;
557        assert!(stdout.contains("$DATABASE_URL"));
558        assert!(!stdout.contains("postgresql://gateway:gateway@localhost:5432/gateway_dev"));
559        assert_eq!(stderr, "");
560        assert_eq!(returncode, 0);
561        Ok(())
562    }
563}