feat: polish workflow studio and workspace assets

src/services/comfyui_process.rs

@@ -1,6 +1,6 @@
 use crate::config::ComfyUIConfig;
 use serde::Serialize;
-use std::{path::Path, sync::Arc, time::Instant};
+use std::{path::Path, process::Stdio, sync::Arc, time::Instant};
 use tokio::{
     process::{Child, Command},
     sync::{Mutex, RwLock},
@@ -67,7 +67,7 @@ impl ComfyUIProcess {
     }
 
     pub fn internal_base_url(&self) -> String {
-        if self.config.enabled {
+        if self.config.enabled || self.config.uses_docker_backend() {
             format!("http://127.0.0.1:{}", self.config.port)
         } else {
             self.config.target_base_url()
@@ -92,7 +92,9 @@ impl ComfyUIProcess {
     pub async fn status_response(&self) -> ComfyUIStatusResponse {
         let status = self.status().await;
         let base_url = self.internal_base_url();
-        let mode = if self.config.enabled {
+        let mode = if self.config.uses_docker_backend() {
+            "docker"
+        } else if self.config.enabled {
             "embedded"
         } else {
             "external"
@@ -154,6 +156,10 @@ impl ComfyUIProcess {
     }
 
     pub async fn start(&self) -> anyhow::Result<()> {
+        if self.config.uses_docker_backend() {
+            return self.start_docker().await;
+        }
+
         if !self.config.enabled {
             *self.status.write().await = ComfyUIStatus::Disabled;
             return Ok(());
@@ -223,6 +229,28 @@ impl ComfyUIProcess {
     }
 
     pub async fn stop(&self) -> anyhow::Result<()> {
+        if self.config.uses_docker_backend() {
+            let name = self.config.container_name.clone();
+            let output = Command::new(docker_bin())
+                .args(["stop", "-t", "10", &name])
+                .stdout(Stdio::null())
+                .stderr(Stdio::piped())
+                .output()
+                .await?;
+            if !output.status.success() {
+                let stderr = String::from_utf8_lossy(&output.stderr);
+                if !stderr.contains("No such container") {
+                    *self.status.write().await = ComfyUIStatus::Unhealthy {
+                        last_error: stderr.trim().to_string(),
+                    };
+                    anyhow::bail!("docker stop failed: {}", stderr.trim());
+                }
+            }
+            *self.started_at.lock().await = None;
+            *self.status.write().await = ComfyUIStatus::Stopped;
+            return Ok(());
+        }
+
         let mut child_slot = self.child.lock().await;
         if let Some(child) = child_slot.as_mut() {
             if let Err(err) = child.start_kill() {
@@ -265,6 +293,89 @@ impl ComfyUIProcess {
             .map(|started| started.elapsed().as_secs())
             .unwrap_or(0)
     }
+
+    async fn start_docker(&self) -> anyhow::Result<()> {
+        if self.health_check().await {
+            *self.started_at.lock().await = Some(Instant::now());
+            *self.status.write().await = ComfyUIStatus::Running {
+                pid: None,
+                uptime_secs: 0,
+            };
+            return Ok(());
+        }
+
+        *self.status.write().await = ComfyUIStatus::Starting;
+        let name = self.config.container_name.clone();
+        let _ = Command::new(docker_bin())
+            .args(["rm", "-f", &name])
+            .stdout(Stdio::null())
+            .stderr(Stdio::null())
+            .status()
+            .await;
+
+        let mut command = Command::new(docker_bin());
+        command
+            .arg("run")
+            .arg("-d")
+            .arg("--rm")
+            .arg("--name")
+            .arg(&name)
+            .arg("-p")
+            .arg(format!("127.0.0.1:{}:8188", self.config.port))
+            .arg("-e")
+            .arg(format!("BOOSAPI_BASE={}", self.model_backend_url()))
+            .arg("-e")
+            .arg(format!("BOOSAPI_BUILTIN_KEY={}", self.builtin_key))
+            .arg("-e")
+            .arg(format!("BOOSAPI_KEY={}", self.builtin_key));
+
+        let data_root = std::env::var("BOOSAPI_COMFY_DATA_ROOT")
+            .unwrap_or_else(|_| "/data/boosapi-comfyui".to_string());
+        for (host, container) in [
+            ("input", "/workspace/ComfyUI/input"),
+            ("output", "/workspace/ComfyUI/output"),
+            ("temp", "/workspace/ComfyUI/temp"),
+            ("user", "/workspace/ComfyUI/user"),
+        ] {
+            let host_path = format!("{data_root}/{host}");
+            tokio::fs::create_dir_all(&host_path).await?;
+            command
+                .arg("--mount")
+                .arg(format!("type=bind,source={host_path},target={container}"));
+        }
+
+        command.arg(&self.config.docker_image);
+        let output = command.output().await?;
+        if !output.status.success() {
+            let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
+            *self.status.write().await = ComfyUIStatus::Unhealthy {
+                last_error: stderr.clone(),
+            };
+            anyhow::bail!("docker run failed: {stderr}");
+        }
+
+        let deadline = Instant::now() + std::time::Duration::from_secs(45);
+        while Instant::now() < deadline {
+            if self.health_check().await {
+                *self.started_at.lock().await = Some(Instant::now());
+                *self.status.write().await = ComfyUIStatus::Running {
+                    pid: None,
+                    uptime_secs: 0,
+                };
+                return Ok(());
+            }
+            tokio::time::sleep(std::time::Duration::from_millis(750)).await;
+        }
+
+        *self.status.write().await = ComfyUIStatus::Unhealthy {
+            last_error: "Docker ComfyUI did not become ready in 45s".to_string(),
+        };
+        anyhow::bail!("Docker ComfyUI did not become ready in 45s")
+    }
+}
+
+fn docker_bin() -> String {
+    std::env::var("BOOSAPI_DOCKER_BIN").unwrap_or_else(|_| "docker".to_string())
 }
 
 fn write_boosapi_config(

src/services/mod.rs

@@ -5,6 +5,7 @@ pub mod comfyui_process;
 pub mod cron;
 pub mod notification;
 pub mod oauth;
+pub mod register_captcha;
 pub mod registration_limiter;
 pub mod route_rebuild;
 pub mod shared;

src/services/register_captcha.rs

@@ -0,0 +1,208 @@
+use std::{
+    collections::HashMap,
+    net::IpAddr,
+    time::{Duration, Instant},
+};
+
+use base64::{engine::general_purpose, Engine as _};
+use parking_lot::Mutex;
+use serde::Serialize;
+use uuid::Uuid;
+
+const CAPTCHA_TTL: Duration = Duration::from_secs(5 * 60);
+const CAPTCHA_ALPHABET: &[u8] = b"ABCDEFGHJKLMNPQRSTUVWXYZ23456789";
+
+#[derive(Debug, Clone, Serialize)]
+pub struct RegisterCaptchaChallenge {
+    pub id: String,
+    pub image: String,
+    #[serde(rename = "expiresInSeconds")]
+    pub expires_in_seconds: u64,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum CaptchaVerifyError {
+    Missing,
+    Expired,
+    Mismatch,
+}
+
+#[derive(Debug, Clone)]
+struct CaptchaEntry {
+    answer: String,
+    expires_at: Instant,
+    ip: Option<IpAddr>,
+}
+
+#[derive(Default)]
+pub struct RegisterCaptchaStore {
+    inner: Mutex<HashMap<String, CaptchaEntry>>,
+}
+
+impl RegisterCaptchaStore {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn issue(&self, ip: Option<IpAddr>) -> RegisterCaptchaChallenge {
+        let now = Instant::now();
+        let id = Uuid::new_v4().simple().to_string();
+        let answer = generate_answer();
+        let expires_at = now + CAPTCHA_TTL;
+
+        let mut guard = self.inner.lock();
+        guard.retain(|_, entry| entry.expires_at > now);
+        guard.insert(
+            id.clone(),
+            CaptchaEntry {
+                answer: answer.clone(),
+                expires_at,
+                ip,
+            },
+        );
+
+        RegisterCaptchaChallenge {
+            id,
+            image: render_svg_data_url(&answer),
+            expires_in_seconds: CAPTCHA_TTL.as_secs(),
+        }
+    }
+
+    pub fn verify(
+        &self,
+        id: &str,
+        code: &str,
+        ip: Option<IpAddr>,
+    ) -> Result<(), CaptchaVerifyError> {
+        let captcha_id = id.trim();
+        let submitted = normalize_code(code);
+        if captcha_id.is_empty() || submitted.is_empty() {
+            return Err(CaptchaVerifyError::Missing);
+        }
+
+        let now = Instant::now();
+        let mut guard = self.inner.lock();
+        guard.retain(|_, entry| entry.expires_at > now);
+
+        let Some(entry) = guard.remove(captcha_id) else {
+            return Err(CaptchaVerifyError::Expired);
+        };
+
+        if entry.expires_at <= now {
+            return Err(CaptchaVerifyError::Expired);
+        }
+
+        if entry.ip.is_some() && ip.is_some() && entry.ip != ip {
+            return Err(CaptchaVerifyError::Mismatch);
+        }
+
+        if entry.answer == submitted {
+            Ok(())
+        } else {
+            Err(CaptchaVerifyError::Mismatch)
+        }
+    }
+}
+
+fn generate_answer() -> String {
+    let mut bytes = [0u8; 16];
+    bytes.copy_from_slice(Uuid::new_v4().as_bytes());
+    (0..5)
+        .map(|index| {
+            let offset = bytes[index] as usize % CAPTCHA_ALPHABET.len();
+            CAPTCHA_ALPHABET[offset] as char
+        })
+        .collect()
+}
+
+fn normalize_code(code: &str) -> String {
+    code.trim()
+        .chars()
+        .filter(|ch| !ch.is_whitespace())
+        .flat_map(|ch| ch.to_uppercase())
+        .collect()
+}
+
+fn render_svg_data_url(answer: &str) -> String {
+    let seed = Uuid::new_v4();
+    let bytes = seed.as_bytes();
+    let mut chars = String::new();
+    for (index, ch) in answer.chars().enumerate() {
+        let x = 18 + index * 28;
+        let y = 42 + (bytes[index] % 9) as i32 - 4;
+        let rotate = (bytes[index + 5] % 25) as i32 - 12;
+        chars.push_str(&format!(
+            r#"<text x="{x}" y="{y}" transform="rotate({rotate} {x} {y})">{ch}</text>"#
+        ));
+    }
+
+    let mut lines = String::new();
+    for index in 0..5 {
+        let a = bytes[index] as u32;
+        let b = bytes[index + 5] as u32;
+        let c = bytes[index + 10] as u32;
+        lines.push_str(&format!(
+            r#"<path d="M{} {} C{} {},{} {},{} {}" />"#,
+            a % 24,
+            12 + b % 44,
+            34 + b % 28,
+            4 + c % 56,
+            76 + a % 28,
+            4 + b % 56,
+            130 + c % 28,
+            12 + a % 44
+        ));
+    }
+
+    let svg = format!(
+        r##"<svg xmlns="http://www.w3.org/2000/svg" width="180" height="64" viewBox="0 0 180 64" role="img" aria-label="captcha">
+<rect width="180" height="64" rx="8" fill="#f8fafc"/>
+<g fill="none" stroke="#94a3b8" stroke-width="1.5" opacity=".55">{lines}</g>
+<g fill="#0f172a" font-family="ui-monospace, SFMono-Regular, Menlo, Consolas, monospace" font-size="28" font-weight="800" letter-spacing="3">{chars}</g>
+<g fill="#64748b" opacity=".45">
+<circle cx="22" cy="14" r="2"/><circle cx="152" cy="18" r="2"/><circle cx="46" cy="54" r="1.7"/><circle cx="128" cy="49" r="1.8"/>
+</g>
+</svg>"##
+    );
+    format!(
+        "data:image/svg+xml;base64,{}",
+        general_purpose::STANDARD.encode(svg.as_bytes())
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::net::{IpAddr, Ipv4Addr};
+
+    #[test]
+    fn captcha_is_one_time() {
+        let store = RegisterCaptchaStore::new();
+        let challenge = store.issue(None);
+        let answer = store.inner.lock().get(&challenge.id).unwrap().answer.clone();
+
+        store.verify(&challenge.id, &answer, None).unwrap();
+        assert_eq!(
+            store.verify(&challenge.id, &answer, None).unwrap_err(),
+            CaptchaVerifyError::Expired
+        );
+    }
+
+    #[test]
+    fn captcha_rejects_other_ip() {
+        let store = RegisterCaptchaStore::new();
+        let challenge = store.issue(Some(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1))));
+        let answer = store.inner.lock().get(&challenge.id).unwrap().answer.clone();
+
+        assert_eq!(
+            store
+                .verify(
+                    &challenge.id,
+                    &answer,
+                    Some(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2)))
+                )
+                .unwrap_err(),
+            CaptchaVerifyError::Mismatch
+        );
+    }
+}

src/services/workspace.rs

@@ -284,6 +284,29 @@ pub async fn record_base64_asset(
     .await
 }
 
+pub async fn record_bytes_asset(
+    state: &AppState,
+    mut input: WorkspaceAssetInput,
+    bytes: &[u8],
+    mime_type: &str,
+) -> Result<StoredWorkspaceAsset, AppError> {
+    let public_id = format!("wa_{}", Uuid::new_v4().simple());
+    input.remote_url = None;
+    let ext = extension_for_mime(mime_type);
+    let path = write_asset_file(input.user_id, &input.source, &public_id, ext, bytes).await?;
+    insert_asset(
+        state,
+        &input,
+        &public_id,
+        "completed",
+        Some(mime_type),
+        Some(&path),
+        bytes.len() as i64,
+        input.metadata.clone(),
+    )
+    .await
+}
+
 pub async fn record_image_assets_from_response(
     state: &AppState,
     user_id: Option<i64>,