forest/rpc/
set_extension_layer.rs

1// Copyright 2019-2025 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use super::ApiPaths;
5use jsonrpsee::MethodResponse;
6use jsonrpsee::core::middleware::{Batch, BatchEntry, Notification};
7use jsonrpsee::server::middleware::rpc::RpcServiceT;
8use tower::Layer;
9
10/// JSON-RPC middleware layer for setting extensions in RPC requests
11#[derive(Clone)]
12pub(super) struct SetExtensionLayer {
13    pub path: ApiPaths,
14}
15
16impl<S> Layer<S> for SetExtensionLayer {
17    type Service = SetExtensionService<S>;
18
19    fn layer(&self, service: S) -> Self::Service {
20        SetExtensionService {
21            service,
22            path: self.path,
23        }
24    }
25}
26
27#[derive(Clone)]
28pub(super) struct SetExtensionService<S> {
29    service: S,
30    path: ApiPaths,
31}
32
33impl<S> RpcServiceT for SetExtensionService<S>
34where
35    S: RpcServiceT<MethodResponse = MethodResponse> + Send + Sync + Clone + 'static,
36{
37    type MethodResponse = S::MethodResponse;
38    type NotificationResponse = S::NotificationResponse;
39    type BatchResponse = S::BatchResponse;
40
41    fn call<'a>(
42        &self,
43        mut req: jsonrpsee::types::Request<'a>,
44    ) -> impl Future<Output = Self::MethodResponse> + Send + 'a {
45        req.extensions_mut().insert(self.path);
46        self.service.call(req)
47    }
48
49    fn batch<'a>(
50        &self,
51        mut batch: Batch<'a>,
52    ) -> impl Future<Output = Self::BatchResponse> + Send + 'a {
53        for req in batch.iter_mut() {
54            match req {
55                Ok(BatchEntry::Call(req)) => {
56                    req.extensions_mut().insert(self.path);
57                }
58                Ok(BatchEntry::Notification(n)) => {
59                    n.extensions_mut().insert(self.path);
60                }
61                Err(_) => {}
62            }
63        }
64        self.service.batch(batch)
65    }
66
67    fn notification<'a>(
68        &self,
69        mut n: Notification<'a>,
70    ) -> impl Future<Output = Self::NotificationResponse> + Send + 'a {
71        n.extensions_mut().insert(self.path);
72        self.service.notification(n)
73    }
74}