1#[cfg(feature = "native-browser")]
2mod native_impl {
3 use crate::browser::backend::BrowserBackend;
4 use crate::diagrams::{DiagramError, DiagramResult};
5 use async_trait::async_trait;
6 use headless_chrome::{Browser, LaunchOptions};
7 use std::ffi::OsStr;
8 use std::io::Write;
9 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
10 use std::sync::Arc;
11 use std::time::Duration;
12 use tempfile::{Builder, NamedTempFile};
13 use tokio::sync::{RwLock, Semaphore};
14
15 const CHROME_ARGS: &[&str] = &[
16 "--no-sandbox",
17 "--disable-setuid-sandbox",
18 "--disable-dev-shm-usage",
19 "--disable-gpu",
20 "--disable-web-security",
21 "--disable-software-rasterizer",
22 "--disable-features=IsolateOrigins,site-per-process",
23 "--font-render-hinting=none",
24 "--allow-file-access-from-files",
25 ];
26
27 const DEFAULT_IDLE_TIMEOUT: Duration = Duration::from_secs(120);
28
29 pub struct NativeBackend {
31 browser: Arc<RwLock<Browser>>,
32 harness_url: String,
33 semaphore: Arc<Semaphore>,
34 _harness_file: NamedTempFile,
35 context_ttl_requests: usize,
36 request_count: AtomicUsize,
37 restarting: AtomicBool,
38 }
39
40 impl NativeBackend {
41 pub async fn new(pool_size: usize, context_ttl_requests: usize) -> Result<Self, String> {
42 let (harness_url, harness_file) = Self::build_harness()?;
43 let browser = Self::spawn_browser().await?;
44 Ok(Self {
45 browser: Arc::new(RwLock::new(browser)),
46 harness_url,
47 semaphore: Arc::new(Semaphore::new(pool_size)),
48 _harness_file: harness_file,
49 context_ttl_requests,
50 request_count: AtomicUsize::new(0),
51 restarting: AtomicBool::new(false),
52 })
53 }
54
55 fn build_harness() -> Result<(String, NamedTempFile), String> {
56 let mut temp_file = Builder::new()
57 .suffix(".html")
58 .tempfile()
59 .map_err(|e| format!("Failed to create temp harness: {}", e))?;
60
61 let mermaid_js = include_str!("../../resources/browser/mermaid.min.js");
62 let bpmn_js = include_str!("../../resources/browser/bpmn-viewer.production.min.js");
63 let index_html = include_str!("../../resources/browser/index.html");
64
65 let html = index_html.replace(
66 "<!-- KROKI_SCRIPTS -->",
67 &("<script>".to_string()
68 + mermaid_js
69 + "</script><script>"
70 + bpmn_js
71 + "</script>"),
72 );
73
74 temp_file
75 .write_all(html.as_bytes())
76 .map_err(|e| format!("Failed to write harness: {}", e))?;
77
78 let path = temp_file
79 .path()
80 .to_str()
81 .ok_or_else(|| "Failed to build harness URL".to_string())?;
82 let harness_url = format!("file://{}", path);
83 tracing::debug!("Local serverless harness created at {}", harness_url);
84
85 Ok((harness_url, temp_file))
86 }
87
88 fn default_launch_options() -> LaunchOptions<'static> {
89 let args: Vec<&'static OsStr> = CHROME_ARGS.iter().map(OsStr::new).collect();
90
91 LaunchOptions {
92 args,
93 idle_browser_timeout: DEFAULT_IDLE_TIMEOUT,
94 ..Default::default()
95 }
96 }
97
98 async fn spawn_browser() -> Result<Browser, String> {
99 let options = Self::default_launch_options();
100 tokio::task::spawn_blocking(move || Browser::new(options))
101 .await
102 .map_err(|e| format!("Browser spawn join failed: {}", e))?
103 .map_err(|e| e.to_string())
104 }
105
106 async fn restart_browser(&self) -> Result<(), String> {
107 let new_browser = Self::spawn_browser().await?;
108 let mut guard = self.browser.write().await;
109 *guard = new_browser;
110 Ok(())
111 }
112
113 fn should_restart(&self) -> bool {
114 if self.context_ttl_requests == 0 {
115 return false;
116 }
117 let count = self.request_count.fetch_add(1, Ordering::Relaxed) + 1;
118 count >= self.context_ttl_requests
119 }
120
121 async fn maybe_restart(&self) {
122 if !self.should_restart() {
123 return;
124 }
125 if self
126 .restarting
127 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
128 .is_err()
129 {
130 return;
131 }
132
133 if let Err(err) = self.restart_browser().await {
134 tracing::error!("Native browser restart failed: {}", err);
135 }
136
137 self.request_count.store(0, Ordering::Relaxed);
138 self.restarting.store(false, Ordering::Release);
139 }
140
141 async fn acquire_browser(&self) -> Browser {
142 let guard = self.browser.read().await;
143 guard.clone()
144 }
145
146 async fn do_render(
147 &self,
148 tab: &headless_chrome::Tab,
149 diagram_type: &str,
150 source: &str,
151 _format: &str,
152 ) -> DiagramResult<Vec<u8>> {
153 tab.navigate_to(&self.harness_url)
154 .map_err(|e| DiagramError::ProcessFailed(format!("Navigation failed: {}", e)))?;
155 tab.wait_for_element("#container")
156 .map_err(|e| DiagramError::ProcessFailed(format!("Harness load timeout: {}", e)))?;
157
158 let font_injection =
159 "const style = document.getElementById('kroki-fonts'); if (style) { style.innerHTML = window.krokiFontCss || ''; }";
160 tab.evaluate(font_injection, false).map_err(|e| {
161 DiagramError::ProcessFailed(format!("Font injection failed: {}", e))
162 })?;
163
164 match diagram_type {
165 "mermaid" => {
166 tab.evaluate(
167 "new Promise(r => { const check = () => window.mermaid ? r() : setTimeout(check, 50); check(); })",
168 true,
169 )
170 .map_err(|e| DiagramError::ProcessFailed(format!("Mermaid load timeout: {}", e)))?;
171 }
172 "bpmn" => {
173 tab.evaluate(
174 "new Promise(r => { const check = () => window.BpmnJS ? r() : setTimeout(check, 50); check(); })",
175 true,
176 )
177 .map_err(|e| DiagramError::ProcessFailed(format!("BPMN load timeout: {}", e)))?;
178 }
179 _ => {
180 return Err(DiagramError::UnsupportedFormat {
181 provider: diagram_type.to_string(),
182 format: _format.to_string(),
183 })
184 }
185 }
186
187 let render_expr = format!(
188 "window.kroki.render{}({})",
189 match diagram_type {
190 "mermaid" => "Mermaid",
191 "bpmn" => "Bpmn",
192 _ => unreachable!(),
193 },
194 serde_json::to_string(source).unwrap()
195 );
196
197 let remote_object = tab.evaluate(&render_expr, true).map_err(|e| {
198 DiagramError::ProcessFailed(format!("Render execution failed: {}", e))
199 })?;
200
201 let result = remote_object
202 .value
203 .and_then(|v| v.as_str().map(|s| s.to_string()))
204 .ok_or_else(|| {
205 DiagramError::ProcessFailed("Render returned null or non-string".to_string())
206 })?;
207
208 if result.is_empty() {
209 return Err(DiagramError::ProcessFailed(
210 "Render produced empty output".to_string(),
211 ));
212 }
213
214 Ok(result.into_bytes())
215 }
216 }
217
218 #[async_trait]
219 impl BrowserBackend for NativeBackend {
220 async fn render(
221 &self,
222 diagram_type: &str,
223 source: &str,
224 format: &str,
225 ) -> DiagramResult<Vec<u8>> {
226 self.maybe_restart().await;
227
228 let _permit = self.semaphore.acquire().await.map_err(|_| {
229 DiagramError::ProcessFailed("Native backend semaphore was closed".to_string())
230 })?;
231
232 tracing::debug!("Creating new tab...");
233 let browser = self.acquire_browser().await;
234 let tab = browser
235 .new_tab()
236 .map_err(|e| DiagramError::ProcessFailed(format!("Failed to create tab: {}", e)))?;
237
238 let result = self.do_render(&tab, diagram_type, source, format).await;
239
240 let _ = tab.close(false);
241
242 result
243 }
244
245 async fn health(&self) -> serde_json::Value {
246 let browser = self.acquire_browser().await;
247 let tabs_count = browser
248 .get_tabs()
249 .lock()
250 .map(|tabs| tabs.len())
251 .unwrap_or(0);
252
253 serde_json::json!({
254 "status": "ok",
255 "backend": "headless_chrome",
256 "tabs": tabs_count,
257 "harness_url": self.harness_url,
258 "concurrency_permits_available": self.semaphore.available_permits()
259 })
260 }
261 }
262
263 pub use NativeBackend as Backend;
264}
265
266#[cfg(feature = "native-browser")]
267pub use native_impl::Backend as NativeBackend;