Add current connections user agent metrics

Author: moehreag

src/endpoints/global_data.rs

@@ -18,7 +18,7 @@ const PROJECT_ID: &str = "p2rxzX0q";
 pub struct GlobalDataContainer {
 	last_full_refresh: DateTime<Utc>,
 	last_player_refresh: DateTime<Utc>,
-	data: GlobalData,
+	pub data: GlobalData,
 }
 
 impl Default for GlobalDataContainer {
@@ -33,7 +33,8 @@ impl Default for GlobalDataContainer {
 					latest_version: String::new(),
 				},
 				notes: String::new(),
-				user_agents: DashMap::new(),
+				request_user_agents: DashMap::new(),
+				gateway_user_agents: DashMap::new(),
 			},
 		}
 	}
@@ -48,7 +49,9 @@ pub struct GlobalData {
 	#[serde(skip_serializing_if = "String::is_empty")]
 	notes: String,
 	#[serde(skip)]
-	pub user_agents: DashMap<String, u32>,
+	pub request_user_agents: DashMap<String, u32>,
+	#[serde(skip)]
+	pub gateway_user_agents: DashMap<String, u32>,
 }
 
 #[derive(Serialize)]
@@ -63,7 +66,8 @@ impl GlobalData {
 			online_players: online,
 			modrinth_data: self.modrinth_data.clone(),
 			notes: self.notes.clone(),
-			user_agents: self.user_agents.clone(),
+			request_user_agents: self.request_user_agents.clone(),
+			gateway_user_agents: self.gateway_user_agents.clone(),
 		}
 	}
 }
@@ -75,7 +79,8 @@ impl Clone for GlobalData {
 			online_players: self.online_players,
 			modrinth_data: self.modrinth_data.clone(),
 			notes: self.notes.clone(),
-			user_agents: self.user_agents.clone(),
+			request_user_agents: self.request_user_agents.clone(),
+			gateway_user_agents: self.gateway_user_agents.clone(),
 		}
 	}
 }
@@ -112,15 +117,17 @@ pub async fn get(
 		return Ok(Json(cloned));
 	}
 	let data = if full_refresh {
-		let agents = data_container.data.user_agents.clone();
+		let request_user_agents = data_container.data.request_user_agents.clone();
+		let gateway_user_agents = data_container.data.gateway_user_agents.clone();
 		GlobalData {
 			total_players: get_total_players(&database).await?,
 			online_players: online_users.len() as u32,
 			modrinth_data: fetch_modrinth_data(client).await?,
 			notes: (cl_args.notes_file.as_ref())
 				.map(|file| read_to_string(file).unwrap_or_else(|_| String::new()))
 				.unwrap_or_else(String::new),
-			user_agents: agents,
+			request_user_agents,
+			gateway_user_agents,
 		}
 	} else {
 		data_container
@@ -163,11 +170,15 @@ pub async fn metrics(
 		writeln!(response, "lifetime_players {lifetime_players}");
 		writeln!(response, "online_players {online_players}");
 		let data_container = global_data.read().await;
-		let agents = data_container.data.user_agents.clone();
-		for (agent, count) in agents {
+		let request_agents = data_container.data.request_user_agents.clone();
+		for (agent, count) in request_agents {
 			writeln!(response, "request_count{{user_agent=\"{agent}\"}} {count}");
 		}
-		data_container.data.user_agents.clear();
+		data_container.data.request_user_agents.clear();
+		let gateway_agents = data_container.data.gateway_user_agents.clone();
+		for (agent, count) in gateway_agents {
+			writeln!(response, "connections{{user_agent=\"{agent}\"}} {count}");
+		}
 	};
 
 	Ok(response)
@@ -205,11 +216,14 @@ async fn fetch_modrinth_data(client: Client) -> Result<ModrinthData, ApiError> {
 	})
 }
 
-pub struct UserAgent;
+pub struct RequestUserAgentCounter;
 
-impl FromRequestParts<ApiState> for UserAgent {
+impl FromRequestParts<ApiState> for RequestUserAgentCounter {
 	type Rejection = ApiError;
-	async fn from_request_parts(parts: &mut Parts, state: &ApiState) -> Result<UserAgent, Self::Rejection> {
+	async fn from_request_parts(
+		parts: &mut Parts,
+		state: &ApiState,
+	) -> Result<RequestUserAgentCounter, Self::Rejection> {
 		if parts.uri.path().ends_with("metrics") {
 			return Ok(Self);
 		}
@@ -222,8 +236,8 @@ impl FromRequestParts<ApiState> for UserAgent {
 			.replace("\\", "")
 			.replace("\"", "");
 
-		let a = state.global_data.read().await;
-		let agents = &a.data.user_agents;
+		let container = state.global_data.read().await;
+		let agents = &container.data.request_user_agents;
 		if agents.contains_key(&agent) {
 			let prev = agents.get(&agent).map(|v| *v.value()).unwrap_or(0);
 			agents.insert(agent, prev + 1);

src/endpoints/image.rs

@@ -9,7 +9,7 @@ use bytes::Buf;
 use chrono::{DateTime, Utc};
 use reqwest::StatusCode;
 use serde::{Deserialize, Serialize};
-use sqlx::{query, PgPool};
+use sqlx::query;
 use uuid::Uuid;
 
 use crate::{

src/extractors.rs

@@ -1,7 +1,7 @@
 use crate::{errors::ApiError, ApiState};
 use axum::{
 	extract::{FromRequestParts, OptionalFromRequestParts},
-	http::{request::Parts, StatusCode},
+	http::{self, request::Parts, StatusCode},
 };
 use base64::{engine::general_purpose::STANDARD_NO_PAD, Engine};
 use sqlx::query;
@@ -63,3 +63,21 @@ impl OptionalFromRequestParts<ApiState> for Authentication {
 			.ok())
 	}
 }
+
+pub struct UserAgent(pub String);
+
+impl FromRequestParts<ApiState> for UserAgent {
+	type Rejection = ApiError;
+	async fn from_request_parts(parts: &mut Parts, _state: &ApiState) -> Result<UserAgent, Self::Rejection> {
+		let agent = parts
+			.headers
+			.get(http::header::USER_AGENT)
+			.map(|v| v.to_str())
+			.ok_or(StatusCode::BAD_REQUEST)?
+			.map_err(|_| StatusCode::BAD_REQUEST)?
+			.replace("\\", "")
+			.replace("\"", "");
+
+		Ok(Self(agent))
+	}
+}

src/gateway.rs

@@ -1,3 +1,4 @@
+use crate::extractors::UserAgent;
 use crate::{errors::ApiError, extractors::Authentication, ApiState};
 use axum::extract::{ws::close_code, ws::CloseFrame, ws::Message, ws::WebSocket, State, WebSocketUpgrade};
 use axum::{body::Body, response::Response};
@@ -12,6 +13,7 @@ use DisconnectReason::*;
 pub async fn gateway(
 	state: State<ApiState>,
 	Authentication(uuid): Authentication,
+	UserAgent(agent): UserAgent,
 	socket: WebSocketUpgrade,
 ) -> Result<Response<Body>, ApiError> {
 	if let Some(socket) = state.socket_sender.remove(&uuid) {
@@ -20,22 +22,36 @@ pub async fn gateway(
 		drop(socket);
 	}
 
-	Ok(socket.on_upgrade(move |socket| gateway_accept_handler(state, uuid, socket)))
+	Ok(socket.on_upgrade(move |socket| gateway_accept_handler(state, uuid, socket, agent)))
 }
 
 async fn gateway_accept_handler(
 	State(ApiState {
 		database,
 		online_users,
 		socket_sender,
+		global_data,
 		..
 	}): State<ApiState>,
 	uuid: Uuid,
 	mut socket: WebSocket,
+	user_agent: String,
 ) {
 	online_users.insert(uuid, None);
 	let (sender, mut receiver) = unbounded_channel();
 	socket_sender.insert(uuid, sender);
+	{
+		let container = global_data.read().await;
+		let agents = &container.data.gateway_user_agents;
+		let user_agent = user_agent.clone();
+		if agents.contains_key(&user_agent) {
+			let prev = agents.get(&user_agent).map(|v| *v.value()).unwrap_or(0);
+			agents.insert(user_agent, prev + 1);
+		} else {
+			agents.insert(user_agent, 1);
+		}
+		drop(container);
+	}
 
 	let disconnect_reason = gateway_accept(&mut socket, &mut receiver).await.unwrap_err();
 	let _ = socket
@@ -50,6 +66,17 @@ async fn gateway_accept_handler(
 	}
 	receiver.close();
 	online_users.remove(&uuid);
+	{
+		let container = global_data.read().await;
+		let agents = &container.data.gateway_user_agents;
+		let prev = agents.get(&user_agent).map(|v| *v.value()).unwrap_or(0);
+		if prev > 1 {
+			agents.insert(user_agent, prev - 1);
+		} else {
+			agents.remove(&user_agent);
+		}
+		drop(container);
+	};
 
 	let _ = query!("UPDATE players SET last_online = 'now' WHERE uuid = $1 AND show_last_online = true", uuid)
 		.execute(&database)

src/main.rs

@@ -1,4 +1,4 @@
-use crate::endpoints::global_data::{self, GlobalDataContainer, UserAgent};
+use crate::endpoints::global_data::{self, GlobalDataContainer, RequestUserAgentCounter};
 use crate::endpoints::user::{self, Activity};
 use crate::endpoints::{account, brew_coffee, channel, get_authenticate, image, not_found};
 use crate::gateway::gateway;
@@ -167,7 +167,7 @@ async fn main() -> anyhow::Result<()> {
 		.route("/hypixel", get(hypixel::get))
 		//.route("/report/:message", post(channel::report_message))
 		.route("/brew_coffee", get(brew_coffee).post(brew_coffee))
-		.layer(axum::middleware::from_extractor_with_state::<UserAgent, ApiState>(state.clone()))
+		.layer(axum::middleware::from_extractor_with_state::<RequestUserAgentCounter, ApiState>(state.clone()))
 		.fallback(not_found)
 		.with_state(state);