use std::collections::HashMap; use std::error::Error; use std::{env, fmt, fs}; use axum::extract::State; use axum::{ body::Full, http::{Response, StatusCode}, response::IntoResponse, routing::get, Router, }; use lazy_static::lazy_static; use maud::html; use maud::Markup; use prometheus::core::{MetricVec, AtomicF64}; use prometheus::{register_gauge_vec, Encoder, GaugeVec, TextEncoder, Registry, Opts}; use reqwest::header::USER_AGENT; use serde::Deserialize; use url::Url; lazy_static!{ static ref PLAYER_COUNT: GaugeVec = register_gauge_vec!( "vrchat_playercount", "Current number of players in instance.", &["instance", "world", "name"], ) .unwrap(); static ref VRCDN_VIEWERS: GaugeVec = register_gauge_vec!( "vrcdn_viewers", "Current number viewers according to VRCDN's API.", &["region"], ) .unwrap(); } #[derive(Debug)] enum WsError { Reqwest(reqwest::Error), Url(url::ParseError), Custom(String), } impl From for WsError { fn from(value: url::ParseError) -> Self { Self::Url(value) } } impl From for WsError { fn from(value: reqwest::Error) -> Self { Self::Reqwest(value) } } impl Error for WsError {} impl fmt::Display for WsError { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { WsError::Reqwest(e) => write!(f, "Reqwest error: {}", e), WsError::Custom(e) => write!(f, "Error: {}", e), WsError::Url(_) => todo!(), } } } #[derive(Clone, Debug, Deserialize)] struct Config { /// Groups we want to track. groups: Option>, /// List of worlds we want to track. worlds: Option>, #[serde(skip)] vrchat_token: Option, } #[derive(Clone, Debug, Deserialize)] struct VrcGroup { id: String, vrcdn: Option, } #[derive(Clone, Debug, Deserialize)] struct VrcInstance { /// Instance ID. instance: Option, /// World ID. world: Option, /// Raw location location: Option, /// Custom name for the instance. name: Option, } #[derive(Clone, Debug, Deserialize)] struct VrcInstanceData { #[serde(rename = "userCount")] user_count: Option, } #[derive(Clone, Debug, Deserialize)] struct VrCdnData { viewers: Vec, } #[derive(Clone, Debug, Deserialize)] struct VrCdnRegion { region: String, total: f64, } #[derive(Clone, Debug, Deserialize)] struct VrcGroupInstance { location: String, } #[derive(Clone, Debug, Deserialize)] struct VrcWorldData { favorites: f64, version: f64, visits: f64, popularity: f64, heat: f64, #[serde(rename = "publicOccupants")] public_occupants: f64, #[serde(rename = "privateOccupants")] private_occupants: f64, #[serde(rename = "occupants")] total_occupants: f64, instances: Vec<(String, f64)>, } #[tokio::main] async fn main() -> Result<(), ()> { let content = fs::read_to_string("config.toml").unwrap(); let mut config: Config = toml::from_str(&content).unwrap(); config.vrchat_token = Some(env::var("VRCHAT_AUTH_TOKEN").unwrap()); let app = Router::new() .route("/", get(homepage)) .route("/metrics", get(metrics_handler)) .with_state(config); let _ = axum::Server::bind(&"0.0.0.0:6534".parse().unwrap()) .serve(app.into_make_service()) .await; Ok(()) } async fn metrics_handler(State(config): State) -> Result { match metrics(config).await { Ok(b) => Ok(Response::builder() .header("content-type", "text/plain") .status(StatusCode::OK) .body(Full::from(b)) .unwrap()), Err(_e) => Err(StatusCode::INTERNAL_SERVER_ERROR), } } async fn homepage() -> Markup { html! { body { main { a href = "/metrics" { "metrics" } } } } } async fn metrics(config: Config) -> Result, WsError> { PLAYER_COUNT.reset(); VRCDN_VIEWERS.reset(); let encoder = TextEncoder::new(); let client = reqwest::Client::new(); let auth_cookie = format!("auth={}", &config.vrchat_token.unwrap()); if config.groups.is_some() { for (name, group) in config.groups.unwrap() { let _ = group_metrics(&client, &auth_cookie, name, group).await; } } if config.worlds.is_some() { for (name, id) in config.worlds.unwrap() { let _ = world_metrics(&client, &auth_cookie, name, id); } } let metric_families = prometheus::gather(); let mut buffer = vec![]; encoder.encode(&metric_families, &mut buffer).unwrap(); Ok(buffer) } async fn group_metrics(client: &reqwest::Client, auth_cookie: &String, name: String, group: VrcGroup) -> Result<(), WsError> { let instance_list_url = format!( "https://api.vrchat.cloud/api/1/groups/{}/instances", group.id ); let url = Url::parse(&instance_list_url).unwrap(); let req = client .get(url) .header( USER_AGENT, "vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)", ) .header("Cookie", auth_cookie) .send() .await?; let data: Vec = req.json().await?; let instances: Vec = data .into_iter() .map(|f| { let spl: Vec<&str> = f.location.split(":").collect(); VrcInstance { instance: Some(spl[0].to_owned()), world: Some(spl[1].to_owned()), location: Some(f.location), name: None, } }) .collect(); for instance in instances { let api_url = format!( "https://api.vrchat.cloud/api/1/instances/{}", &instance.location.clone().unwrap() ); let url = Url::parse(&api_url).unwrap(); let req = client .get(url) .header( USER_AGENT, "vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)", ) .header("Cookie", auth_cookie) .send() .await?; let data: VrcInstanceData = req.json().await?; let instance_name = instance.name.unwrap_or(instance.location.unwrap()); PLAYER_COUNT .with_label_values(&[&instance.world.unwrap(), &instance.instance.unwrap(), &name]) .set(data.user_count.unwrap_or(0 as f64)); } if group.vrcdn.is_some() { let vrcdn_url = format!( "https://api.vrcdn.live/v1/viewers/{}", group.vrcdn.unwrap() ); let req = client .get(vrcdn_url) .header( USER_AGENT, "vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)", ) .send() .await .unwrap(); let vrcdn_data: VrCdnData = req.json().await.unwrap(); for region in vrcdn_data.viewers { VRCDN_VIEWERS .with_label_values(&[®ion.region]) .set(region.total); } } Ok(()) } async fn world_metrics(client: &reqwest::Client, auth_cookie: &String, name: String, id: String) -> Result<(), WsError> { let api_url = format!( "https://api.vrchat.cloud/api/1/worlds/{}", &id ); let url = Url::parse(&api_url).unwrap(); let world_data: VrcWorldData = client .get(url) .header( USER_AGENT, "vr-event-tracker(git.gmem.ca/arch/vr-event-tracker)", ) .header("Cookie", auth_cookie) .send() .await?.json().await?; for instance in world_data.instances { PLAYER_COUNT .with_label_values(&[&instance.0, &id, &name]) .set(instance.1); } Ok(()) }