use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::net::Ipv4Addr;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::Instant;
use base64::Engine as _;
use axum::Router;
use axum::extract::State;
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::response::{Html, Response};
use axum::routing::get;
use futures::{SinkExt, StreamExt};
use serde::Serialize;
use tokio::sync::broadcast;
use tokio::time::Duration;
pub use crate::solver::observe::{BrowserId, Phase};
#[derive(Clone, Serialize)]
pub struct Snapshot {
pub id: BrowserId,
pub phase: Phase,
pub url: String,
pub clicks: u32,
pub elapsed_ms: u64,
pub last_click: Option<(i32, i32)>,
pub live: bool,
pub user_agent: Option<String>,
#[serde(default)]
pub exit: Option<String>,
#[serde(default)]
pub retired: bool,
}
struct Entry {
snap: Snapshot,
display: Option<String>,
since: Instant,
}
#[derive(Clone, Serialize)]
pub struct ExitRow {
pub code: String,
pub country: String,
pub health: String,
pub activity: String,
pub latency_ms: Option<u64>,
pub last_probe_unix: Option<f64>,
pub proxy_url: String,
pub over_latency: bool,
pub warm: bool,
pub paced: bool,
pub cooling: bool,
pub cooling_reason: Option<crate::store::Cooling>,
#[serde(flatten)]
pub stats: crate::store::Stats,
}
pub struct Introspector {
state: Mutex<BTreeMap<BrowserId, Entry>>,
next_id: AtomicU32,
tx: broadcast::Sender<Snapshot>,
exits: Mutex<BTreeMap<String, ExitRow>>,
exits_tx: broadcast::Sender<String>,
events: Mutex<VecDeque<Event>>,
events_tx: broadcast::Sender<Event>,
ghosts: Mutex<VecDeque<Arc<Ghost>>>,
ghosts_tx: broadcast::Sender<GhostEvent>,
progress: Mutex<Option<JobProgress>>,
progress_tx: broadcast::Sender<JobProgress>,
mullvad_required: std::sync::atomic::AtomicBool,
solving: std::sync::atomic::AtomicBool,
server: Mutex<Option<tokio::task::JoinHandle<()>>>,
}
#[derive(Clone, Serialize)]
pub struct Event {
pub ts_unix: f64,
pub msg: String,
}
#[derive(Clone, Serialize)]
pub struct JobProgress {
pub total: Option<u64>,
pub done: u64,
pub ok: u64,
pub err: u64,
pub in_flight: u64,
pub per_sec: f64,
pub eta_secs: Option<u64>,
}
const MAX_EVENTS: usize = 300;
const MAX_GHOSTS: usize = 128;
#[derive(Clone)]
struct Ghost {
snap: Snapshot,
summary: String,
shot: String,
latency_ms: Option<u64>,
}
#[derive(Clone)]
enum GhostEvent {
Added(Arc<Ghost>),
Retired(BrowserId),
}
impl Introspector {
pub fn new() -> Arc<Self> {
let (tx, _) = broadcast::channel(256);
let (exits_tx, _) = broadcast::channel(4096);
Arc::new(Introspector {
state: Mutex::new(BTreeMap::new()),
next_id: AtomicU32::new(0),
tx,
exits: Mutex::new(BTreeMap::new()),
exits_tx,
events: Mutex::new(VecDeque::new()),
events_tx: broadcast::channel(256).0,
ghosts: Mutex::new(VecDeque::new()),
ghosts_tx: broadcast::channel(64).0,
progress: Mutex::new(None),
progress_tx: broadcast::channel(16).0,
mullvad_required: std::sync::atomic::AtomicBool::new(false),
solving: std::sync::atomic::AtomicBool::new(false),
server: Mutex::new(None),
})
}
pub fn event(&self, msg: impl Into<String>) {
let msg = msg.into();
let ts_unix = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs_f64())
.unwrap_or(0.0);
let ev = Event { ts_unix, msg };
let mut q = self.events.lock().unwrap();
if q.len() >= MAX_EVENTS {
q.pop_front();
}
q.push_back(ev.clone());
drop(q);
let _ = self.events_tx.send(ev);
}
pub fn set_mullvad_required(&self, required: bool) {
self.mullvad_required.store(required, Ordering::Relaxed);
}
pub fn set_solving(&self, solving: bool) {
if solving {
self.solving.store(true, Ordering::Relaxed);
}
}
pub fn publish_progress(&self, p: JobProgress) {
{
let mut slot = self.progress.lock().unwrap();
if let Some(cur) = slot.as_ref()
&& p.done < cur.done
&& p.done != 0
{
return;
}
*slot = Some(p.clone());
}
let _ = self.progress_tx.send(p);
}
pub fn set_user_agent(&self, id: BrowserId, ua: String) {
self.update(id, |s| s.user_agent = Some(ua));
}
pub fn set_exit(&self, id: BrowserId, code: String) {
self.update(id, |s| s.exit = Some(code));
}
pub fn failed(
&self,
id: BrowserId,
screenshot: Option<Vec<u8>>,
summary: String,
elapsed_ms: u64,
) {
let mut snap = self
.state
.lock()
.unwrap()
.get(&id)
.map(|e| e.snap.clone())
.unwrap_or(Snapshot {
id,
phase: Phase::Failed,
url: String::new(),
clicks: 0,
elapsed_ms: 0,
last_click: None,
live: false,
user_agent: None,
exit: None,
retired: false,
});
snap.phase = Phase::Failed;
snap.elapsed_ms = elapsed_ms;
let latency_ms = snap.exit.as_ref().and_then(|code| self.exit_latency(code));
let shot = screenshot
.map(|b| {
format!(
"data:image/png;base64,{}",
base64::engine::general_purpose::STANDARD.encode(b)
)
})
.unwrap_or_default();
let ghost = Arc::new(Ghost {
snap,
summary,
shot,
latency_ms,
});
let evicted = {
let mut g = self.ghosts.lock().unwrap();
g.push_back(ghost.clone());
(g.len() > MAX_GHOSTS).then(|| g.pop_front()).flatten()
};
if let Some(e) = evicted {
let _ = self.ghosts_tx.send(GhostEvent::Retired(e.snap.id));
}
let _ = self.ghosts_tx.send(GhostEvent::Added(ghost));
}
pub fn dismiss_ghost(&self, id: BrowserId) {
let removed = {
let mut g = self.ghosts.lock().unwrap();
let before = g.len();
g.retain(|x| x.snap.id != id);
before != g.len()
};
if removed {
let _ = self.ghosts_tx.send(GhostEvent::Retired(id));
}
}
fn ghost_json(g: &Ghost) -> serde_json::Value {
serde_json::json!({
"id": g.snap.id,
"exit": g.snap.exit,
"clicks": g.snap.clicks,
"elapsed_ms": g.snap.elapsed_ms,
"latency_ms": g.latency_ms,
"summary": g.summary,
"shot": g.shot,
})
}
fn ghosts_json(&self) -> Vec<serde_json::Value> {
self.ghosts
.lock()
.unwrap()
.iter()
.map(|g| Self::ghost_json(g))
.collect()
}
fn exit_rows(&self) -> Vec<ExitRow> {
self.exits.lock().unwrap().values().cloned().collect()
}
fn state_json(&self) -> serde_json::Value {
serde_json::json!({
"mullvad_required": self.mullvad_required.load(Ordering::Relaxed),
"solving": self.solving.load(Ordering::Relaxed),
"tls_profile": crate::slim::PROFILE,
"browsers": self.all(),
"exits": self.exit_rows(),
"events": self.events.lock().unwrap().iter().cloned().collect::<Vec<_>>(),
"ghosts": self.ghosts_json(),
"progress": self.progress.lock().unwrap().clone(),
})
}
pub fn publish_exit(&self, row: ExitRow) {
let code = row.code.clone();
self.exits.lock().unwrap().insert(code.clone(), row);
let _ = self.exits_tx.send(code);
}
fn exit_row(&self, code: &str) -> Option<ExitRow> {
self.exits.lock().unwrap().get(code).cloned()
}
pub fn exit_latency(&self, code: &str) -> Option<u64> {
self.exits
.lock()
.unwrap()
.get(code)
.and_then(|r| r.latency_ms)
}
pub fn register(&self, display: Option<String>) -> BrowserId {
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
let snap = Snapshot {
id,
phase: Phase::Idle,
url: String::new(),
clicks: 0,
elapsed_ms: 0,
last_click: None,
live: display.is_some(),
user_agent: None,
exit: None,
retired: false,
};
let entry = Entry {
snap: snap.clone(),
display,
since: Instant::now(),
};
self.state.lock().unwrap().insert(id, entry);
let _ = self.tx.send(snap);
id
}
pub fn deregister(&self, id: BrowserId) {
let retired = self.state.lock().unwrap().remove(&id).map(|mut e| {
e.snap.retired = true;
e.snap
});
if let Some(snap) = retired {
let _ = self.tx.send(snap);
}
}
fn update(&self, id: BrowserId, f: impl FnOnce(&mut Snapshot)) {
let snap = {
let mut state = self.state.lock().unwrap();
let Some(e) = state.get_mut(&id) else { return };
f(&mut e.snap);
e.snap.elapsed_ms = e.since.elapsed().as_millis() as u64;
e.snap.clone()
};
let _ = self.tx.send(snap);
}
pub fn phase(&self, id: BrowserId, phase: Phase) {
match phase {
Phase::Cleared | Phase::Blocked | Phase::Failed => {
tracing::info!(browser = id, phase = ?phase, "browser phase")
}
_ => tracing::debug!(browser = id, phase = ?phase, "browser phase"),
}
{
let mut state = self.state.lock().unwrap();
if let Some(e) = state.get_mut(&id) {
e.since = Instant::now();
}
}
self.update(id, |s| s.phase = phase);
}
pub fn navigating(&self, id: BrowserId, url: &str) {
let url = url.to_string();
self.phase(id, Phase::Navigating);
self.update(id, |s| s.url = url);
}
pub fn clicked(&self, id: BrowserId, x: i32, y: i32) {
self.phase(id, Phase::Verifying);
self.update(id, |s| {
s.clicks += 1;
s.last_click = Some((x, y));
});
}
fn all(&self) -> Vec<Snapshot> {
self.state
.lock()
.unwrap()
.values()
.map(|e| e.snap.clone())
.collect()
}
fn display_of(&self, id: BrowserId) -> Option<String> {
self.state
.lock()
.unwrap()
.get(&id)
.and_then(|e| e.display.clone())
}
fn watchable_ids(&self) -> Vec<BrowserId> {
self.state
.lock()
.unwrap()
.iter()
.filter(|(_, e)| e.display.is_some())
.map(|(id, _)| *id)
.collect()
}
async fn frame(&self, id: BrowserId) -> Option<Vec<u8>> {
let display = self.display_of(id)?;
tokio::task::spawn_blocking(move || crate::solver::frame::grab(&display))
.await
.ok()
.flatten()
}
pub async fn serve(self: Arc<Self>) -> Option<String> {
let listener = bind_localhost(7878).await?;
let port = listener.local_addr().ok()?.port();
let app = Router::new()
.route("/", get(index))
.route("/ws", get(ws_upgrade))
.route("/api/state", get(state_api))
.with_state(self.clone());
let handle = tokio::spawn(async move {
if let Err(e) = axum::serve(listener, app).await {
tracing::warn!("introspection server stopped: {e}");
}
});
*self.server.lock().unwrap() = Some(handle);
Some(format!("http://localhost:{port}"))
}
pub fn stop(&self) {
if let Some(handle) = self.server.lock().unwrap().take() {
handle.abort();
}
}
}
async fn bind_localhost(base: u16) -> Option<tokio::net::TcpListener> {
for port in base..base.saturating_add(64) {
if let Ok(l) = tokio::net::TcpListener::bind((Ipv4Addr::LOCALHOST, port)).await {
return Some(l);
}
}
None
}
async fn index() -> Html<&'static str> {
Html(include_str!("dashboard.html"))
}
async fn state_api(State(intro): State<Arc<Introspector>>) -> axum::Json<serde_json::Value> {
axum::Json(intro.state_json())
}
async fn ws_upgrade(ws: WebSocketUpgrade, State(intro): State<Arc<Introspector>>) -> Response {
ws.on_upgrade(move |socket| dashboard_socket(socket, intro))
}
type SharedSink = Arc<tokio::sync::Mutex<futures::stream::SplitSink<WebSocket, Message>>>;
async fn dashboard_socket(socket: WebSocket, intro: Arc<Introspector>) {
let (mut sink, mut rx) = socket.split();
let init = serde_json::json!({
"type": "init",
"solving": intro.solving.load(Ordering::Relaxed),
"browsers": intro.all(),
"exits": intro.exit_rows(),
"ghosts": intro.ghosts_json(),
"progress": intro.progress.lock().unwrap().clone(),
});
if sink
.send(Message::Text(init.to_string().into()))
.await
.is_err()
{
return;
}
let sink: SharedSink = Arc::new(tokio::sync::Mutex::new(sink));
let (sel_tx, sel_rx) = tokio::sync::watch::channel::<Option<BrowserId>>(None);
let mut state = tokio::spawn(state_stream(sink.clone(), intro.clone()));
let mut frames = tokio::spawn(frame_stream(sink, intro.clone(), sel_rx));
loop {
tokio::select! {
incoming = rx.next() => match incoming {
Some(Ok(Message::Text(t))) => {
if let Ok(v) = serde_json::from_str::<serde_json::Value>(&t) {
if let Some(id) = v.get("dismiss").and_then(|s| s.as_u64()) {
intro.dismiss_ghost(id as BrowserId);
} else {
let _ = sel_tx.send(v.get("select").and_then(|s| s.as_u64()).map(|n| n as BrowserId));
}
}
}
Some(Ok(Message::Close(_))) | None => break,
Some(Err(_)) => break,
_ => {}
},
_ = &mut state => break,
_ = &mut frames => break,
}
}
state.abort();
frames.abort();
}
async fn state_stream(sink: SharedSink, intro: Arc<Introspector>) {
use broadcast::error::{RecvError, TryRecvError};
let mut browser_events = intro.tx.subscribe();
let mut exit_events = intro.exits_tx.subscribe();
let mut ghost_events = intro.ghosts_tx.subscribe();
let mut progress_events = intro.progress_tx.subscribe();
let mut dirty_exits: HashSet<String> = HashSet::new();
let mut dirty_browsers: HashMap<BrowserId, Snapshot> = HashMap::new();
let mut ghost_msgs: Vec<serde_json::Value> = Vec::new();
let mut progress_dirty = false;
let mut reseed = false;
loop {
tokio::select! {
r = exit_events.recv() => match r {
Ok(code) => { dirty_exits.insert(code); }
Err(RecvError::Lagged(_)) => reseed = true,
Err(RecvError::Closed) => break,
},
r = progress_events.recv() => match r {
Ok(_) | Err(RecvError::Lagged(_)) => progress_dirty = true,
Err(RecvError::Closed) => break,
},
r = browser_events.recv() => match r {
Ok(s) => { dirty_browsers.insert(s.id, s); }
Err(RecvError::Lagged(_)) => {}
Err(RecvError::Closed) => break,
},
r = ghost_events.recv() => match r {
Ok(GhostEvent::Added(g)) => ghost_msgs.push(serde_json::json!({ "type": "ghost", "ghost": Introspector::ghost_json(&g) })),
Ok(GhostEvent::Retired(id)) => ghost_msgs.push(serde_json::json!({ "type": "ghost_retired", "id": id })),
Err(RecvError::Lagged(_)) => {}
Err(RecvError::Closed) => break,
},
}
loop {
match exit_events.try_recv() {
Ok(code) => {
dirty_exits.insert(code);
}
Err(TryRecvError::Lagged(_)) => reseed = true,
Err(TryRecvError::Empty) | Err(TryRecvError::Closed) => break,
}
}
while let Ok(r) = progress_events.try_recv() {
let _ = r;
progress_dirty = true;
}
while let Ok(s) = browser_events.try_recv() {
dirty_browsers.insert(s.id, s);
}
while let Ok(ev) = ghost_events.try_recv() {
match ev {
GhostEvent::Added(g) => ghost_msgs.push(
serde_json::json!({ "type": "ghost", "ghost": Introspector::ghost_json(&g) }),
),
GhostEvent::Retired(id) => {
ghost_msgs.push(serde_json::json!({ "type": "ghost_retired", "id": id }))
}
}
}
let mut msgs: Vec<serde_json::Value> = Vec::new();
if reseed {
reseed = false;
dirty_exits.clear();
msgs.push(serde_json::json!({ "type": "exits_full", "exits": intro.exit_rows() }));
}
for code in dirty_exits.drain() {
if let Some(row) = intro.exit_row(&code) {
msgs.push(serde_json::json!({ "type": "exit", "row": row }));
}
}
for (_, snap) in dirty_browsers.drain() {
msgs.push(serde_json::json!({ "type": "update", "browser": snap }));
}
msgs.append(&mut ghost_msgs);
if progress_dirty {
progress_dirty = false;
if let Some(p) = intro.progress.lock().unwrap().clone() {
msgs.push(serde_json::json!({ "type": "progress", "progress": p }));
}
}
let mut s = sink.lock().await;
for m in msgs {
if s.send(Message::Text(m.to_string().into())).await.is_err() {
return;
}
}
}
}
async fn frame_stream(
sink: SharedSink,
intro: Arc<Introspector>,
sel_rx: tokio::sync::watch::Receiver<Option<BrowserId>>,
) {
let mut thumb_tick = tokio::time::interval(Duration::from_millis(600));
let mut sel_tick = tokio::time::interval(Duration::from_millis(150));
loop {
tokio::select! {
_ = thumb_tick.tick() => {
let grabs = intro.watchable_ids().into_iter().map(|id| grab_frame(&intro, id));
for frame in futures::future::join_all(grabs).await.into_iter().flatten() {
if sink.lock().await.send(Message::Binary(frame.into())).await.is_err() {
return;
}
}
}
_ = sel_tick.tick() => {
let selected = *sel_rx.borrow(); if let Some(id) = selected
&& let Some(frame) = grab_frame(&intro, id).await
&& sink.lock().await.send(Message::Binary(frame.into())).await.is_err()
{
return;
}
}
}
}
}
async fn grab_frame(intro: &Arc<Introspector>, id: BrowserId) -> Option<Vec<u8>> {
let png = tokio::time::timeout(Duration::from_millis(200), intro.frame(id))
.await
.ok()
.flatten()?;
Some(tag_frame(id, png))
}
fn tag_frame(id: BrowserId, mut png: Vec<u8>) -> Vec<u8> {
let mut buf = id.to_be_bytes().to_vec();
buf.append(&mut png);
buf
}
impl crate::solver::observe::Observer for Introspector {
fn register(&self, display: Option<String>) -> BrowserId {
self.register(display)
}
fn phase(&self, id: BrowserId, phase: Phase) {
self.phase(id, phase)
}
fn navigating(&self, id: BrowserId, url: &str) {
self.navigating(id, url)
}
fn set_user_agent(&self, id: BrowserId, ua: String) {
self.set_user_agent(id, ua)
}
fn event(&self, msg: String) {
self.event(msg)
}
fn clicked(&self, id: BrowserId, x: i32, y: i32) {
self.clicked(id, x, y)
}
fn failed(&self, id: BrowserId, screenshot: Option<Vec<u8>>, summary: String, elapsed_ms: u64) {
self.failed(id, screenshot, summary, elapsed_ms)
}
fn deregister(&self, id: BrowserId) {
self.deregister(id)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn deregister_removes_the_row_so_dead_browsers_dont_accumulate() {
let intro = Introspector::new();
let a = intro.register(Some(":1".into()));
let b = intro.register(Some(":2".into()));
assert_eq!(intro.all().len(), 2);
intro.deregister(a);
let rows = intro.all();
assert_eq!(rows.len(), 1, "retired browser's row must be dropped");
assert_eq!(rows[0].id, b);
intro.deregister(a);
assert_eq!(intro.all().len(), 1);
}
#[test]
fn deregister_broadcasts_a_retired_snapshot_so_live_dashboards_drop_the_row() {
let intro = Introspector::new();
let mut rx = intro.tx.subscribe();
let a = intro.register(Some(":1".into()));
assert!(
!rx.try_recv().unwrap().retired,
"register broadcasts a live snapshot"
);
intro.deregister(a);
let snap = rx
.try_recv()
.expect("deregister must broadcast, not just mutate state");
assert_eq!(snap.id, a);
assert!(
snap.retired,
"the broadcast marks the row retired so the dashboard removes it"
);
}
}