1use axum::body::Body;
23use axum::extract::State;
24use axum::http::{HeaderMap, HeaderValue, StatusCode, header};
25use axum::response::{IntoResponse, Response};
26use std::sync::Arc;
27
28pub const DEFAULT_UI_UPSTREAM: &str = "https://app.opencode.ai";
29
30pub fn is_allowed_cors_origin(origin: &str, extra: &[String]) -> bool {
32 let origin = origin.trim();
33 if origin.is_empty() {
34 return true; }
36 if origin.starts_with("http://localhost:") || origin.starts_with("http://127.0.0.1:") {
37 return true;
38 }
39 if origin.starts_with("oc://renderer") {
40 return true;
41 }
42 if origin == "tauri://localhost"
43 || origin == "http://tauri.localhost"
44 || origin == "https://tauri.localhost"
45 {
46 return true;
47 }
48 if is_opencode_ai_origin(origin) {
49 return true;
50 }
51 extra.iter().any(|e| e == origin)
52}
53
54fn is_opencode_ai_origin(origin: &str) -> bool {
58 let Some(rest) = origin.strip_prefix("https://") else {
59 return false;
60 };
61 if rest.contains('/') || rest.contains(':') {
62 return false;
63 }
64 let labels: Vec<&str> = rest.split('.').collect();
65 if labels.len() < 2 {
66 return false;
67 }
68 if labels[labels.len() - 2] != "opencode" || labels[labels.len() - 1] != "ai" {
69 return false;
70 }
71 labels.iter().all(|l| {
72 !l.is_empty()
73 && l.chars()
74 .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-')
75 })
76}
77
78pub async fn cors_gate(
85 State(st): State<Arc<crate::AppState>>,
86 req: axum::extract::Request,
87 next: axum::middleware::Next,
88) -> Response {
89 use axum::http::Method;
90 let origin = req
91 .headers()
92 .get(header::ORIGIN)
93 .and_then(|v| v.to_str().ok())
94 .map(str::to_string);
95 let allowed = origin
96 .as_deref()
97 .map(|o| is_allowed_cors_origin(o, &st.cors_extra))
98 .unwrap_or(false);
99
100 let is_preflight = req.method() == Method::OPTIONS
101 && req.headers().contains_key("access-control-request-method");
102 if is_preflight && allowed {
103 let Some(origin) = origin else {
104 return next.run(req).await;
105 };
106 let mut resp = StatusCode::NO_CONTENT.into_response();
107 let h = resp.headers_mut();
108 h.insert(
109 "access-control-allow-origin",
110 HeaderValue::from_str(&origin).unwrap_or(HeaderValue::from_static("*")),
111 );
112 h.insert(
113 "vary",
114 HeaderValue::from_static("Access-Control-Request-Headers, Origin"),
115 );
116 h.insert(
117 "access-control-allow-methods",
118 HeaderValue::from_static("GET, HEAD, PUT, PATCH, POST, DELETE"),
119 );
120 if let Some(reqh) = req.headers().get("access-control-request-headers").cloned() {
122 h.insert("access-control-allow-headers", reqh);
123 } else {
124 h.insert("access-control-allow-headers", HeaderValue::from_static(""));
125 }
126 h.insert("access-control-max-age", HeaderValue::from_static("86400"));
127 return resp;
128 }
129
130 let mut resp = next.run(req).await;
131 if allowed
132 && let Some(origin) = origin
133 && let Ok(v) = HeaderValue::from_str(&origin)
134 {
135 resp.headers_mut().insert("access-control-allow-origin", v);
136 }
137 resp
138}
139
140pub const ASSET_CACHE_CAP: usize = 32 * 1024 * 1024;
150pub const ASSET_ENTRY_CAP: usize = 8 * 1024 * 1024;
154
155#[derive(Clone)]
156pub struct AssetEntry {
157 pub body: bytes::Bytes, pub br: bytes::Bytes,
159 pub gzip: bytes::Bytes,
160 pub content_type: String,
161 pub etag: String,
162}
163
164struct CacheState {
165 map: std::collections::HashMap<String, std::sync::Arc<AssetEntry>>,
166 order: std::collections::VecDeque<String>,
167 bytes: usize,
168}
169
170pub struct UiProxy {
173 pub enabled: bool,
174 pub upstream: String,
175 client: reqwest::Client,
176 cache: parking_lot::Mutex<CacheState>,
177 pub prewarm: bool,
181 pub embedded: bool,
185}
186
187impl UiProxy {
188 pub fn from_env() -> Self {
189 let enabled = std::env::var("OCSERVE_UI")
190 .map(|v| v != "0" && !v.eq_ignore_ascii_case("off"))
191 .unwrap_or(true);
192 let upstream = std::env::var("OCSERVE_UI_UPSTREAM")
193 .ok()
194 .filter(|u| !u.is_empty())
195 .unwrap_or_else(|| DEFAULT_UI_UPSTREAM.to_string());
196 let prewarm = std::env::var("OCSERVE_UI_PREWARM")
197 .map(|v| v == "1")
198 .unwrap_or(false);
199 let embedded = std::env::var("OCSERVE_DISABLE_EMBEDDED_WEB_UI")
200 .map(|v| v == "0" || v.eq_ignore_ascii_case("off"))
201 .unwrap_or(true);
202 let mut p = Self::new(enabled, upstream);
203 p.prewarm = prewarm;
204 p.embedded = embedded;
205 p
206 }
207
208 pub fn new(enabled: bool, upstream: String) -> Self {
209 let client = reqwest::Client::builder()
210 .timeout(std::time::Duration::from_secs(30))
211 .build()
212 .unwrap_or_default();
213 Self {
214 enabled,
215 upstream,
216 client,
217 cache: parking_lot::Mutex::new(CacheState {
218 map: std::collections::HashMap::new(),
219 order: std::collections::VecDeque::new(),
220 bytes: 0,
221 }),
222 prewarm: false,
223 embedded: true,
224 }
225 }
226
227 fn cache_get(&self, key: &str) -> Option<std::sync::Arc<AssetEntry>> {
230 let mut c = self.cache.lock();
231 let e = c.map.get(key).cloned()?;
232 if let Some(pos) = c.order.iter().position(|k| k == key) {
234 let k = c.order.remove(pos).expect("position");
235 c.order.push_back(k);
236 }
237 Some(e)
238 }
239
240 fn cache_put(&self, key: String, entry: std::sync::Arc<AssetEntry>) {
243 let sz = entry.body.len() + entry.br.len() + entry.gzip.len();
244 if sz > ASSET_ENTRY_CAP {
245 return;
246 }
247 let mut c = self.cache.lock();
248 if c.map.contains_key(&key) {
249 return;
250 }
251 while c.bytes + sz > ASSET_CACHE_CAP {
252 let Some(front) = c.order.pop_front() else {
253 break;
254 };
255 if let Some(old) = c.map.remove(&front) {
256 c.bytes = c
257 .bytes
258 .saturating_sub(old.body.len() + old.br.len() + old.gzip.len());
259 }
260 }
261 c.bytes += sz;
262 c.map.insert(key.clone(), entry);
263 c.order.push_back(key);
264 ocserve_metrics::gauge("ocserve_webui_asset_cache_bytes", c.bytes as i64);
265 }
266
267 fn is_cacheable_asset(path: &str) -> bool {
271 path.starts_with("/assets/")
272 }
273}
274
275fn csp(hash: &str) -> String {
278 format!(
279 "default-src 'self'; script-src 'self' 'wasm-unsafe-eval'{}; style-src 'self' 'unsafe-inline'; img-src 'self' data: https: blob:; font-src 'self' data:; media-src 'self' data:; connect-src * data: blob:",
280 if hash.is_empty() {
281 String::new()
282 } else {
283 format!(" 'sha256-{hash}'")
284 }
285 )
286}
287
288fn theme_preload_script(body: &str) -> Option<&str> {
291 let needle = "id=\"oc-theme-preload-script\"";
292 let alt = "id='oc-theme-preload-script'";
293 let start = body.find(needle).or_else(|| body.find(alt))?;
294 let open_end = body[start..].find('>')? + start;
296 let close = body[open_end..].find("</script>")? + open_end;
297 Some(&body[open_end + 1..close])
298}
299
300fn csp_for_html(body: &str) -> String {
301 use base64::Engine as _;
302 use sha2::Digest as _;
303 match theme_preload_script(body) {
304 Some(script) => {
305 let digest = sha2::Sha256::digest(script.as_bytes());
306 csp(&base64::engine::general_purpose::STANDARD.encode(digest))
307 }
308 None => csp(""),
309 }
310}
311
312pub async fn ui_fallback(
316 State(st): State<Arc<crate::AppState>>,
317 req: axum::extract::Request,
318) -> Response {
319 let ui = &st.ui;
320 if !ui.enabled {
321 return (
322 StatusCode::NOT_FOUND,
323 axum::Json(serde_json::json!({"error": "Not Found"})),
324 )
325 .into_response();
326 }
327 if ui.embedded && crate::embedded_ui::available() {
331 return embedded_response(req.uri().path());
332 }
333 let path = req
334 .uri()
335 .path_and_query()
336 .map(|pq| pq.as_str())
337 .unwrap_or("/");
338 let url = format!("{}{}", ui.upstream.trim_end_matches('/'), path);
339 let req_path = req.uri().path().to_string();
341 let req_headers_ae = req
342 .headers()
343 .get(header::ACCEPT_ENCODING)
344 .and_then(|v| v.to_str().ok())
345 .map(str::to_string);
346 let req_headers_inm = req
347 .headers()
348 .get(header::IF_NONE_MATCH)
349 .and_then(|v| v.to_str().ok())
350 .map(str::to_string);
351
352 let method = reqwest::Method::from_bytes(req.method().as_str().as_bytes())
353 .unwrap_or(reqwest::Method::GET);
354 let mut builder = ui.client.request(method, &url);
355 builder = builder.headers(forward_headers(req.headers()));
357 let body = axum::body::to_bytes(req.into_body(), 64 * 1024 * 1024)
359 .await
360 .unwrap_or_default();
361 if !body.is_empty() {
362 builder = builder.body(body);
363 }
364
365 let resp = match builder.send().await {
366 Ok(r) => r,
367 Err(e) => {
368 tracing::warn!("ui proxy: {url} failed: {e:#}");
369 return (
370 StatusCode::BAD_GATEWAY,
371 axum::Json(serde_json::json!({"error": "UI upstream unreachable"})),
372 )
373 .into_response();
374 }
375 };
376
377 let status = StatusCode::from_u16(resp.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
378 let content_type = resp
379 .headers()
380 .get(reqwest::header::CONTENT_TYPE)
381 .and_then(|v| v.to_str().ok())
382 .map(str::to_string);
383 let is_html = content_type
384 .as_deref()
385 .map(|c| c.contains("text/html"))
386 .unwrap_or(false);
387 let upstream_cache_control = resp
388 .headers()
389 .get(reqwest::header::CACHE_CONTROL)
390 .and_then(|v| v.to_str().ok())
391 .map(str::to_string);
392 let is_asset = UiProxy::is_cacheable_asset(&req_path);
393
394 if is_asset && status == StatusCode::OK {
397 let cache_key = req_path.to_string();
398 let ae = req_headers_ae.as_deref();
399 let inm = req_headers_inm.as_deref();
400 if let Some(entry) = ui.cache_get(&cache_key) {
401 ocserve_metrics::counter("ocserve_webui_asset_cache_hits_total", 1);
402 return asset_response(&entry, ae, inm, true);
403 }
404 ocserve_metrics::counter("ocserve_webui_asset_cache_misses_total", 1);
405 let bytes = resp.bytes().await.unwrap_or_default();
406 let ct = content_type
407 .clone()
408 .unwrap_or_else(|| "application/octet-stream".into());
409 let entry = tokio::task::spawn_blocking(move || {
413 let etag = crate::compress::etag_for(&bytes);
414 std::sync::Arc::new(AssetEntry {
415 br: bytes::Bytes::from(crate::compress::compress(
416 &bytes,
417 crate::compress::Encoding::Brotli,
418 )),
419 gzip: bytes::Bytes::from(crate::compress::compress(
420 &bytes,
421 crate::compress::Encoding::Gzip,
422 )),
423 body: bytes,
424 content_type: ct,
425 etag,
426 })
427 })
428 .await
429 .unwrap_or_else(|_| {
430 std::sync::Arc::new(AssetEntry {
431 body: bytes::Bytes::new(),
432 br: bytes::Bytes::new(),
433 gzip: bytes::Bytes::new(),
434 content_type: "application/octet-stream".into(),
435 etag: "\"\"".into(),
436 })
437 });
438 if !entry.body.is_empty() {
440 ui.cache_put(cache_key, entry.clone());
441 }
442 return asset_response(&entry, ae, inm, false);
443 }
444
445 if is_html {
450 let bytes = resp.bytes().await.unwrap_or_default();
451 let csp_value = csp_for_html(&String::from_utf8_lossy(&bytes));
452 let mut out = Response::builder().status(status);
453 if let Some(ct) = &content_type {
454 out = out.header(header::CONTENT_TYPE, ct);
455 }
456 if let Some(cc) = upstream_cache_control {
458 out = out.header(header::CACHE_CONTROL, cc);
459 }
460 out = out.header("content-security-policy", csp_value);
461 return out
462 .body(Body::from(bytes))
463 .unwrap_or_else(|_| StatusCode::BAD_GATEWAY.into_response());
464 }
465
466 let mut out = Response::builder().status(status);
468 if let Some(ct) = &content_type {
469 out = out.header(header::CONTENT_TYPE, ct);
470 }
471 if let Some(cc) = upstream_cache_control {
472 out = out.header(header::CACHE_CONTROL, cc);
473 }
474 out = out.header("content-security-policy", csp(""));
475 out.body(Body::from_stream(resp.bytes_stream()))
476 .unwrap_or_else(|_| StatusCode::BAD_GATEWAY.into_response())
477}
478
479fn asset_response(
482 entry: &AssetEntry,
483 accept_encoding: Option<&str>,
484 if_none_match: Option<&str>,
485 cached: bool,
486) -> Response {
487 use axum::http::{StatusCode, header};
488 let _ = cached;
489 if crate::compress::if_none_match_matches(if_none_match, &entry.etag) {
490 return Response::builder()
491 .status(StatusCode::NOT_MODIFIED)
492 .header(header::ETAG, &entry.etag)
493 .header(header::CACHE_CONTROL, "public, max-age=31536000, immutable")
494 .header(header::VARY, "Accept-Encoding")
495 .body(Body::empty())
496 .expect("static 304");
497 }
498 let enc = crate::compress::negotiate(accept_encoding);
499 let mut b = Response::builder()
500 .status(StatusCode::OK)
501 .header(header::CONTENT_TYPE, &entry.content_type)
502 .header(header::ETAG, &entry.etag)
503 .header(header::CACHE_CONTROL, "public, max-age=31536000, immutable")
505 .header(header::VARY, "Accept-Encoding");
506 let body = match enc {
507 Some(crate::compress::Encoding::Brotli) => {
508 b = b.header(header::CONTENT_ENCODING, "br");
509 entry.br.clone()
510 }
511 Some(crate::compress::Encoding::Gzip) => {
512 b = b.header(header::CONTENT_ENCODING, "gzip");
513 entry.gzip.clone()
514 }
515 None => entry.body.clone(),
516 };
517 b.body(Body::from(body)).expect("static asset")
518}
519
520fn forward_headers(src: &HeaderMap) -> HeaderMap {
521 let mut out = HeaderMap::new();
522 for (name, value) in src {
523 if matches!(
525 name.as_str(),
526 "host" | "content-length" | "transfer-encoding" | "connection" | "accept-encoding"
527 ) {
528 continue;
529 }
530 out.insert(name.clone(), value.clone());
531 }
532 out
533}
534
535fn embedded_response(path: &str) -> Response {
541 let Some((resolved, body)) = crate::embedded_ui::lookup(path) else {
542 return (
543 StatusCode::NOT_FOUND,
544 axum::Json(serde_json::json!({"error": "Not Found"})),
545 )
546 .into_response();
547 };
548 let ct = crate::embedded_ui::mime_for(resolved);
549 let mut out = Response::builder()
550 .status(StatusCode::OK)
551 .header(header::CONTENT_TYPE, ct);
552 if ct.starts_with("text/html") {
553 static CSP: std::sync::OnceLock<String> = std::sync::OnceLock::new();
559 let csp_value = CSP
560 .get_or_init(|| csp_for_html(&String::from_utf8_lossy(body)))
561 .clone();
562 out = out.header("content-security-policy", csp_value);
563 } else {
564 out = out.header("content-security-policy", csp(""));
565 }
566 out.body(Body::from(bytes::Bytes::from_static(body)))
567 .unwrap_or_else(|_| StatusCode::BAD_GATEWAY.into_response())
568}
569
570pub async fn prewarm_assets(st: &Arc<crate::AppState>) -> anyhow::Result<()> {
574 let ui = &st.ui;
575 if !ui.enabled {
576 return Ok(());
577 }
578 let base = ui.upstream.trim_end_matches('/');
579 let html = ui
580 .client
581 .get(format!("{base}/"))
582 .send()
583 .await?
584 .text()
585 .await?;
586 let mut assets: Vec<String> = Vec::new();
587 for part in html.split('"').filter(|s| s.starts_with("/assets/")) {
588 let a = part.split(['"', '?']).next().unwrap_or("");
589 if a.starts_with("/assets/") && !assets.iter().any(|x| x == a) {
590 assets.push(a.to_string());
591 }
592 }
593 for a in &assets {
594 let resp = ui.client.get(format!("{base}{a}")).send().await;
597 let Ok(resp) = resp else { continue };
598 if resp.status().as_u16() != 200 {
599 continue;
600 }
601 let ct = resp
602 .headers()
603 .get(reqwest::header::CONTENT_TYPE)
604 .and_then(|v| v.to_str().ok())
605 .unwrap_or("application/octet-stream")
606 .to_string();
607 let bytes = resp.bytes().await.unwrap_or_default();
608 if bytes.is_empty() || bytes.len() > ASSET_ENTRY_CAP {
609 continue;
610 }
611 let ct_c = ct.clone();
612 let entry = tokio::task::spawn_blocking(move || {
613 let etag = crate::compress::etag_for(&bytes);
614 std::sync::Arc::new(AssetEntry {
615 br: bytes::Bytes::from(crate::compress::compress(
616 &bytes,
617 crate::compress::Encoding::Brotli,
618 )),
619 gzip: bytes::Bytes::from(crate::compress::compress(
620 &bytes,
621 crate::compress::Encoding::Gzip,
622 )),
623 body: bytes,
624 content_type: ct_c,
625 etag,
626 })
627 })
628 .await?;
629 ui.cache_put(a.clone(), entry);
630 tracing::info!("ui prewarm: cached {a}");
631 }
632 Ok(())
633}
634
635#[cfg(test)]
636mod tests {
637 use super::*;
638
639 #[test]
640 fn cors_matches_upstream_rules() {
641 let none: Vec<String> = vec![];
642 assert!(is_allowed_cors_origin("http://localhost:5173", &none));
644 assert!(is_allowed_cors_origin("http://127.0.0.1:4096", &none));
645 assert!(is_allowed_cors_origin("oc://renderer", &none));
646 assert!(is_allowed_cors_origin("tauri://localhost", &none));
647 assert!(is_allowed_cors_origin("http://tauri.localhost", &none));
648 assert!(is_allowed_cors_origin("https://tauri.localhost", &none));
649 assert!(is_allowed_cors_origin("https://opencode.ai", &none));
650 assert!(is_allowed_cors_origin("https://app.opencode.ai", &none));
651 assert!(is_allowed_cors_origin("https://a-b.c.opencode.ai", &none));
652 assert!(!is_allowed_cors_origin("https://evil.com", &none));
654 assert!(!is_allowed_cors_origin("http://opencode.ai", &none)); assert!(!is_allowed_cors_origin(
656 "https://opencode.ai.evil.com",
657 &none
658 ));
659 assert!(!is_allowed_cors_origin("https://OPEnc0de.ai", &none)); assert!(!is_allowed_cors_origin("https://opencode.ai:8080", &none));
661 let extra = vec!["https://my.box".to_string()];
663 assert!(is_allowed_cors_origin("https://my.box", &extra));
664 assert!(!is_allowed_cors_origin("https://my.box", &none));
665 }
666
667 #[test]
668 fn theme_preload_hash_extraction() {
669 let html = r#"<script id="oc-theme-preload-script">window.x=1</script>"#;
670 assert_eq!(theme_preload_script(html), Some("window.x=1"));
671 let html2 = r#"<script defer id='oc-theme-preload-script'>y</script>"#;
672 assert_eq!(theme_preload_script(html2), Some("y"));
673 assert_eq!(theme_preload_script("<html></html>"), None);
674 let with = csp_for_html(html);
676 assert!(with.contains("sha256-"), "{with}");
677 assert!(with.contains("'sha256-"), "{with}");
678 let without = csp_for_html("<html></html>");
679 assert!(!without.contains("sha256"));
680 }
681}