1 files changed, 17 insertions, 22 deletions

diff --git a/src/webmentions/mod.rs b/src/webmentions/mod.rs
index 3e9b094..d5a617e 100644
--- a/src/webmentions/mod.rs
+++ b/src/webmentions/mod.rs
@@ -1,4 +1,4 @@
-use axum::{Form, response::{IntoResponse, Response}, Extension};
+use axum::{extract::{FromRef, State}, response::{IntoResponse, Response}, routing::post, Form};
 use axum::http::StatusCode;
 use tracing::error;
 
@@ -20,7 +20,7 @@ impl queue::PostgresJobItem for Webmention {
 }
 
 async fn accept_webmention<Q: JobQueue<Webmention>>(
-    Extension(queue): Extension<Q>,
+    State(queue): State<Q>,
     Form(webmention): Form<Webmention>,
 ) -> Response {
     if let Err(err) = webmention.source.parse::<url::Url>() {
@@ -31,27 +31,16 @@ async fn accept_webmention<Q: JobQueue<Webmention>>(
     }
 
     match queue.put(&webmention).await {
-        Ok(id) => StatusCode::ACCEPTED.into_response(),
+        Ok(_id) => StatusCode::ACCEPTED.into_response(),
         Err(err) => (StatusCode::INTERNAL_SERVER_ERROR, [
             ("Content-Type", "text/plain")
         ], err.to_string()).into_response()
     }
 }
 
-pub fn router<Q: JobQueue<Webmention>, S: Storage + 'static>(
-    queue: Q, db: S, http: reqwest::Client,
-    cancellation_token: tokio_util::sync::CancellationToken
-) -> (axum::Router, SupervisedTask) {
-    // Automatically spawn a background task to handle webmentions
-    let bgtask_handle = supervised_webmentions_task(queue.clone(), db, http, cancellation_token);
-
-    let router = axum::Router::new()
-        .route("/.kittybox/webmention",
-            axum::routing::post(accept_webmention::<Q>)
-        )
-        .layer(Extension(queue));
-
-    (router, bgtask_handle)
+pub fn router<St: Clone + Send + Sync + 'static, Q: JobQueue<Webmention> + FromRef<St>>() -> axum::Router<St> {
+    axum::Router::new()
+        .route("/.kittybox/webmention", post(accept_webmention::<Q>))
 }
 
 #[derive(thiserror::Error, Debug)]
@@ -184,10 +173,16 @@ async fn process_webmentions_from_queue<Q: JobQueue<Webmention>, S: Storage + 's
     unreachable!()
 }
 
-fn supervised_webmentions_task<Q: JobQueue<Webmention>, S: Storage + 'static>(
-    queue: Q, db: S,
-    http: reqwest::Client,
+pub fn supervised_webmentions_task<St: Send + Sync + 'static, S: Storage + FromRef<St> + 'static, Q: JobQueue<Webmention> + FromRef<St> + 'static>(
+    state: &St,
     cancellation_token: tokio_util::sync::CancellationToken
-) -> SupervisedTask {
-    supervisor::<Error<Q::Error>, _, _>(move || process_webmentions_from_queue(queue.clone(), db.clone(), http.clone()), cancellation_token)
+) -> SupervisedTask
+where reqwest::Client: FromRef<St>
+{
+    let queue = Q::from_ref(state);
+    let storage = S::from_ref(state);
+    let http = reqwest::Client::from_ref(state);
+    supervisor::<Error<Q::Error>, _, _>(move || process_webmentions_from_queue(
+        queue.clone(), storage.clone(), http.clone()
+    ), cancellation_token)
 }