RustでWebview - 闘うITエンジニアの覚え書き

Cargo.toml

[package]
name = "rust_webview_gui"
version = "0.1.0"
edition = "2021"

[dependencies]
wry = "0.53"
tao = "0.29"
libc = "0.2"
once_cell = "1.19"
ctrlc = "3.4"

[target.'cfg(target_os = "windows")'.dependencies]
windows-sys = { version = "0.52", features = [
    "Win32_System_JobObjects",
    "Win32_System_Threading",
    "Win32_Foundation",
    "Win32_Security",
] }

src/main.rs

use std::{
    env,
    net::TcpStream,
    path::PathBuf,
    process::{Child, Command, Stdio},
    sync::{Arc, Mutex},
    thread,
    time::Duration,
};
#[cfg(unix)]
use std::os::unix::process::CommandExt;
use wry::{Result, WebViewBuilder};
use tao::{
    event::{Event, WindowEvent},
    event_loop::{ControlFlow, EventLoop},
    window::WindowBuilder,
};
use once_cell::sync::OnceCell;

// 設定情報
struct Config {
    streamlit_path: String,
    server_host: String,
    server_port: u16,
    nav_url: String,
    probe_host: String,
}

#[cfg(target_os = "windows")]
mod windows_job {
    use once_cell::sync::OnceCell;
    use std::{mem::size_of, os::windows::io::AsRawHandle, process::Child};
    use windows_sys::Win32::{
        Foundation::{GetLastError, HANDLE},
        System::JobObjects::{
            AssignProcessToJobObject, CreateJobObjectW, SetInformationJobObject,
            TerminateJobObject, JOBOBJECT_EXTENDED_LIMIT_INFORMATION,
            JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, JobObjectExtendedLimitInformation,
        },
    };

    static JOB_HANDLE: OnceCell<HANDLE> = OnceCell::new();

    fn ensure_job() -> Option<HANDLE> {
        let handle = *JOB_HANDLE.get_or_init(|| unsafe {
            let job = CreateJobObjectW(std::ptr::null_mut(), std::ptr::null());
            if job == 0 {
                eprintln!(
                    "Failed to create Windows job object: {}",
                    GetLastError()
                );
                return 0;
            }

            let mut limits: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = std::mem::zeroed();
            limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
            let result = SetInformationJobObject(
                job,
                JobObjectExtendedLimitInformation,
                &limits as *const _ as *const _,
                size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
            );
            if result == 0 {
                eprintln!(
                    "Failed to configure Windows job object limits: {}",
                    GetLastError()
                );
            }

            job
        });

        (handle != 0).then_some(handle)
    }

    pub fn assign(child: &Child) {
        if let Some(job) = ensure_job() {
            let process_handle = child.as_raw_handle();
            if process_handle.is_null() {
                eprintln!("Streamlit child process has no handle to assign to job object");
                return;
            }

            unsafe {
                if AssignProcessToJobObject(job, process_handle as HANDLE) == 0 {
                    eprintln!(
                        "Failed to assign Streamlit process to Windows job object: {}",
                        GetLastError()
                    );
                }
            }
        }
    }

    pub fn terminate() {
        if let Some(&job) = JOB_HANDLE.get() {
            if job != 0 {
                unsafe {
                    if TerminateJobObject(job, 1) == 0 {
                        eprintln!(
                            "Failed to terminate Windows job object: {}",
                            GetLastError()
                        );
                    }
                }
            }
        }
    }
}

#[cfg(not(target_os = "windows"))]
mod windows_job {
    use std::process::Child;

    pub fn assign(_child: &Child) {}
    pub fn terminate() {}
}

static CHILD_HANDLE: OnceCell>> = OnceCell::new();

// Streamlit 起動待機の状態
enum ReadyState {
    Pending,
    Ready(String),
    Failed(String),
}

fn wait_for_port(host: &str, port: u16, timeout_secs: u64) -> bool {
    let start = std::time::Instant::now();
    let addr = format!("{}:{}", host, port);

    while start.elapsed().as_secs() < timeout_secs {
        if TcpStream::connect(&addr).is_ok() {
            println!("Streamlit server is ready at {}", addr);
            return true;
        }
        thread::sleep(Duration::from_millis(300));
    }

    println!("Streamlit did not start within {} seconds", timeout_secs);
    false
}

fn parse_config() -> Config {
    // コマンドライン引数・環境変数から設定を取得（streamlit 実行ファイル・ホスト・ポート・URL）
    let args: Vec<String> = env::args().collect();
    let mut arg_streamlit_path: Option<String> = None;
    let mut arg_host: Option<String> = None;
    let mut arg_port: Option<u16> = None;
    let mut arg_url: Option<String> = None;
    let mut i = 1;
    while i < args.len() {
        match args[i].as_str() {
            "--host" if i + 1 < args.len() => { arg_host = Some(args[i+1].clone()); i += 2; }
            "--port" if i + 1 < args.len() => { if let Ok(p) = args[i+1].parse::<u16>() { arg_port = Some(p); } i += 2; }
            "--url"  if i + 1 < args.len() => { arg_url = Some(args[i+1].clone()); i += 2; }
            s if s.starts_with('-') => { i += 1; }
            other => { if arg_streamlit_path.is_none() { arg_streamlit_path = Some(other.to_string()); } i += 1; }
        }
    }

    let streamlit_path = determine_streamlit_path(arg_streamlit_path);
    let env_host = env::var("STREAMLIT_SERVER_ADDRESS").ok().or_else(|| env::var("STREAMLIT_HOST").ok());
    let env_port = env::var("STREAMLIT_SERVER_PORT").ok().and_then(|v| v.parse::<u16>().ok())
        .or_else(|| env::var("STREAMLIT_PORT").ok().and_then(|v| v.parse::<u16>().ok()));
    let env_url = env::var("STREAMLIT_URL").ok();
    let server_host = arg_host.or(env_host).unwrap_or_else(|| "127.0.0.1".to_string());
    let server_port = arg_port.or(env_port).unwrap_or(8501u16);
    let default_url_host = if server_host == "0.0.0.0" { "127.0.0.1" } else { server_host.as_str() };
    let nav_url = arg_url.or(env_url).unwrap_or_else(|| format!("http://{}:{}", default_url_host, server_port));
    let probe_host = if server_host == "0.0.0.0" { "127.0.0.1".to_string() } else { server_host.clone() };
    println!("Using Streamlit executable: {} (host={}, port={}, url={})", streamlit_path, server_host, server_port, nav_url);

    Config { streamlit_path, server_host, server_port, nav_url, probe_host }
}

// streamlit 実行ファイルのパスを優先順位で解決
// 1) 引数で指定されたもの
// 2) 実行可能ファイルと同階層にある "run_app"
// 3) "../streamlit_app/dist/run_app"
fn determine_streamlit_path(arg_path: Option<String>) -> String {
    if let Some(p) = arg_path { return p; }

    // 2) 実行可能ファイルと同階層にある run_app / run_app.exe
    if let Ok(exe) = env::current_exe() {
        if let Some(dir) = exe.parent() {
            // 優先候補（プラットフォームに応じた順序）
            #[cfg(target_os = "windows")]
            let names = ["run_app.exe", "run_app"];
            #[cfg(not(target_os = "windows"))]
            let names = ["run_app", "run_app.exe"];

            for name in names { 
                let candidate = dir.join(name);
                if candidate.exists() {
                    return candidate.to_string_lossy().into_owned();
                }
            }
        }
    }

    // 3) 既定の相対パス（どちらも試す）
    #[cfg(target_os = "windows")]
    let defaults = ["../streamlit_app/dist/run_app.exe", "../streamlit_app/dist/run_app"];
    #[cfg(not(target_os = "windows"))]
    let defaults = ["../streamlit_app/dist/run_app", "../streamlit_app/dist/run_app.exe"];
    for d in defaults { 
        if PathBuf::from(d).exists() { 
            return d.to_string(); 
        }
    }

    // 最後のフォールバック
    #[cfg(target_os = "windows")]
    { return "../streamlit_app/dist/run_app.exe".to_string(); }
    #[cfg(not(target_os = "windows"))]
    { return "../streamlit_app/dist/run_app".to_string(); }
}

fn loading_html() -> &'static str {
    r#"
        <!DOCTYPE html>
        <html lang=\"ja\">
        <head>
          <meta charset=\"UTF-8\" />
          <meta name=\"viewport\" content=\"width=device-width, initial-scale=1.0\" />
          <title>起動中...</title>
          <style>
            html, body { height: 100%; margin: 0; font-family: system-ui, -apple-system, Segoe UI, Roboto, Ubuntu, Cantarell, Noto Sans, Helvetica, Arial, \"Apple Color Emoji\", \"Segoe UI Emoji\"; }
            .wrap { height: 100%; display: grid; place-items: center; background: #0f172a; color: #e2e8f0; }
            .box { text-align: center; }
            .spinner { width: 54px; height: 54px; border: 6px solid #334155; border-top-color: #60a5fa; border-radius: 50%; animation: spin 1s linear infinite; margin: 0 auto 14px; }
            @keyframes spin { to { transform: rotate(360deg); } }
            .msg { font-size: 14px; opacity: 0.9; }
          </style>
        </head>
        <body>
          <div class=\"wrap\">
            <div class=\"box\">
              <div class=\"spinner\"></div>
              <div class=\"msg\">現在、アプリケーションを起動しています...</div>
            </div>
          </div>
        </body>
        </html>
    "#
}

fn start_streamlit(cfg: &Config, state: &Arc<Mutex<ReadyState>>) {
    // Streamlit サブプロセス起動（ローディング表示後にすぐバックグラウンドで）
    let pathbuf = PathBuf::from(&cfg.streamlit_path);
    let abs_path = if pathbuf.is_absolute() {
        pathbuf
    } else {
        env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
            .join(pathbuf)
    };
    if !abs_path.exists() {
        if let Ok(mut st) = state.lock() {
            *st = ReadyState::Failed(format!("Streamlit 実行ファイルが見つかりません: {:?}", abs_path));
        }
        return;
    }

    let mut cmd = Command::new(&abs_path);
    if let Some(dir) = abs_path.parent() { cmd.current_dir(dir); }
    // ログを確認できるように親の標準入出力を継承
    cmd.stdout(Stdio::inherit()).stderr(Stdio::inherit());
    // Streamlit のホスト・ポートを環境変数で渡す（run_app.py 側で参照）
    cmd.env("STREAMLIT_SERVER_ADDRESS", &cfg.server_host)
        .env("STREAMLIT_SERVER_PORT", cfg.server_port.to_string())
        .env("STREAMLIT_HOST", &cfg.server_host)
        .env("STREAMLIT_PORT", cfg.server_port.to_string());
    println!("Starting Streamlit: {:?} (cwd: {:?})", &abs_path, abs_path.parent());
    #[cfg(unix)]
    unsafe {
        // 子プロセスを新しいプロセスグループに入れる
        cmd.pre_exec(|| {
            let _ = libc::setpgid(0, 0);
            Ok(())
        });
    }
    match cmd.spawn() {
        Ok(child_spawned) => {
            windows_job::assign(&child_spawned);
            let child_arc = Arc::new(Mutex::new(child_spawned));
            let _ = CHILD_HANDLE.set(Arc::clone(&child_arc));

            // Ctrl+C (SIGINT) 対応: 子プロセス（PG）を終了してからプロセス終了
            ctrlc::set_handler(|| {
                if let Some(c) = CHILD_HANDLE.get() {
                    terminate_child(c);
                }
                std::process::exit(0);
            }).expect("Failed to set Ctrl-C handler");

            // パニック時のクリーンアップ
            let default_hook = std::panic::take_hook();
            std::panic::set_hook(Box::new(move |info| {
                if let Some(c) = CHILD_HANDLE.get() {
                    terminate_child(c);
                }
                default_hook(info);
            }));

            // 別スレッドで起動待機して状態を更新
            let state_clone = Arc::clone(state);
            let nav_url_clone = cfg.nav_url.clone();
            let probe_host_clone = cfg.probe_host.clone();
            let port = cfg.server_port;
            thread::spawn(move || {
                if wait_for_port(&probe_host_clone, port, 120) {
                    if let Ok(mut st) = state_clone.lock() {
                        *st = ReadyState::Ready(nav_url_clone);
                    }
                } else {
                    if let Ok(mut st) = state_clone.lock() {
                        *st = ReadyState::Failed("Streamlit が起動しませんでした".to_string());
                    }
                }
            });

            // 子プロセスの早期終了を監視（すぐに失敗表示へ切替）
            let state_clone2 = Arc::clone(state);
            thread::spawn(move || {
                if let Some(ch) = CHILD_HANDLE.get() {
                    loop {
                        // 既に結果が出ていれば終了
                        if let Ok(st) = state_clone2.lock() {
                            if !matches!(*st, ReadyState::Pending) { break; }
                        }
                        if let Ok(mut child) = ch.lock() {
                            if let Ok(Some(status)) = child.try_wait() {
                                let msg = if let Some(code) = status.code() {
                                    format!("Streamlit プロセスが終了しました (exit code: {})", code)
                                } else {
                                    "Streamlit プロセスが終了しました".to_string()
                                };
                                drop(child);
                                if let Ok(mut st) = state_clone2.lock() {
                                    if matches!(*st, ReadyState::Pending) {
                                        *st = ReadyState::Failed(msg);
                                    }
                                }
                                break;
                            }
                        }
                        thread::sleep(Duration::from_millis(200));
                    }
                }
            });
        }
        Err(err) => {
            if let Ok(mut st) = state.lock() {
                *st = ReadyState::Failed(format!("Streamlit 起動に失敗しました: {}", err));
            }
        }
    }
}

fn apply_state_in_webview(webview: &Option<wry::WebView>, state: &Arc<Mutex<ReadyState>>) {
    if let Some(ref wv) = webview {
        if let Ok(mut st) = state.lock() {
            match std::mem::replace(&mut *st, ReadyState::Pending) {
                ReadyState::Ready(url) => {
                    let js = format!("window.location.replace('{}');", js_escape_single(&url));
                    let _ = wv.evaluate_script(&js);
                }
                ReadyState::Failed(msg) => {
                    let html = format!(
                        "<html><body style='display:grid;place-items:center;height:100%;background:#1f2937;color:#f3f4f6;font-family:system-ui'>\
                        <div style='text-align:center'><div style='font-size:18px;margin-bottom:8px'>起動に失敗しました</div>\
                        <div style='opacity:.85'>{}</div></div></body></html>",
                        msg.replace('\n', "<br/>")
                    );
                    let data_url = format!("data:text/html;charset=utf-8,{}", url_encode(&html));
                    let js = format!("window.location.replace('{}');", js_escape_single(&data_url));
                    let _ = wv.evaluate_script(&js);
                }
                ReadyState::Pending => {}
            }
        }
    }
}

fn run_event_loop(
    event_loop: EventLoop<()>,
    window: tao::window::Window,
    mut webview: Option<wry::WebView>,
    state: Arc<Mutex<ReadyState>>,
) {
    // window を保持しておく（破棄されないようにスコープに保持）
    let _window = window;
    event_loop.run(move |event, _, control_flow| {
        *control_flow = ControlFlow::Wait;
        match event {
            Event::Resumed => {}
            Event::WindowEvent { event: WindowEvent::CloseRequested, .. } => {
                println!("Window closed. Killing Streamlit subprocess...");
                if let Some(wv) = webview.take() { drop(wv); }
                if let Some(c) = CHILD_HANDLE.get() { terminate_child(c); }
                *control_flow = ControlFlow::Exit;
            }
            Event::LoopDestroyed => {
                if let Some(wv) = webview.take() { drop(wv); }
                if let Some(c) = CHILD_HANDLE.get() { terminate_child(c); }
            }
            Event::MainEventsCleared => {
                apply_state_in_webview(&webview, &state);
            }
            _ => (),
        }
    });
}

fn main() -> Result<()> {
    // 設定のパース
    let cfg = parse_config();

    // EventLoop と Window 作成
    let event_loop: EventLoop<()> = EventLoop::new();
    let window = WindowBuilder::new()
        .with_title("Streamlit WebView")
        .build(&event_loop)
        .expect("Failed to build window");

    // まずはローディングページを表示
    let webview = Some(
        WebViewBuilder::new()
            .with_html(loading_html())
            .build(&window)?,
    );

    // 状態の保持と子プロセス起動
    let state = Arc::new(Mutex::new(ReadyState::Pending));
    start_streamlit(&cfg, &state);

    // イベントループ実行
    run_event_loop(event_loop, window, webview, state);
    #[allow(unreachable_code)]
    Ok(())
}

fn terminate_child(child_arc: &Arc<Mutex<Child>>) {
    let mut child = match child_arc.lock() {
        Ok(guard) => guard,
        Err(_) => return,
    };

    // すでに終了していれば何もしない
    if let Ok(Some(_)) = child.try_wait() {
        return;
    }

    #[cfg(unix)]
    unsafe {
        // プロセスグループ全体にSIGTERM → 待機 → SIGKILL
        let pgid = child.id() as i32; // setpgid(0,0)により pid == pgid
        let _ = libc::killpg(pgid, libc::SIGTERM);
        // 短い猶予
        for _ in 0..10 {
            if let Ok(Some(_)) = child.try_wait() {
                break;
            }
            thread::sleep(Duration::from_millis(100));
        }
        if let Ok(None) = child.try_wait() {
            let _ = libc::killpg(pgid, libc::SIGKILL);
        }
    }

    #[cfg(not(unix))]
    {
        #[cfg(target_os = "windows")]
        {
            windows_job::terminate();
        }
        let _ = child.kill();
    }

    // ゾンビ防止のためwaitで回収
    let _ = child.wait();
}

// URL エンコード（data URL 用に簡易実装）
fn url_encode(s: &str) -> String {
    fn is_unreserved(b: u8) -> bool {
        matches!(b, b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~')
    }
    let mut out = String::with_capacity(s.len() * 3);
    for &b in s.as_bytes() {
        if is_unreserved(b) {
            out.push(b as char);
        } else {
            out.push('%');
            out.push_str(&format!("{:02X}", b));
        }
    }
    out
}

// JS のシングルクォート文字列用に最低限のエスケープ
fn js_escape_single(s: &str) -> String {
    s.replace('\\', "\\\\").replace('\'', "\\'").replace('\n', "\\n")
}