use std::collections::HashMap;
use std::process::Command;
use crate::error::{AXError, AXResult};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum BrowserType {
Chromium,
Firefox,
Brave,
Edge,
}
impl BrowserType {
#[must_use]
pub fn image_suffix(self) -> &'static str {
match self {
Self::Chromium => "chromium",
Self::Firefox => "firefox",
Self::Brave => "brave",
Self::Edge => "microsoft-edge",
}
}
#[must_use]
pub fn docker_image(self) -> String {
format!("ghcr.io/m1k1o/neko/{}:latest", self.image_suffix())
}
#[must_use]
pub fn display_name(self) -> &'static str {
match self {
Self::Chromium => "Chromium",
Self::Firefox => "Firefox",
Self::Brave => "Brave",
Self::Edge => "Microsoft Edge",
}
}
#[must_use]
pub fn supports_full_cdp(self) -> bool {
matches!(self, Self::Chromium | Self::Edge)
}
}
#[derive(Debug, Clone)]
pub struct NekoConfig {
pub browser: BrowserType,
pub width: u32,
pub height: u32,
pub cdp_port: u16,
pub vnc_port: u16,
pub admin_password: String,
pub name_prefix: String,
}
impl NekoConfig {
#[must_use]
pub fn chromium() -> Self {
Self::for_browser(BrowserType::Chromium)
}
#[must_use]
pub fn firefox() -> Self {
Self::for_browser(BrowserType::Firefox)
}
#[must_use]
pub fn brave() -> Self {
Self::for_browser(BrowserType::Brave)
}
#[must_use]
pub fn edge() -> Self {
Self::for_browser(BrowserType::Edge)
}
#[must_use]
pub fn builder() -> NekoConfigBuilder {
NekoConfigBuilder::default()
}
fn for_browser(browser: BrowserType) -> Self {
let base_cdp = 9222u16;
let base_vnc = 5900u16;
Self {
browser,
width: 1920,
height: 1080,
cdp_port: base_cdp,
vnc_port: base_vnc,
admin_password: "admin".into(),
name_prefix: "neko".into(),
}
}
}
#[derive(Debug, Default)]
pub struct NekoConfigBuilder {
browser: Option<BrowserType>,
width: Option<u32>,
height: Option<u32>,
cdp_port: Option<u16>,
vnc_port: Option<u16>,
admin_password: Option<String>,
name_prefix: Option<String>,
}
impl NekoConfigBuilder {
#[must_use]
pub fn browser(mut self, browser: BrowserType) -> Self {
self.browser = Some(browser);
self
}
#[must_use]
pub fn dimensions(mut self, width: u32, height: u32) -> Self {
self.width = Some(width);
self.height = Some(height);
self
}
#[must_use]
pub fn cdp_port(mut self, port: u16) -> Self {
self.cdp_port = Some(port);
self
}
#[must_use]
pub fn vnc_port(mut self, port: u16) -> Self {
self.vnc_port = Some(port);
self
}
#[must_use]
pub fn admin_password(mut self, password: impl Into<String>) -> Self {
self.admin_password = Some(password.into());
self
}
#[must_use]
pub fn name_prefix(mut self, prefix: impl Into<String>) -> Self {
self.name_prefix = Some(prefix.into());
self
}
#[must_use]
pub fn build(self) -> NekoConfig {
let base = NekoConfig::for_browser(self.browser.unwrap_or(BrowserType::Chromium));
NekoConfig {
browser: self.browser.unwrap_or(base.browser),
width: self.width.unwrap_or(base.width),
height: self.height.unwrap_or(base.height),
cdp_port: self.cdp_port.unwrap_or(base.cdp_port),
vnc_port: self.vnc_port.unwrap_or(base.vnc_port),
admin_password: self.admin_password.unwrap_or(base.admin_password),
name_prefix: self.name_prefix.unwrap_or(base.name_prefix),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NekoBrowser {
pub(crate) container_id: String,
pub(crate) cdp_port: u16,
pub(crate) vnc_port: u16,
pub(crate) browser: BrowserType,
}
impl NekoBrowser {
#[must_use]
pub fn cdp_url(&self) -> String {
format!("ws://127.0.0.1:{}/devtools/browser", self.cdp_port)
}
#[must_use]
pub fn vnc_addr(&self) -> String {
format!("127.0.0.1:{}", self.vnc_port)
}
#[must_use]
pub fn browser(&self) -> BrowserType {
self.browser
}
#[must_use]
pub fn container_id(&self) -> &str {
&self.container_id
}
pub fn screenshot(&self) -> AXResult<Vec<u8>> {
self.screenshot_via(&RealCdpShooter)
}
#[cfg(test)]
#[must_use]
pub fn new_for_test(
container_id: &str,
cdp_port: u16,
vnc_port: u16,
browser: BrowserType,
) -> Self {
Self {
container_id: container_id.to_string(),
cdp_port,
vnc_port,
browser,
}
}
pub(crate) fn screenshot_via(&self, shooter: &dyn ScreenshotShooter) -> AXResult<Vec<u8>> {
shooter.capture(self.cdp_port)
}
}
pub(crate) trait ScreenshotShooter {
fn capture(&self, cdp_port: u16) -> AXResult<Vec<u8>>;
}
struct RealCdpShooter;
impl ScreenshotShooter for RealCdpShooter {
fn capture(&self, cdp_port: u16) -> AXResult<Vec<u8>> {
use crate::electron_cdp::probe_cdp_port;
use tungstenite::connect as ws_connect;
if !probe_cdp_port(cdp_port) {
return Err(AXError::AppNotFound(format!(
"No CDP endpoint on port {cdp_port}"
)));
}
let ws_url = format!("ws://127.0.0.1:{cdp_port}/devtools/browser");
let (mut socket, _) =
ws_connect(&ws_url).map_err(|e| AXError::SystemError(format!("CDP connect: {e}")))?;
let request = serde_json::json!({
"id": 1,
"method": "Page.captureScreenshot",
"params": { "format": "png" }
});
socket
.send(tungstenite::Message::Text(request.to_string().into()))
.map_err(|e| AXError::SystemError(format!("CDP send: {e}")))?;
loop {
let msg = socket
.read()
.map_err(|e| AXError::SystemError(format!("CDP read: {e}")))?;
if let tungstenite::Message::Text(text) = msg {
let resp: serde_json::Value = serde_json::from_str(&text)
.map_err(|e| AXError::SystemError(format!("CDP parse: {e}")))?;
if resp["id"].as_u64() == Some(1) {
let b64 = resp["result"]["data"]
.as_str()
.ok_or_else(|| AXError::SystemError("No screenshot data".into()))?;
return base64_decode(b64);
}
}
}
}
}
fn base64_decode(input: &str) -> AXResult<Vec<u8>> {
const TABLE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
let mut lookup = [255u8; 256];
for (i, &c) in TABLE.iter().enumerate() {
lookup[c as usize] = i as u8;
}
let clean: Vec<u8> = input.bytes().filter(|&b| b != b'=').collect();
let mut out = Vec::with_capacity(clean.len() * 3 / 4);
for chunk in clean.chunks(4) {
let vals: Vec<u8> = chunk.iter().map(|&b| lookup[b as usize]).collect();
if vals.contains(&255) {
return Err(AXError::SystemError("Invalid base64 character".into()));
}
match vals.as_slice() {
[a, b, c, d] => {
out.push((a << 2) | (b >> 4));
out.push((b << 4) | (c >> 2));
out.push((c << 6) | d);
}
[a, b, c] => {
out.push((a << 2) | (b >> 4));
out.push((b << 4) | (c >> 2));
}
[a, b] => {
out.push((a << 2) | (b >> 4));
}
_ => {}
}
}
Ok(out)
}
pub(crate) trait DockerRunner {
fn run_container(&self, args: &[&str]) -> AXResult<String>;
fn stop_container(&self, container_id: &str) -> AXResult<()>;
fn rm_container(&self, container_id: &str) -> AXResult<()>;
fn list_neko_containers(&self) -> AXResult<Vec<String>>;
}
struct RealDockerRunner;
impl DockerRunner for RealDockerRunner {
fn run_container(&self, args: &[&str]) -> AXResult<String> {
let output = Command::new("docker")
.arg("run")
.args(args)
.output()
.map_err(|e| AXError::SystemError(format!("docker run: {e}")))?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
return Err(AXError::SystemError(format!("docker run failed: {stderr}")));
}
Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
}
fn stop_container(&self, container_id: &str) -> AXResult<()> {
let status = Command::new("docker")
.args(["stop", container_id])
.status()
.map_err(|e| AXError::SystemError(format!("docker stop: {e}")))?;
status
.success()
.then_some(())
.ok_or_else(|| AXError::SystemError(format!("docker stop {container_id} failed")))
}
fn rm_container(&self, container_id: &str) -> AXResult<()> {
let status = Command::new("docker")
.args(["rm", "-f", container_id])
.status()
.map_err(|e| AXError::SystemError(format!("docker rm: {e}")))?;
status
.success()
.then_some(())
.ok_or_else(|| AXError::SystemError(format!("docker rm {container_id} failed")))
}
fn list_neko_containers(&self) -> AXResult<Vec<String>> {
let output = Command::new("docker")
.args(["ps", "-a", "-q", "--filter", "label=axterminator.neko=1"])
.output()
.map_err(|e| AXError::SystemError(format!("docker ps: {e}")))?;
let ids = String::from_utf8_lossy(&output.stdout)
.lines()
.filter(|l| !l.is_empty())
.map(str::to_string)
.collect();
Ok(ids)
}
}
pub struct DockerManager {
runner: Box<dyn DockerRunner>,
active: HashMap<String, NekoBrowser>,
}
impl std::fmt::Debug for DockerManager {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DockerManager")
.field("active_count", &self.active.len())
.finish()
}
}
impl DockerManager {
#[must_use]
pub fn new() -> Self {
Self::with_runner(Box::new(RealDockerRunner))
}
#[must_use]
pub(crate) fn with_runner(runner: Box<dyn DockerRunner>) -> Self {
Self {
runner,
active: HashMap::new(),
}
}
pub fn launch(&mut self, config: NekoConfig) -> AXResult<NekoBrowser> {
let container_name = self.unique_name(&config);
let args = build_docker_args(&config, &container_name);
let args_refs: Vec<&str> = args.iter().map(String::as_str).collect();
let container_id = self.runner.run_container(&args_refs)?;
let browser = NekoBrowser {
container_id: container_id.clone(),
cdp_port: config.cdp_port,
vnc_port: config.vnc_port,
browser: config.browser,
};
self.active.insert(container_id, browser.clone());
Ok(browser)
}
pub fn stop(&mut self, browser: &NekoBrowser) -> AXResult<()> {
self.runner.stop_container(&browser.container_id)?;
self.runner.rm_container(&browser.container_id)?;
self.active.remove(&browser.container_id);
Ok(())
}
pub fn cleanup(&mut self) -> usize {
let tracked: Vec<String> = self.active.keys().cloned().collect();
let discovered = self.runner.list_neko_containers().unwrap_or_default();
let mut ids: Vec<String> = tracked;
for id in discovered {
if !ids.contains(&id) {
ids.push(id);
}
}
let count = ids.len();
for id in &ids {
let _ = self.runner.stop_container(id);
let _ = self.runner.rm_container(id);
}
self.active.clear();
count
}
#[must_use]
pub fn active_count(&self) -> usize {
self.active.len()
}
fn unique_name(&self, config: &NekoConfig) -> String {
format!(
"{}-{}-{}",
config.name_prefix,
config.browser.image_suffix(),
config.cdp_port,
)
}
}
impl Default for DockerManager {
fn default() -> Self {
Self::new()
}
}
fn build_docker_args(config: &NekoConfig, name: &str) -> Vec<String> {
vec![
"--detach".into(),
"--name".into(),
name.to_string(),
"--label".into(),
"axterminator.neko=1".into(),
"--publish".into(),
format!("{}:9222", config.cdp_port),
"--publish".into(),
format!("{}:5900", config.vnc_port),
"--shm-size".into(),
"2g".into(),
"--env".into(),
format!("NEKO_SCREEN={}x{}@30", config.width, config.height),
"--env".into(),
format!("NEKO_PASSWORD_ADMIN={}", config.admin_password),
"--env".into(),
"NEKO_REMOTE_DEBUGGING_PORT=9222".into(),
config.browser.docker_image(),
]
}
#[cfg(test)]
mod tests {
use super::*;
struct MockDockerRunner {
next_id: std::cell::Cell<u64>,
fail_launch: bool,
discoverable: Vec<String>,
}
impl MockDockerRunner {
fn new() -> Self {
Self {
next_id: std::cell::Cell::new(1),
fail_launch: false,
discoverable: vec![],
}
}
fn failing() -> Self {
Self {
fail_launch: true,
..Self::new()
}
}
fn with_discoverable(ids: Vec<String>) -> Self {
Self {
discoverable: ids,
..Self::new()
}
}
}
impl DockerRunner for MockDockerRunner {
fn run_container(&self, _args: &[&str]) -> AXResult<String> {
if self.fail_launch {
return Err(AXError::SystemError("docker daemon not running".into()));
}
let id = self.next_id.get();
self.next_id.set(id + 1);
Ok(format!("mock-container-{id:04}"))
}
fn stop_container(&self, _container_id: &str) -> AXResult<()> {
Ok(())
}
fn rm_container(&self, _container_id: &str) -> AXResult<()> {
Ok(())
}
fn list_neko_containers(&self) -> AXResult<Vec<String>> {
Ok(self.discoverable.clone())
}
}
struct MockCdpShooter {
pixels: Vec<u8>,
}
impl MockCdpShooter {
fn new(pixels: Vec<u8>) -> Self {
Self { pixels }
}
}
impl ScreenshotShooter for MockCdpShooter {
fn capture(&self, _cdp_port: u16) -> AXResult<Vec<u8>> {
Ok(self.pixels.clone())
}
}
#[test]
fn browser_type_image_suffix_maps_correctly() {
assert_eq!(BrowserType::Chromium.image_suffix(), "chromium");
assert_eq!(BrowserType::Firefox.image_suffix(), "firefox");
assert_eq!(BrowserType::Brave.image_suffix(), "brave");
assert_eq!(BrowserType::Edge.image_suffix(), "microsoft-edge");
}
#[test]
fn browser_type_docker_image_uses_ghcr_prefix() {
let browser = BrowserType::Chromium;
let image = browser.docker_image();
assert_eq!(image, "ghcr.io/m1k1o/neko/chromium:latest");
assert!(image.starts_with("ghcr.io/m1k1o/neko/"));
}
#[test]
fn browser_type_cdp_support_chromium_and_edge_only() {
assert!(BrowserType::Chromium.supports_full_cdp());
assert!(BrowserType::Edge.supports_full_cdp());
assert!(!BrowserType::Firefox.supports_full_cdp());
assert!(!BrowserType::Brave.supports_full_cdp());
}
#[test]
fn neko_config_chromium_defaults_sensible() {
let cfg = NekoConfig::chromium();
assert_eq!(cfg.browser, BrowserType::Chromium);
assert_eq!(cfg.width, 1920);
assert_eq!(cfg.height, 1080);
assert_eq!(cfg.cdp_port, 9222);
assert_eq!(cfg.vnc_port, 5900);
assert_eq!(cfg.admin_password, "admin");
assert_eq!(cfg.name_prefix, "neko");
}
#[test]
fn neko_config_builder_overrides_browser_and_ports() {
let cfg = NekoConfig::builder()
.browser(BrowserType::Firefox)
.cdp_port(9333)
.vnc_port(5910)
.build();
assert_eq!(cfg.browser, BrowserType::Firefox);
assert_eq!(cfg.cdp_port, 9333);
assert_eq!(cfg.vnc_port, 5910);
assert_eq!(cfg.width, 1920);
assert_eq!(cfg.height, 1080);
}
#[test]
fn neko_config_builder_full_customisation() {
let cfg = NekoConfig::builder()
.browser(BrowserType::Brave)
.dimensions(1280, 720)
.cdp_port(9400)
.vnc_port(5920)
.admin_password("s3cr3t")
.name_prefix("ci-neko")
.build();
assert_eq!(cfg.browser, BrowserType::Brave);
assert_eq!(cfg.width, 1280);
assert_eq!(cfg.height, 720);
assert_eq!(cfg.cdp_port, 9400);
assert_eq!(cfg.vnc_port, 5920);
assert_eq!(cfg.admin_password, "s3cr3t");
assert_eq!(cfg.name_prefix, "ci-neko");
}
#[test]
fn neko_browser_cdp_url_format_correct() {
let browser = NekoBrowser::new_for_test("abc123", 9222, 5900, BrowserType::Chromium);
let url = browser.cdp_url();
assert_eq!(url, "ws://127.0.0.1:9222/devtools/browser");
}
#[test]
fn neko_browser_cdp_url_reflects_custom_port() {
let browser = NekoBrowser::new_for_test("def456", 9400, 5920, BrowserType::Firefox);
let url = browser.cdp_url();
assert_eq!(url, "ws://127.0.0.1:9400/devtools/browser");
}
#[test]
fn neko_browser_vnc_addr_format_correct() {
let browser = NekoBrowser::new_for_test("ghi789", 9222, 5901, BrowserType::Chromium);
assert_eq!(browser.vnc_addr(), "127.0.0.1:5901");
}
#[test]
fn neko_browser_screenshot_via_mock_returns_expected_bytes() {
let browser = NekoBrowser::new_for_test("mock-id", 9222, 5900, BrowserType::Chromium);
let expected = vec![0x89u8, 0x50, 0x4E, 0x47]; let shooter = MockCdpShooter::new(expected.clone());
let result = browser.screenshot_via(&shooter);
assert!(result.is_ok());
assert_eq!(result.unwrap(), expected);
}
#[test]
fn docker_manager_launch_tracks_container() {
let mut mgr = DockerManager::with_runner(Box::new(MockDockerRunner::new()));
let browser = mgr.launch(NekoConfig::chromium()).unwrap();
assert_eq!(mgr.active_count(), 1);
assert!(!browser.container_id().is_empty());
assert_eq!(browser.cdp_port, 9222);
assert_eq!(browser.browser(), BrowserType::Chromium);
}
#[test]
fn docker_manager_stop_removes_from_tracking() {
let mut mgr = DockerManager::with_runner(Box::new(MockDockerRunner::new()));
let browser = mgr.launch(NekoConfig::chromium()).unwrap();
assert_eq!(mgr.active_count(), 1);
mgr.stop(&browser).unwrap();
assert_eq!(mgr.active_count(), 0);
}
#[test]
fn docker_manager_multiple_browsers_tracked_independently() {
let mut mgr = DockerManager::with_runner(Box::new(MockDockerRunner::new()));
let chrome = mgr.launch(NekoConfig::chromium()).unwrap();
let firefox = mgr.launch(NekoConfig::firefox()).unwrap();
let brave = mgr.launch(NekoConfig::brave()).unwrap();
assert_eq!(mgr.active_count(), 3);
let ids = [
chrome.container_id(),
firefox.container_id(),
brave.container_id(),
];
assert_ne!(ids[0], ids[1]);
assert_ne!(ids[1], ids[2]);
assert_ne!(ids[0], ids[2]);
}
#[test]
fn docker_manager_cleanup_removes_all_tracked_containers() {
let mut mgr = DockerManager::with_runner(Box::new(MockDockerRunner::new()));
mgr.launch(NekoConfig::chromium()).unwrap();
mgr.launch(NekoConfig::firefox()).unwrap();
assert_eq!(mgr.active_count(), 2);
let removed = mgr.cleanup();
assert_eq!(removed, 2);
assert_eq!(mgr.active_count(), 0);
}
#[test]
fn docker_manager_cleanup_also_removes_discovered_orphans() {
let runner =
MockDockerRunner::with_discoverable(vec!["orphan-001".into(), "orphan-002".into()]);
let mut mgr = DockerManager::with_runner(Box::new(runner));
let removed = mgr.cleanup();
assert_eq!(removed, 2);
}
#[test]
fn docker_manager_launch_failure_returns_error() {
let mut mgr = DockerManager::with_runner(Box::new(MockDockerRunner::failing()));
let result = mgr.launch(NekoConfig::chromium());
assert!(result.is_err());
assert_eq!(mgr.active_count(), 0);
let msg = result.unwrap_err().to_string();
assert!(msg.contains("docker daemon"));
}
#[test]
fn build_docker_args_includes_required_flags() {
let cfg = NekoConfig::chromium();
let args = build_docker_args(&cfg, "test-neko-chromium-9222");
assert!(args.contains(&"--detach".into()));
assert!(args.contains(&"--shm-size".into()));
assert!(args.contains(&"2g".into()));
assert!(args.contains(&"--label".into()));
assert!(args.contains(&"axterminator.neko=1".into()));
assert!(args.iter().any(|a| a.contains("9222:9222")));
assert!(args.iter().any(|a| a.contains("5900:5900")));
assert!(args
.iter()
.any(|a| a.contains("ghcr.io/m1k1o/neko/chromium")));
}
#[test]
fn build_docker_args_respects_custom_ports_and_resolution() {
let cfg = NekoConfig::builder()
.browser(BrowserType::Firefox)
.cdp_port(9500)
.vnc_port(5950)
.dimensions(1280, 720)
.build();
let args = build_docker_args(&cfg, "neko-firefox-9500");
assert!(args.iter().any(|a| a.contains("9500:9222")));
assert!(args.iter().any(|a| a.contains("5950:5900")));
assert!(args.iter().any(|a| a.contains("1280x720")));
assert!(args
.iter()
.any(|a| a.contains("ghcr.io/m1k1o/neko/firefox")));
}
#[test]
fn base64_decode_round_trips_hello() {
let encoded = "SGVsbG8=";
let result = base64_decode(encoded).unwrap();
assert_eq!(result, b"Hello");
}
#[test]
fn base64_decode_rejects_invalid_characters() {
let result = base64_decode("SGVs!G8=");
assert!(result.is_err());
}
}