vrchat-prometheus-adapter/src/main.rs

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<url::ParseError> for WsError {
    fn from(value: url::ParseError) -> Self {
        Self::Url(value)
    }
}

impl From<reqwest::Error> 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<String>,
    group: Option<String>,
    vrchat_token: Option<String>,
}

#[derive(Clone, Debug, Deserialize)]
struct VrcInstance {
    /// Instance ID.
    instance: Option<String>,
    /// World ID.
    world: Option<String>,
    /// Raw location
    location: Option<String>,
    /// Custom name for the instance.
    #[serde(skip)]
    name: Option<String>,
}

#[derive(Clone, Debug, Deserialize)]
struct VrcInstanceData {
    #[serde(rename = "userCount")]
    user_count: Option<f64>,
}

#[derive(Clone, Debug, Deserialize)]
struct VrCdnData {
    viewers: Vec<VrCdnRegion>,
}
#[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<Config>) -> Result<impl IntoResponse, StatusCode> {
    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<Vec<u8>, 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());
    let mut instances: Vec<VrcInstance> = 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<VrcGroupInstance> = 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(&[&region.region])
            .set(region.total);
    }

    let metric_families = prometheus::gather();
    let mut buffer = vec![];
    encoder.encode(&metric_families, &mut buffer).unwrap();
    Ok(buffer)
}