From 8587b3e9b3bd99b798d46e0c09842c3294a7be17 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Mon, 23 Aug 2021 10:30:26 +0200
Subject: [PATCH] feat: add basic policy middleware (#400)

* feat: initial policy design

* docs: add some docs
---
 ethers-middleware/src/lib.rs    |   5 ++
 ethers-middleware/src/policy.rs | 111 ++++++++++++++++++++++++++++++++
 2 files changed, 116 insertions(+)
 create mode 100644 ethers-middleware/src/policy.rs

diff --git a/ethers-middleware/src/lib.rs b/ethers-middleware/src/lib.rs
index c779a042..7d4acc49 100644
--- a/ethers-middleware/src/lib.rs
+++ b/ethers-middleware/src/lib.rs
@@ -78,3 +78,8 @@ pub mod transformer;
 /// instead of using eth_sendTransaction and eth_sign
 pub mod signer;
 pub use signer::SignerMiddleware;
+
+/// The [Policy](crate::PolicyMiddleware) is used to ensure transactions comply with the rules
+/// configured in the `PolicyMiddleware` before sending them.
+pub mod policy;
+pub use policy::PolicyMiddleware;
diff --git a/ethers-middleware/src/policy.rs b/ethers-middleware/src/policy.rs
new file mode 100644
index 00000000..ef99e279
--- /dev/null
+++ b/ethers-middleware/src/policy.rs
@@ -0,0 +1,111 @@
+use ethers_core::types::{transaction::eip2718::TypedTransaction, BlockId};
+use ethers_providers::{FromErr, Middleware, PendingTransaction};
+
+use async_trait::async_trait;
+use std::fmt::Debug;
+use thiserror::Error;
+
+/// Basic trait to ensure that transactions about to be sent follow certain rules.
+#[async_trait]
+pub trait Policy: Sync + Send + Debug {
+    type Error: Sync + Send + Debug;
+
+    /// Evaluates the transactions.
+    ///
+    /// Returns Ok with the `tx` or an Err otherwise.
+    async fn ensure_can_send(&self, tx: TypedTransaction) -> Result<TypedTransaction, Self::Error>;
+}
+
+/// A policy that does not restrict anything.
+#[derive(Debug, Clone, Copy)]
+pub struct AllowEverything;
+
+#[async_trait]
+impl Policy for AllowEverything {
+    type Error = ();
+
+    async fn ensure_can_send(&self, tx: TypedTransaction) -> Result<TypedTransaction, Self::Error> {
+        Ok(tx)
+    }
+}
+
+/// A policy that rejects all transactions.
+#[derive(Debug, Clone, Copy)]
+pub struct RejectEverything;
+
+#[async_trait]
+impl Policy for RejectEverything {
+    type Error = ();
+
+    async fn ensure_can_send(&self, _: TypedTransaction) -> Result<TypedTransaction, Self::Error> {
+        Err(())
+    }
+}
+
+/// Middleware used to enforce certain policies for transactions.
+#[derive(Clone, Debug)]
+pub struct PolicyMiddleware<M, P> {
+    pub(crate) inner: M,
+    pub(crate) policy: P,
+}
+
+impl<M: Middleware, P: Policy> FromErr<M::Error> for PolicyMiddlewareError<M, P> {
+    fn from(src: M::Error) -> PolicyMiddlewareError<M, P> {
+        PolicyMiddlewareError::MiddlewareError(src)
+    }
+}
+
+impl<M, P> PolicyMiddleware<M, P>
+where
+    M: Middleware,
+    P: Policy,
+{
+    /// Creates a new client from the provider and policy.
+    pub fn new(inner: M, policy: P) -> Self {
+        Self { inner, policy }
+    }
+}
+
+#[derive(Error, Debug)]
+/// Error thrown when the client interacts with the policy middleware.
+pub enum PolicyMiddlewareError<M: Middleware, P: Policy> {
+    /// Thrown when the internal policy errors
+    #[error("{0:?}")]
+    PolicyError(P::Error),
+    /// Thrown when an internal middleware errors
+    #[error(transparent)]
+    MiddlewareError(M::Error),
+}
+
+#[async_trait]
+impl<M, P> Middleware for PolicyMiddleware<M, P>
+where
+    M: Middleware,
+    P: Policy,
+{
+    type Error = PolicyMiddlewareError<M, P>;
+    type Provider = M::Provider;
+    type Inner = M;
+
+    fn inner(&self) -> &M {
+        &self.inner
+    }
+
+    /// This ensures the tx complies with the registered policy.
+    /// If so then this simply delegates the transaction to the inner middleware
+    async fn send_transaction<T: Into<TypedTransaction> + Send + Sync>(
+        &self,
+        tx: T,
+        block: Option<BlockId>,
+    ) -> Result<PendingTransaction<'_, Self::Provider>, Self::Error> {
+        let tx = self
+            .policy
+            .ensure_can_send(tx.into())
+            .await
+            .map_err(PolicyMiddlewareError::PolicyError)?;
+        self.inner
+            .send_transaction(tx, block)
+            .await
+            .map_err(PolicyMiddlewareError::MiddlewareError)
+    }
+}