use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use axum::Router;
use axum::extract::Request;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::routing::post;
#[derive(Clone)]
struct Script {
status: StatusCode,
content_type: &'static str,
body: String,
}
#[derive(Clone)]
pub(super) struct ScriptedServer {
url: String,
stream: Arc<std::sync::Mutex<Vec<String>>>,
script: Arc<std::sync::Mutex<Script>>,
get_status: Arc<std::sync::Mutex<(StatusCode, &'static str)>>,
gets: Arc<AtomicU64>,
last_get: Arc<std::sync::Mutex<String>>,
}
impl ScriptedServer {
pub(super) async fn start(stream: Vec<String>) -> Result<Self, Box<dyn std::error::Error>> {
let state = Self {
url: String::new(),
stream: Arc::new(std::sync::Mutex::new(stream)),
script: Arc::new(std::sync::Mutex::new(Script {
status: StatusCode::OK,
content_type: "application/json",
body: String::new(),
})),
get_status: Arc::new(std::sync::Mutex::new((StatusCode::OK, "text/event-stream"))),
gets: Arc::new(AtomicU64::new(0)),
last_get: Arc::new(std::sync::Mutex::new(String::new())),
};
let router = Router::new()
.route(
"/probe",
post({
let state = state.clone();
move || {
let state = state.clone();
async move { state.post() }
}
})
.get({
let state = state.clone();
move |request: Request| {
let state = state.clone();
async move { state.get(request.uri().path()) }
}
}),
)
.with_state(state.clone());
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
let port = listener.local_addr()?.port();
tokio::spawn(async move {
let _ = axum::serve(listener, router).await;
});
Ok(Self {
url: format!("http://127.0.0.1:{port}/probe"),
..state
})
}
pub(super) fn url(&self) -> &str {
&self.url
}
pub(super) fn with_post_status(
self,
status: StatusCode,
content_type: &'static str,
body: &str,
) -> Self {
*self
.script
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Script {
status,
content_type,
body: body.to_owned(),
};
self
}
pub(super) fn with_get_status(self, status: StatusCode, content_type: &'static str) -> Self {
*self
.get_status
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = (status, content_type);
self
}
pub(super) fn gets(&self) -> u64 {
self.gets.load(Ordering::SeqCst)
}
pub(super) fn last_get_path(&self) -> String {
self.last_get
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
}
fn post(&self) -> axum::response::Response {
let script = self
.script
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone();
(
script.status,
[(axum::http::header::CONTENT_TYPE, script.content_type)],
script.body,
)
.into_response()
}
fn get(&self, path: &str) -> axum::response::Response {
self.gets.fetch_add(1, Ordering::SeqCst);
let path = path.to_owned();
*self
.last_get
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = path;
let (status, content_type) = *self
.get_status
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let body = self
.stream
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.join("");
(
status,
[(axum::http::header::CONTENT_TYPE, content_type)],
body,
)
.into_response()
}
}