mod catalog;
mod composition;
mod dev;
use anyhow::{Context, Result, bail};
use clap::{Parser, Subcommand};
use omochi_transport::{
pairing::Invite,
store::{Role, Store},
};
use std::{
os::unix::fs::{MetadataExt, PermissionsExt},
path::{Path, PathBuf},
time::Duration,
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
const DEFAULT_WEB_URL: &str = "https://app.omochi.fyi/";
#[derive(Parser)]
#[command(
name = "omochi-ui",
version,
about = "Attach browser UIs to local and paired applications"
)]
struct Cli {
#[arg(long, global = true, env = "OMOCHI_CONFIG_DIR")]
config_dir: Option<PathBuf>,
#[arg(long, global = true, env = "OMOCHI_RUNTIME_ROOT")]
runtime_dir: Option<PathBuf>,
#[command(subcommand)]
command: Command,
}
#[derive(Subcommand)]
enum Command {
Status {
#[arg(long, default_value = "127.0.0.1")]
host: std::net::IpAddr,
#[arg(long, default_value_t = 5555)]
port: u16,
},
Dev {
#[arg(long, default_value = "127.0.0.1")]
host: std::net::IpAddr,
#[arg(long, default_value_t = 5555)]
port: u16,
#[arg(last = true, required = true)]
command: Vec<String>,
},
Generate {
#[command(subcommand)]
command: GenerateCommand,
},
Gateway {
#[arg(long)]
remote: bool,
#[arg(long, requires = "remote")]
pair: bool,
#[arg(long, requires = "remote", default_value = DEFAULT_WEB_URL, value_parser = validate_web_url)]
web_url: String,
#[arg(long)]
offline: bool,
#[arg(long)]
relay: Vec<String>,
#[arg(long, default_value = "gateway")]
name: String,
},
Hub {
#[arg(long, default_value = "127.0.0.1")]
host: std::net::IpAddr,
#[arg(long, default_value_t = 5555)]
port: u16,
#[arg(long)]
allow_origin: Vec<String>,
#[arg(long)]
no_auto_gateway: bool,
#[arg(long)]
assets: Option<PathBuf>,
#[arg(long, conflicts_with = "offline")]
remote: bool,
#[arg(long, requires = "remote", default_value = DEFAULT_WEB_URL, value_parser = validate_web_url)]
web_url: String,
#[arg(long)]
offline: bool,
#[arg(long)]
relay: Vec<String>,
},
Up {
#[arg(default_value = "omochi.toml")]
file: PathBuf,
#[arg(long, default_value = "127.0.0.1")]
host: std::net::IpAddr,
#[arg(long, default_value_t = 5555)]
port: u16,
#[arg(long)]
allow_origin: Vec<String>,
#[arg(long)]
open: bool,
},
Pair {
invitation: Option<String>,
#[arg(long)]
name: String,
#[arg(long, default_value = "hub")]
hub_name: String,
#[arg(long)]
offline: bool,
#[arg(long)]
relay: Vec<String>,
},
Id {
#[arg(long)]
hub: bool,
},
Trust {
#[command(subcommand)]
command: TrustCommand,
},
}
#[derive(Subcommand)]
enum GenerateCommand {
Catalog(catalog::Options),
}
#[derive(Subcommand)]
enum TrustCommand {
List,
Forget {
peer: String,
},
}
fn config_root(explicit: Option<PathBuf>) -> Result<PathBuf> {
if let Some(path) = explicit {
return Ok(path);
}
let base = std::env::var_os("XDG_CONFIG_HOME")
.map(PathBuf::from)
.or_else(|| std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".config")))
.context("set XDG_CONFIG_HOME or --config-dir")?;
Ok(base.join("omochi/rust-v1"))
}
fn validate_web_url(text: &str) -> std::result::Result<String, String> {
let invalid = || {
"--web-url requires an absolute HTTPS URL without credentials, query or fragment (maximum 2048 characters)".to_owned()
};
let parsed = url::Url::parse(text).map_err(|_| invalid())?;
if text.chars().count() > 2048
|| text.trim() != text
|| !text.to_ascii_lowercase().starts_with("https://")
|| text
.split('/')
.nth(2)
.is_some_and(|host| host.contains('@'))
|| parsed.scheme() != "https"
|| parsed.host_str().is_none()
|| !parsed.username().is_empty()
|| parsed.password().is_some()
|| parsed.query().is_some()
|| parsed.fragment().is_some()
{
return Err(invalid());
}
Ok(text.to_owned())
}
fn web_invitation(base: &str, invitation: &str) -> Result<String> {
let payload = invitation
.strip_prefix("omochi://invite/")
.context("invalid invitation link")?;
Ok(format!("{base}#g:{payload}"))
}
fn private_directory(path: &Path) -> Result<()> {
if !path.exists() {
std::fs::create_dir_all(path)?;
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700))?;
}
let meta = std::fs::symlink_metadata(path)?;
unsafe extern "C" {
fn getuid() -> u32;
}
if !meta.is_dir() || meta.mode() & 0o077 != 0 || meta.uid() != unsafe { getuid() } {
bail!(
"runtime directory must be owned and private: {}",
path.display()
)
}
Ok(())
}
struct ShutdownSignal {
term: tokio::signal::unix::Signal,
interrupt: tokio::signal::unix::Signal,
}
impl ShutdownSignal {
fn new() -> Result<Self> {
use tokio::signal::unix::{SignalKind, signal};
Ok(Self {
term: signal(SignalKind::terminate())?,
interrupt: signal(SignalKind::interrupt())?,
})
}
async fn wait(&mut self) {
tokio::select! { _ = self.interrupt.recv() => {}, _ = self.term.recv() => {} }
}
}
async fn probe(path: &Path) -> Result<()> {
use omochi_protocol::{OpenStream, VERSION, read_frame, session::Control, write_frame};
tokio::time::timeout(Duration::from_secs(2), async {
let mut socket = tokio::net::UnixStream::connect(path).await?;
write_frame(&mut socket, &OpenStream::Control { protocol: VERSION }).await?;
match read_frame::<_, Control>(&mut socket).await? {
Some(Control::Hello {
protocol: VERSION, ..
}) => Ok(()),
_ => bail!("incompatible gateway; stop the old gateway before starting this version"),
}
})
.await
.context("gateway did not respond; stop the old gateway before retrying")?
}
#[tokio::main]
async fn main() {
if let Err(error) = run(Cli::parse()).await {
eprintln!("[omochi-ui] {error:#}");
std::process::exit(1);
}
}
async fn run(cli: Cli) -> Result<()> {
if let Command::Generate {
command: GenerateCommand::Catalog(options),
} = cli.command
{
return catalog::generate(options);
}
if let Command::Status { host, port } = cli.command {
let _ = rustls::crypto::ring::default_provider().install_default();
let address = std::net::SocketAddr::new(host, port);
let status: serde_json::Value = reqwest::Client::builder()
.no_proxy()
.redirect(reqwest::redirect::Policy::none())
.timeout(Duration::from_secs(5))
.build()?
.get(format!("http://{address}/api/connections"))
.send()
.await
.context("cannot reach the hub; start omochi-ui hub first")?
.error_for_status()?
.json()
.await?;
println!("{}", serde_json::to_string_pretty(&status)?);
return Ok(());
}
let root = config_root(cli.config_dir)?;
let store = Store::open(&root)?;
let runtime = match cli.runtime_dir {
Some(root) => root,
None => omochi_ui::prepare_omochi_runtime_root()?,
};
match cli.command {
Command::Dev {
host,
port,
command,
} => {
let code = dev::run(command, std::net::SocketAddr::new(host, port), root).await?;
if code != 0 {
std::process::exit(code);
}
}
Command::Generate { .. } | Command::Status { .. } => unreachable!(),
Command::Id { hub } => println!(
"{}",
store
.identity(if hub { Role::Hub } else { Role::Gateway })?
.public()
),
Command::Trust {
command: TrustCommand::List,
} => {
for (id, peer) in store.peers()?.1 {
println!("{id}\t{:?}\t{}", peer.role, peer.name);
}
}
Command::Trust {
command: TrustCommand::Forget { peer },
} => {
let id = match peer.parse::<omochi_transport::EndpointId>() {
Ok(id) => id,
Err(_) => {
let matches: Vec<_> = store
.peers()?
.1
.into_iter()
.filter(|(_, p)| p.name == peer)
.collect();
if matches.len() != 1 {
anyhow::bail!("peer name must identify exactly one registered peer");
}
matches[0].0.parse()?
}
};
omochi_transport::live::forget(store, id, Duration::from_secs(10)).await?;
println!("forgot {id}");
}
Command::Pair {
invitation,
name,
hub_name,
offline,
relay,
} => {
let text = if let Some(text) = invitation {
text
} else {
eprintln!("[omochi-ui] paste the invitation and end input (Ctrl+D):");
let mut bytes = Vec::new();
tokio::io::stdin()
.take(16385)
.read_to_end(&mut bytes)
.await?;
String::from_utf8(bytes).context("invitation must be UTF-8")?
};
let invite = Invite::decode(&text)?;
let relays = if relay.is_empty() {
None
} else {
Some(
relay
.into_iter()
.map(|s| s.parse())
.collect::<std::result::Result<Vec<_>, _>>()?,
)
};
let endpoint =
omochi_transport::bind(store.identity(Role::Hub)?, offline, relays).await?;
let result =
omochi_transport::pairing::redeem(&endpoint, store, invite, &name, &hub_name).await;
endpoint.close().await;
result?;
eprintln!("[omochi-ui] registered {name}");
}
Command::Gateway {
remote,
pair,
web_url,
offline,
relay,
name,
} => {
let mut shutdown = ShutdownSignal::new()?;
let relays = if remote && !relay.is_empty() {
Some(
relay
.into_iter()
.map(|s| s.parse())
.collect::<std::result::Result<Vec<_>, _>>()?,
)
} else {
None
};
private_directory(&runtime)?;
let gateway = omochi_gateway::Gateway::new();
let identity = store.identity(Role::Gateway)?;
let node = identity.public().to_string();
let local =
gateway.listen_unix(runtime.join("gateway.sock"), node.clone(), name.clone())?;
let discovery = gateway.discover(runtime.clone());
let mut remote_endpoint = None;
let mut remote_task = None;
let mut result: Result<()> = async {
if remote {
let endpoint = omochi_transport::bind(identity, offline, relays).await?;
remote_endpoint = Some(endpoint.clone());
if pair {
let invitation = Invite::issue(&store, &endpoint)?.encode()?;
let url = web_invitation(&web_url, &invitation)?;
let mut stdout = tokio::io::stdout();
stdout
.write_all(format!("{invitation}\n{url}\n").as_bytes())
.await?;
stdout.flush().await?;
match qrcode::QrCode::new(url.as_bytes()) {
Ok(code) => eprintln!(
"{}\nInvitation expires in 180 seconds.",
code.render::<qrcode::render::unicode::Dense1x2>().build()
),
Err(_) => eprintln!(
"[omochi-ui] invitation is too large for a QR code; use the URL"
),
}
}
remote_task = Some(gateway.listen_remote_async(endpoint, store, name).await?);
}
eprintln!(
"[omochi-ui] gateway {node} at {}",
runtime.join("gateway.sock").display()
);
shutdown.wait().await;
Ok(())
}
.await;
gateway.shutdown();
for task in [local, discovery] {
result = result.and(task.await.unwrap_or_else(|error| Err(error.into())));
}
if let Some(task) = remote_task {
result = result.and(task.await.map_err(anyhow::Error::from));
}
if let Some(endpoint) = remote_endpoint {
endpoint.close().await;
}
result?;
}
Command::Hub {
host,
port,
allow_origin,
no_auto_gateway,
assets,
remote,
web_url,
offline,
relay,
} => {
serve_hub(
store,
root,
runtime,
HubOptions {
host,
port,
allow_origin,
no_auto_gateway,
assets,
remote,
web_url,
offline,
relay,
presets: vec![],
open: false,
},
)
.await?;
}
Command::Up {
file,
host,
port,
allow_origin,
open,
} => {
let node = store.identity(Role::Gateway)?.public().to_string();
let composition = composition::load(&file, &node)?;
private_directory(&runtime)?;
let runtime = runtime.canonicalize()?;
let mut children = composition::Supervisor::start(&composition.services, &runtime)?;
let result = serve_hub(
store,
root,
runtime,
HubOptions {
host,
port,
allow_origin,
no_auto_gateway: false,
assets: None,
remote: false,
web_url: DEFAULT_WEB_URL.to_owned(),
offline: false,
relay: vec![],
presets: vec![composition.preset],
open,
},
)
.await;
children.close().await;
result?;
}
}
Ok(())
}
struct HubOptions {
host: std::net::IpAddr,
port: u16,
allow_origin: Vec<String>,
no_auto_gateway: bool,
assets: Option<PathBuf>,
remote: bool,
web_url: String,
offline: bool,
relay: Vec<String>,
presets: Vec<serde_json::Value>,
open: bool,
}
async fn serve_hub(
store: Store,
root: PathBuf,
runtime: PathBuf,
options: HubOptions,
) -> Result<()> {
let mut shutdown = ShutdownSignal::new()?;
let HubOptions {
host,
port,
allow_origin,
no_auto_gateway,
assets,
remote: remote_dashboard,
web_url,
offline,
relay,
presets,
open,
} = options;
private_directory(&runtime)?;
let socket = runtime.join("gateway.sock");
let owned = if socket.exists() {
probe(&socket).await?;
None
} else if !no_auto_gateway {
let gateway = omochi_gateway::Gateway::new();
let listener = gateway.listen_unix(
socket.clone(),
store.identity(Role::Gateway)?.public().to_string(),
"local".into(),
)?;
let discovery = gateway.discover(runtime);
Some((gateway, listener, discovery))
} else {
None
};
let hub = omochi_hub::Hub::new();
let attached = hub.attach_unix(socket);
let mut tasks = Vec::new();
let mut remote_task = None;
let mut result: Result<()> = async {
let dashboards =
omochi_hub::dashboards::DashboardStore::open(root).map_err(|e| anyhow::anyhow!(e.1))?;
let remote_boards = dashboards.clone();
let remote_hub = hub.clone();
let remote_store = store.clone();
remote_task = Some(tokio::spawn(async move {
let stop = remote_hub.cancellation_token();
loop {
let store = remote_store.clone();
if remote_dashboard
|| omochi_transport::blocking::run(move || store.peers())
.await??
.1
.values()
.any(|p| p.role == Role::Gateway)
{
break;
}
tokio::select! { _ = stop.cancelled() => return Ok::<(), anyhow::Error>(()), _ = tokio::time::sleep(Duration::from_millis(250)) => {} }
}
let relays = if relay.is_empty() {
None
} else {
Some(
relay
.into_iter()
.map(|s| s.parse())
.collect::<std::result::Result<Vec<_>, _>>()?,
)
};
let identity_store = remote_store.clone();
let key =
omochi_transport::blocking::run(move || identity_store.identity(Role::Hub)).await??;
let endpoint = omochi_transport::bind(key, offline, relays).await?;
remote_hub.configure_pairing(endpoint.clone(), remote_store.clone());
let listener = if remote_dashboard {
remote_hub.set_dashboard_host(endpoint.clone(), web_url)?;
let listener = omochi_hub::dashboard_remote::listen(
remote_hub.clone(),
remote_boards,
endpoint.clone(),
)?;
Some(listener)
} else {
None
};
let result = remote_hub
.attach_registered(endpoint, remote_store)
.0
.await
.map_err(anyhow::Error::from);
if let Some(listener) = listener {
result.and(listener.await.unwrap_or_else(|error| Err(error.into())))
} else {
result
}
}));
let selected_preset = presets
.first()
.and_then(|preset| preset["id"].as_str())
.map(str::to_owned);
let (address, http) = omochi_hub::http::serve(
omochi_hub::http::HttpState {
hub: hub.clone(),
dashboards,
allow_origins: allow_origin,
assets,
presets,
},
std::net::SocketAddr::new(host, port),
)
.await?;
tasks.push(http);
let mut dashboard_url = url::Url::parse(&format!("http://{address}/"))?;
if let Some(id) = selected_preset {
dashboard_url.query_pairs_mut().append_pair("preset", &id);
}
eprintln!("[omochi-ui] dashboard {dashboard_url}");
if open {
dashboard_url = dashboard_browser_url(dashboard_url, address)?;
composition::open_browser(dashboard_url.as_str());
}
tokio::select! {
_ = shutdown.wait() => Ok(()),
finished = remote_task.as_mut().unwrap() => {
remote_task.take();
finished.context("hub remote task failed")??;
bail!("hub remote task stopped unexpectedly")
}
}
}
.await;
hub.shutdown();
if let Some(task) = remote_task {
tasks.push(task);
}
if let Some((gateway, listener, discovery)) = owned {
gateway.shutdown();
tasks.extend([listener, discovery]);
}
result = result.and(attached.await.map_err(anyhow::Error::from));
for task in tasks {
result = result.and(task.await.unwrap_or_else(|error| Err(error.into())));
}
result
}
fn dashboard_browser_url(
mut dashboard_url: url::Url,
address: std::net::SocketAddr,
) -> Result<url::Url> {
let browser_ip = match address.ip() {
std::net::IpAddr::V4(ip) if ip.is_unspecified() => std::net::Ipv4Addr::LOCALHOST.into(),
std::net::IpAddr::V6(ip) if ip.is_unspecified() => std::net::Ipv6Addr::LOCALHOST.into(),
ip => ip,
};
dashboard_url
.set_ip_host(browser_ip)
.map_err(|()| anyhow::anyhow!("invalid dashboard browser URL"))?;
Ok(dashboard_url)
}
#[cfg(test)]
mod shutdown_tests {
use super::*;
#[test]
fn dashboard_browser_url_preserves_ip_family_port_and_preset() {
for (address, expected) in [
("127.0.0.1:43210", "http://127.0.0.1:43210/"),
("192.0.2.1:43210", "http://192.0.2.1:43210/"),
("0.0.0.0:43210", "http://127.0.0.1:43210/"),
("[::1]:43210", "http://[::1]:43210/"),
("[2001:db8::1]:43210", "http://[2001:db8::1]:43210/"),
("[::]:43210", "http://[::1]:43210/"),
] {
let address = address.parse().unwrap();
let mut dashboard_url = url::Url::parse(&format!("http://{address}/")).unwrap();
dashboard_url
.query_pairs_mut()
.append_pair("preset", "demo & preview");
let browser_url = dashboard_browser_url(dashboard_url, address).unwrap();
assert_eq!(
browser_url.as_str(),
format!("{expected}?preset=demo+%26+preview"),
"listener address: {address}"
);
}
}
#[test]
fn hosted_invitation_requires_valid_https_base_and_remote_mode() {
for base in [
"http://example.com/",
"https://u:p@example.com/",
"https://example.com/?",
"https://example.com/#",
"relative",
" https://example.com/",
] {
assert!(validate_web_url(base).is_err());
}
assert!(validate_web_url(&format!("https://example.com/{}", "x".repeat(2048))).is_err());
assert!(
Cli::try_parse_from(["omochi-ui", "gateway", "--web-url", "https://example.com/"])
.is_err()
);
let base = "https://example.com/ui/";
assert_eq!(validate_web_url(base).unwrap(), base);
assert_eq!(
web_invitation(base, "omochi://invite/YQ").unwrap(),
"https://example.com/ui/#g:YQ"
);
assert!(
Cli::try_parse_from([
"omochi-ui",
"gateway",
"--remote",
"--pair",
"--web-url",
base
])
.is_ok()
);
}
#[test]
fn downstream_dashboard_listener_is_explicit_and_keeps_offline_ws_separate() {
assert!(Cli::try_parse_from(["omochi-ui", "hub", "--remote"]).is_ok());
assert!(Cli::try_parse_from(["omochi-ui", "hub", "--remote", "--offline"]).is_err());
assert!(
Cli::try_parse_from(["omochi-ui", "hub", "--web-url", "https://ui.example/"]).is_err()
);
assert!(
Cli::try_parse_from([
"omochi-ui",
"hub",
"--remote",
"--web-url",
"https://ui.example/",
"--offline"
])
.is_err()
);
assert!(Cli::try_parse_from(["omochi-ui", "hub", "--offline"]).is_ok());
assert!(
Cli::try_parse_from([
"omochi-ui",
"hub",
"--remote",
"--web-url",
"https://ui.example/dashboard/"
])
.is_ok()
);
}
#[test]
fn hosted_dashboard_defaults_allow_local_mode_and_custom_urls() {
for command in ["hub", "gateway"] {
for remote in [false, true] {
let mut args = vec!["omochi-ui", command];
if remote {
args.push("--remote");
}
let cli = Cli::try_parse_from(args).unwrap();
match cli.command {
Command::Hub {
web_url,
remote: actual_remote,
..
}
| Command::Gateway {
web_url,
remote: actual_remote,
..
} => {
assert_eq!(web_url, "https://app.omochi.fyi/");
assert_eq!(actual_remote, remote);
}
_ => unreachable!(),
}
}
let cli = Cli::try_parse_from([
"omochi-ui",
command,
"--remote",
"--web-url",
"https://ui.example/dashboard/",
])
.unwrap();
match cli.command {
Command::Hub { web_url, .. } | Command::Gateway { web_url, .. } => {
assert_eq!(web_url, "https://ui.example/dashboard/");
}
_ => unreachable!(),
}
}
}
#[tokio::test]
async fn dashboard_store_failure_cleans_up_owned_gateway() {
let root = tempfile::tempdir().unwrap();
let store = Store::open(root.path().join("config")).unwrap();
let dashboards = root.path().join("dashboards");
std::fs::write(&dashboards, "not a directory").unwrap();
let runtime = root.path().join("runtime");
let error = tokio::time::timeout(
Duration::from_secs(10),
serve_hub(
store,
dashboards,
runtime.clone(),
HubOptions {
host: std::net::Ipv4Addr::LOCALHOST.into(),
port: 0,
allow_origin: vec![],
no_auto_gateway: false,
assets: None,
remote: false,
web_url: DEFAULT_WEB_URL.into(),
offline: true,
relay: vec![],
presets: vec![],
open: false,
},
),
)
.await
.unwrap()
.unwrap_err();
assert_eq!(
error.to_string(),
"Dashboard directory must be owned and private"
);
assert!(!runtime.join("gateway.sock").exists());
}
#[tokio::test]
async fn shutdown_is_registered_before_wait_is_polled() {
let mut shutdown = ShutdownSignal::new().unwrap();
assert_eq!(unsafe { libc::raise(libc::SIGTERM) }, 0);
tokio::time::timeout(Duration::from_secs(3), shutdown.wait())
.await
.expect("signal delivered before wait was lost");
}
}