#![expect(
clippy::unwrap_used,
clippy::expect_used,
reason = "example/test/bench: panic-on-error and print-for-output are the standard patterns for demos and harnesses"
)]
use rama::{
Layer,
error::{BoxError, BoxErrorExt},
http::{
BodyExtractExt,
client::{EasyHttpWebClient, HttpPooledConnectorConfig},
server::HttpServer,
service::{client::HttpClientExt, web::WebService},
},
layer::{
LimitLayer,
limit::{Policy, PolicyOutput, policy::PolicyResult},
},
rt::Executor,
tcp::server::TcpListener,
telemetry::tracing::{
self,
level_filters::LevelFilter,
subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt},
},
};
use std::sync::{
Arc,
atomic::{AtomicBool, Ordering},
};
use tokio::{sync::oneshot::Sender, sync::oneshot::channel};
const ADDRESS: &str = "127.0.0.1:62024";
#[tokio::main]
async fn main() {
setup_tracing();
let (ready_tx, ready_rx) = channel();
tokio::spawn(run_server(ADDRESS, ready_tx));
ready_rx.await.unwrap();
let client = EasyHttpWebClient::connector_builder()
.with_default_transport_connector()
.with_default_dns_connector()
.without_tls_proxy_support()
.with_proxy_support()
.without_tls_support()
.with_default_http_connector(Executor::default())
.try_with_connection_pool(HttpPooledConnectorConfig {
max_concurrent_streams: 20,
..Default::default()
})
.expect("connection pool")
.build_client();
let resp = client
.get(format!("http://{ADDRESS}/"))
.send()
.await
.unwrap();
let body = resp.try_into_string().await.unwrap();
tracing::info!("body: {body}");
assert_eq!(body, "Hello, World!");
let _resp = client
.get(format!("http://{ADDRESS}/"))
.send()
.await
.unwrap();
let client = EasyHttpWebClient::default();
let result = client.get(format!("http://{ADDRESS}/")).send().await;
_ = result.unwrap_err();
}
fn setup_tracing() {
tracing::subscriber::registry()
.with(fmt::layer())
.with(
EnvFilter::builder()
.with_default_directive(LevelFilter::DEBUG.into())
.from_env_lossy(),
)
.init();
}
async fn run_server(addr: &str, ready: Sender<()>) {
tracing::info!("running service at: {addr}");
let http_service =
HttpServer::default().service(WebService::default().with_get("/", "Hello, World!"));
let serve = TcpListener::build(Executor::default())
.bind_address(addr)
.await
.expect("bind TCP Listener")
.serve(LimitLayer::new(FirstConnOnly::new()).into_layer(http_service));
ready.send(()).unwrap();
serve.await;
}
#[derive(Clone)]
struct FirstConnOnly(Arc<AtomicBool>);
impl FirstConnOnly {
fn new() -> Self {
Self(Arc::new(AtomicBool::new(false)))
}
}
impl<Request> Policy<Request> for FirstConnOnly
where
Request: Send + 'static,
{
type Guard = ();
type Error = BoxError;
async fn check(&self, request: Request) -> PolicyResult<Request, Self::Guard, Self::Error> {
let output = match !self.0.swap(true, Ordering::AcqRel) {
true => PolicyOutput::Ready(()),
false => PolicyOutput::Abort(BoxError::from_static_str(
"only first connection is allowed",
)),
};
PolicyResult {
input: request,
output,
}
}
}