From 11496fff2b448d013d9aae3d46316b4d874bcb7a Mon Sep 17 00:00:00 2001
From: videogame hacker <half-kh-hacker@hackery.site>
Date: Mon, 13 Mar 2023 17:04:49 +0000
Subject: [PATCH] [wish-server-rs] Add endpoint to see view count

---
 wish-server-rs/Cargo.lock               |  1 +
 wish-server-rs/Cargo.toml               |  1 +
 wish-server-rs/src/main.rs              | 14 +++++++++++--
 wish-server-rs/src/wish/channel_info.rs | 26 +++++++++++++++++++++++++
 wish-server-rs/src/wish/mod.rs          |  1 +
 5 files changed, 41 insertions(+), 2 deletions(-)
 create mode 100644 wish-server-rs/src/wish/channel_info.rs

diff --git a/wish-server-rs/Cargo.lock b/wish-server-rs/Cargo.lock
index 4402a40..f6942fd 100644
--- a/wish-server-rs/Cargo.lock
+++ b/wish-server-rs/Cargo.lock
@@ -3147,6 +3147,7 @@ dependencies = [
  "dotenv",
  "miette",
  "once_cell",
+ "serde",
  "sqlx",
  "tokio",
  "tower-http",
diff --git a/wish-server-rs/Cargo.toml b/wish-server-rs/Cargo.toml
index e5b826b..0f2a456 100644
--- a/wish-server-rs/Cargo.toml
+++ b/wish-server-rs/Cargo.toml
@@ -10,6 +10,7 @@ axum = { version = "0.6.11", features = ["macros"] }
 dotenv = "0.15.0"
 miette = { version = "5.5.0", features = ["fancy"] }
 once_cell = "1.17.1"
+serde = { version = "1.0.155", features = ["derive"] }
 sqlx = { version = "0.6.2", features = ["runtime-tokio-rustls", "sqlite"] }
 tokio = { version = "1.26.0", features = ["full"] }
 tower-http = { version = "0.4.0", features = ["trace", "cors"] }
diff --git a/wish-server-rs/src/main.rs b/wish-server-rs/src/main.rs
index 2409c7e..facb49b 100644
--- a/wish-server-rs/src/main.rs
+++ b/wish-server-rs/src/main.rs
@@ -12,7 +12,12 @@ mod streams;
 mod util;
 mod wish;
 
-use crate::wish::{setup_webrtc, whep::handle_whep, whip::handle_whip};
+use crate::wish::{
+    channel_info::{self, get_channel_info},
+    setup_webrtc,
+    whep::handle_whep,
+    whip::handle_whip,
+};
 
 #[derive(Clone)]
 pub struct AppState {
@@ -55,6 +60,10 @@ async fn main() -> Result<()> {
     let app = Router::new()
         .route("/api/wish-server/whip", routing::any(handle_whip))
         .route("/api/wish-server/whep", routing::any(handle_whep))
+        .route(
+            "/api/wish-server/stream/:channel",
+            routing::get(get_channel_info),
+        )
         .layer(TraceLayer::new_for_http())
         .layer(
             CorsLayer::new()
@@ -75,7 +84,8 @@ async fn main() -> Result<()> {
     axum::Server::bind(&bind_addr)
         .serve(app.into_make_service())
         .await
-        .unwrap();
+        .into_diagnostic()
+        .wrap_err("Server listening failed")?;
 
     Ok(())
 }
diff --git a/wish-server-rs/src/wish/channel_info.rs b/wish-server-rs/src/wish/channel_info.rs
new file mode 100644
index 0000000..65bf184
--- /dev/null
+++ b/wish-server-rs/src/wish/channel_info.rs
@@ -0,0 +1,26 @@
+use std::sync::atomic::Ordering;
+
+use axum::{
+    extract::{Path, State},
+    http::StatusCode,
+    response::{IntoResponse, Response},
+    Json,
+};
+use serde::Serialize;
+
+use crate::{streams::get_ongoing_stream, util::log_http_error};
+
+#[derive(Serialize)]
+struct ChannelInfo {
+    view_count: u64,
+}
+
+pub async fn get_channel_info(Path(channel): Path<String>) -> Response {
+    let channel = match get_ongoing_stream(&channel).await {
+        Ok(c) => c,
+        Err(e) => return log_http_error(StatusCode::INTERNAL_SERVER_ERROR, e),
+    };
+    let view_count = channel.viewer_count.load(Ordering::Relaxed);
+    let info = ChannelInfo { view_count };
+    Json(info).into_response()
+}
diff --git a/wish-server-rs/src/wish/mod.rs b/wish-server-rs/src/wish/mod.rs
index 86cd13a..11dbce4 100644
--- a/wish-server-rs/src/wish/mod.rs
+++ b/wish-server-rs/src/wish/mod.rs
@@ -11,6 +11,7 @@ use webrtc::{
     peer_connection::configuration::RTCConfiguration,
 };
 
+pub mod channel_info;
 pub mod whep;
 pub mod whip;