unify lock logic

1 files changed, 3 insertions, 7 deletions

diff --git a/src/transport/firewall.rs b/src/transport/firewall.rs
index b8e05cd..c8592bd 100644
--- a/src/transport/firewall.rs
+++ b/src/transport/firewall.rs
@@ -6,6 +6,7 @@ use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 use utoipa::ToSchema;
 
+use crate::locks::RwLockExt;
 use crate::transport::EncodedMessage;
 use crate::transport::{MessageHandler, ParsedMessage, Transport};
 
@@ -27,9 +28,7 @@ pub struct Firewall {
 
 impl Firewall {
     fn read(&self) -> RwLockReadGuard<'_, FirewallRules> {
-        self.rules
-            .read()
-            .unwrap_or_else(|poisoned| poisoned.into_inner())
+        self.rules.read_or_recover()
     }
 
     pub fn rules(&self) -> FirewallRules {
@@ -37,10 +36,7 @@ impl Firewall {
     }
 
     pub fn set_rules(&self, rules: FirewallRules) {
-        *self
-            .rules
-            .write()
-            .unwrap_or_else(|poisoned| poisoned.into_inner()) = rules;
+        *self.rules.write_or_recover() = rules;
     }
 
     pub fn blocks_inbound(&self, peer: &str) -> bool {