From eeba5647f9e7e7fc2a3f0c2ffec17f570bd7306a Mon Sep 17 00:00:00 2001
From: TotallyNot <44345987+TotallyNot@users.noreply.github.com>
Date: Thu, 4 Apr 2024 16:19:44 +0200
Subject: [PATCH] changed before/after hook signatures

---
 torn-key-pool/Cargo.toml  |  2 +-
 torn-key-pool/src/lib.rs  |  2 +-
 torn-key-pool/src/send.rs | 46 ++++++++++++++++++++++++++-------------
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/torn-key-pool/Cargo.toml b/torn-key-pool/Cargo.toml
index 1612f1d..816824f 100644
--- a/torn-key-pool/Cargo.toml
+++ b/torn-key-pool/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "torn-key-pool"
-version = "0.8.0"
+version = "0.8.1"
 edition = "2021"
 authors = ["Pyrit [2111649]"]
 license = "MIT"
diff --git a/torn-key-pool/src/lib.rs b/torn-key-pool/src/lib.rs
index 7b9c961..2e6d428 100644
--- a/torn-key-pool/src/lib.rs
+++ b/torn-key-pool/src/lib.rs
@@ -43,7 +43,7 @@ where
     }
 }
 
-pub trait ApiKey: Sync + Send + std::fmt::Debug + Clone {
+pub trait ApiKey: Sync + Send + std::fmt::Debug + Clone + 'static {
     type IdType: PartialEq + Eq + std::hash::Hash + Send + Sync + std::fmt::Debug + Clone;
 
     fn value(&self) -> &str;
diff --git a/torn-key-pool/src/send.rs b/torn-key-pool/src/send.rs
index 5b528d8..6572c9c 100644
--- a/torn-key-pool/src/send.rs
+++ b/torn-key-pool/src/send.rs
@@ -8,7 +8,8 @@ use torn_api::{
 };
 
 use crate::{
-    ApiKey, IntoSelector, KeyAction, KeyPoolError, KeyPoolExecutor, KeyPoolStorage, PoolOptions,
+    ApiKey, IntoSelector, KeyAction, KeyDomain, KeyPoolError, KeyPoolExecutor, KeyPoolStorage,
+    KeySelector, PoolOptions,
 };
 
 #[async_trait]
@@ -30,9 +31,11 @@ where
     {
         request.comment = self.options.comment.clone();
         if let Some(hook) = self.options.hooks_before.get(&std::any::TypeId::of::<A>()) {
-            let concrete = hook.downcast_ref::<BeforeHook<A>>().unwrap();
+            let concrete = hook
+                .downcast_ref::<BeforeHook<A, S::Key, S::Domain>>()
+                .unwrap();
 
-            (concrete.body)(&mut request);
+            (concrete.body)(&mut request, &self.selector);
         }
         loop {
             let key = self
@@ -58,9 +61,11 @@ where
                 Ok(res) => {
                     let res = res.into();
                     if let Some(hook) = self.options.hooks_after.get(&std::any::TypeId::of::<A>()) {
-                        let concrete = hook.downcast_ref::<AfterHook<A, S::Domain>>().unwrap();
+                        let concrete = hook
+                            .downcast_ref::<AfterHook<A, S::Key, S::Domain>>()
+                            .unwrap();
 
-                        match (concrete.body)(&res) {
+                        match (concrete.body)(&res, &self.selector) {
                             Err(KeyAction::Delete) => {
                                 self.storage
                                     .remove_key(key.selector())
@@ -156,20 +161,28 @@ where
 }
 
 #[allow(clippy::type_complexity)]
-pub struct BeforeHook<A>
+pub struct BeforeHook<A, K, D>
 where
     A: ApiSelection,
+    K: ApiKey,
+    D: KeyDomain,
 {
-    body: Box<dyn Fn(&mut ApiRequest<A>) + Send + Sync + 'static>,
+    body: Box<dyn Fn(&mut ApiRequest<A>, &KeySelector<K, D>) + Send + Sync + 'static>,
 }
 
 #[allow(clippy::type_complexity)]
-pub struct AfterHook<A, D>
+pub struct AfterHook<A, K, D>
 where
     A: ApiSelection,
-    D: crate::KeyDomain,
+    K: ApiKey,
+    D: KeyDomain,
 {
-    body: Box<dyn Fn(&A::Response) -> Result<(), crate::KeyAction<D>> + Send + Sync + 'static>,
+    body: Box<
+        dyn Fn(&A::Response, &KeySelector<K, D>) -> Result<(), crate::KeyAction<D>>
+            + Send
+            + Sync
+            + 'static,
+    >,
 }
 
 pub struct PoolBuilder<C, S>
@@ -202,7 +215,7 @@ where
 
     pub fn hook_before<A>(
         mut self,
-        hook: impl Fn(&mut ApiRequest<A>) + Send + Sync + 'static,
+        hook: impl Fn(&mut ApiRequest<A>, &KeySelector<S::Key, S::Domain>) + Send + Sync + 'static,
     ) -> Self
     where
         A: ApiSelection + 'static,
@@ -218,14 +231,17 @@ where
 
     pub fn hook_after<A>(
         mut self,
-        hook: impl Fn(&A::Response) -> Result<(), KeyAction<S::Domain>> + Send + Sync + 'static,
+        hook: impl Fn(&A::Response, &KeySelector<S::Key, S::Domain>) -> Result<(), KeyAction<S::Domain>>
+            + Send
+            + Sync
+            + 'static,
     ) -> Self
     where
         A: ApiSelection + 'static,
     {
         self.options.hooks_after.insert(
             std::any::TypeId::of::<A>(),
-            Box::new(AfterHook::<A, S::Domain> {
+            Box::new(AfterHook::<A, S::Key, S::Domain> {
                 body: Box::new(hook),
             }),
         );
@@ -331,7 +347,7 @@ mod test {
         let (storage, _) = setup(pool).await;
 
         let pool = PoolBuilder::new(reqwest::Client::default(), storage)
-            .hook_before::<torn_api::user::UserSelection>(|req| {
+            .hook_before::<torn_api::user::UserSelection>(|req, _s| {
                 req.selections.push("crimes");
             })
             .build();
@@ -345,7 +361,7 @@ mod test {
         let (storage, _) = setup(pool).await;
 
         let pool = PoolBuilder::new(reqwest::Client::default(), storage)
-            .hook_after::<torn_api::user::UserSelection>(|_res| Err(KeyAction::Delete))
+            .hook_after::<torn_api::user::UserSelection>(|_res, _s| Err(KeyAction::Delete))
             .build();
 
         let key = pool.storage.read_key(KeySelector::Id(1)).await.unwrap();