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::{register_gauge_vec, Encoder, GaugeVec, TextEncoder}; 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 { vrcdn: Option, group: Option, vrchat_token: Option, } #[derive(Clone, Debug, Deserialize)] struct VrcInstance { /// Instance ID. instance: Option, /// World ID. world: Option, /// Raw location location: Option, /// Custom name for the instance. #[serde(skip)] 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, } #[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> { let encoder = TextEncoder::new(); let client = reqwest::Client::new(); let auth_cookie = format!("auth={}", &config.vrchat_token.unwrap()); let mut instances: Vec = vec![]; // Check if we can fetch instances from a group if set. if config.group.is_some() { let api_url = format!( "https://api.vrchat.cloud/api/1/groups/{}/instances", config.group.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: Vec = req.json().await?; instances = 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 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)); } let vrcdn_url = format!( "https://api.vrcdn.live/v1/viewers/{}", config.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); } let metric_families = prometheus::gather(); let mut buffer = vec![]; encoder.encode(&metric_families, &mut buffer).unwrap(); Ok(buffer) }