tonic/service/
interceptor.rs

1/*
2 *
3 * Copyright 2025 gRPC authors.
4 *
5 * Permission is hereby granted, free of charge, to any person obtaining a copy
6 * of this software and associated documentation files (the "Software"), to
7 * deal in the Software without restriction, including without limitation the
8 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
9 * sell copies of the Software, and to permit persons to whom the Software is
10 * furnished to do so, subject to the following conditions:
11 *
12 * The above copyright notice and this permission notice shall be included in
13 * all copies or substantial portions of the Software.
14 *
15 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
20 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
21 * IN THE SOFTWARE.
22 *
23 */
24
25//! gRPC interceptors which are a kind of middleware.
26//!
27//! See [`Interceptor`] for more details.
28
29use crate::{Status, request::SanitizeHeaders};
30use pin_project::pin_project;
31use std::{
32    fmt,
33    future::Future,
34    pin::Pin,
35    task::{Context, Poll},
36};
37use tower_layer::Layer;
38use tower_service::Service;
39
40/// A gRPC interceptor.
41///
42/// gRPC interceptors are similar to middleware but have less flexibility. An interceptor allows
43/// you to do two main things, one is to add/remove/check items in the `MetadataMap` of each
44/// request. Two, cancel a request with a `Status`.
45///
46/// Any function that satisfies the bound `FnMut(Request<()>) -> Result<Request<()>, Status>` can be
47/// used as an `Interceptor`.
48///
49/// An interceptor can be used on both the server and client side through the `tonic-build` crate's
50/// generated structs.
51///
52/// See the [interceptor example][example] for more details.
53///
54/// If you need more powerful middleware, [tower] is the recommended approach. You can find
55/// examples of how to use tower with tonic [here][tower-example].
56///
57/// Additionally, interceptors is not the recommended way to add logging to your service. For that
58/// a [tower] middleware is more appropriate since it can also act on the response. For example
59/// tower-http's [`Trace`](https://docs.rs/tower-http/latest/tower_http/trace/index.html)
60/// middleware supports gRPC out of the box.
61///
62/// [tower]: https://crates.io/crates/tower
63/// [example]: https://github.com/grpc/grpc-rust/tree/master/examples/src/interceptor
64/// [tower-example]: https://github.com/grpc/grpc-rust/tree/master/examples/src/tower
65pub trait Interceptor {
66    /// Intercept a request before it is sent, optionally cancelling it.
67    fn call(&mut self, request: crate::Request<()>) -> Result<crate::Request<()>, Status>;
68}
69
70impl<F> Interceptor for F
71where
72    F: FnMut(crate::Request<()>) -> Result<crate::Request<()>, Status>,
73{
74    fn call(&mut self, request: crate::Request<()>) -> Result<crate::Request<()>, Status> {
75        self(request)
76    }
77}
78
79/// A gRPC interceptor that can be used as a [`Layer`],
80///
81/// See [`Interceptor`] for more details.
82#[derive(Debug, Clone, Copy)]
83pub struct InterceptorLayer<I> {
84    interceptor: I,
85}
86
87impl<I> InterceptorLayer<I> {
88    /// Create a new interceptor layer.
89    ///
90    /// See [`Interceptor`] for more details.
91    pub fn new(interceptor: I) -> Self {
92        Self { interceptor }
93    }
94}
95
96impl<S, I> Layer<S> for InterceptorLayer<I>
97where
98    I: Clone,
99{
100    type Service = InterceptedService<S, I>;
101
102    fn layer(&self, service: S) -> Self::Service {
103        InterceptedService::new(service, self.interceptor.clone())
104    }
105}
106
107/// A service wrapped in an interceptor middleware.
108///
109/// See [`Interceptor`] for more details.
110#[derive(Clone, Copy)]
111pub struct InterceptedService<S, I> {
112    inner: S,
113    interceptor: I,
114}
115
116impl<S, I> InterceptedService<S, I> {
117    /// Create a new `InterceptedService` that wraps `S` and intercepts each request with the
118    /// function `F`.
119    pub fn new(service: S, interceptor: I) -> Self {
120        Self {
121            inner: service,
122            interceptor,
123        }
124    }
125}
126
127impl<S, I> fmt::Debug for InterceptedService<S, I>
128where
129    S: fmt::Debug,
130{
131    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
132        f.debug_struct("InterceptedService")
133            .field("inner", &self.inner)
134            .field("f", &format_args!("{}", std::any::type_name::<I>()))
135            .finish()
136    }
137}
138
139impl<S, I, ReqBody, ResBody> Service<http::Request<ReqBody>> for InterceptedService<S, I>
140where
141    S: Service<http::Request<ReqBody>, Response = http::Response<ResBody>>,
142    I: Interceptor,
143{
144    type Response = http::Response<ResponseBody<ResBody>>;
145    type Error = S::Error;
146    type Future = ResponseFuture<S::Future>;
147
148    #[inline]
149    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
150        self.inner.poll_ready(cx)
151    }
152
153    fn call(&mut self, req: http::Request<ReqBody>) -> Self::Future {
154        // It is bad practice to modify the body (i.e. Message) of the request via an interceptor.
155        // To avoid exposing the body of the request to the interceptor function, we first remove it
156        // here, allow the interceptor to modify the metadata and extensions, and then recreate the
157        // HTTP request with the body. Tonic requests do not preserve the URI, HTTP version, and
158        // HTTP method of the HTTP request, so we extract them here and then add them back in below.
159        let uri = req.uri().clone();
160        let method = req.method().clone();
161        let version = req.version();
162        let req = crate::Request::from_http(req);
163        let (metadata, extensions, msg) = req.into_parts();
164
165        match self
166            .interceptor
167            .call(crate::Request::from_parts(metadata, extensions, ()))
168        {
169            Ok(req) => {
170                let (metadata, extensions, _) = req.into_parts();
171                let req = crate::Request::from_parts(metadata, extensions, msg);
172                let req = req.into_http(uri, method, version, SanitizeHeaders::No);
173                ResponseFuture::future(self.inner.call(req))
174            }
175            Err(status) => ResponseFuture::status(status),
176        }
177    }
178}
179
180// required to use `InterceptedService` with `Router`
181impl<S, I> crate::server::NamedService for InterceptedService<S, I>
182where
183    S: crate::server::NamedService,
184{
185    const NAME: &'static str = S::NAME;
186}
187
188/// Response future for [`InterceptedService`].
189#[pin_project]
190#[derive(Debug)]
191pub struct ResponseFuture<F> {
192    #[pin]
193    kind: Kind<F>,
194}
195
196impl<F> ResponseFuture<F> {
197    fn future(future: F) -> Self {
198        Self {
199            kind: Kind::Future(future),
200        }
201    }
202
203    fn status(status: Status) -> Self {
204        Self {
205            kind: Kind::Status(Some(status)),
206        }
207    }
208}
209
210#[pin_project(project = KindProj)]
211#[derive(Debug)]
212enum Kind<F> {
213    Future(#[pin] F),
214    Status(Option<Status>),
215}
216
217impl<F, E, B> Future for ResponseFuture<F>
218where
219    F: Future<Output = Result<http::Response<B>, E>>,
220{
221    type Output = Result<http::Response<ResponseBody<B>>, E>;
222
223    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
224        match self.project().kind.project() {
225            KindProj::Future(future) => future.poll(cx).map_ok(|res| res.map(ResponseBody::wrap)),
226            KindProj::Status(status) => {
227                let (parts, ()) = status.take().unwrap().into_http::<()>().into_parts();
228                let response = http::Response::from_parts(parts, ResponseBody::<B>::empty());
229                Poll::Ready(Ok(response))
230            }
231        }
232    }
233}
234
235/// Response body for [`InterceptedService`].
236#[pin_project]
237#[derive(Debug)]
238pub struct ResponseBody<B> {
239    #[pin]
240    kind: ResponseBodyKind<B>,
241}
242
243#[pin_project(project = ResponseBodyKindProj)]
244#[derive(Debug)]
245enum ResponseBodyKind<B> {
246    Empty,
247    Wrap(#[pin] B),
248}
249
250impl<B> ResponseBody<B> {
251    fn new(kind: ResponseBodyKind<B>) -> Self {
252        Self { kind }
253    }
254
255    fn empty() -> Self {
256        Self::new(ResponseBodyKind::Empty)
257    }
258
259    fn wrap(body: B) -> Self {
260        Self::new(ResponseBodyKind::Wrap(body))
261    }
262}
263
264impl<B: http_body::Body> http_body::Body for ResponseBody<B> {
265    type Data = B::Data;
266    type Error = B::Error;
267
268    fn poll_frame(
269        self: Pin<&mut Self>,
270        cx: &mut Context<'_>,
271    ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
272        match self.project().kind.project() {
273            ResponseBodyKindProj::Empty => Poll::Ready(None),
274            ResponseBodyKindProj::Wrap(body) => body.poll_frame(cx),
275        }
276    }
277
278    fn size_hint(&self) -> http_body::SizeHint {
279        match &self.kind {
280            ResponseBodyKind::Empty => http_body::SizeHint::with_exact(0),
281            ResponseBodyKind::Wrap(body) => body.size_hint(),
282        }
283    }
284
285    fn is_end_stream(&self) -> bool {
286        match &self.kind {
287            ResponseBodyKind::Empty => true,
288            ResponseBodyKind::Wrap(body) => body.is_end_stream(),
289        }
290    }
291}
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296    use tower::ServiceExt;
297
298    #[tokio::test]
299    async fn doesnt_remove_headers_from_requests() {
300        let svc = tower::service_fn(|request: http::Request<()>| async move {
301            assert_eq!(
302                request
303                    .headers()
304                    .get("user-agent")
305                    .expect("missing in leaf service"),
306                "test-tonic"
307            );
308
309            Ok::<_, Status>(http::Response::new(()))
310        });
311
312        let svc = InterceptedService::new(svc, |request: crate::Request<()>| {
313            assert_eq!(
314                request
315                    .metadata()
316                    .get("user-agent")
317                    .expect("missing in interceptor"),
318                "test-tonic"
319            );
320
321            Ok(request)
322        });
323
324        let request = http::Request::builder()
325            .header("user-agent", "test-tonic")
326            .body(())
327            .unwrap();
328
329        svc.oneshot(request).await.unwrap();
330    }
331
332    #[tokio::test]
333    async fn handles_intercepted_status_as_response() {
334        let message = "Blocked by the interceptor";
335        let expected = Status::permission_denied(message).into_http::<()>();
336
337        let svc = tower::service_fn(|_: http::Request<()>| async {
338            Ok::<_, Status>(http::Response::new(()))
339        });
340
341        let svc = InterceptedService::new(svc, |_: crate::Request<()>| {
342            Err(Status::permission_denied(message))
343        });
344
345        let request = http::Request::builder().body(()).unwrap();
346        let response = svc.oneshot(request).await.unwrap();
347
348        assert_eq!(expected.status(), response.status());
349        assert_eq!(expected.version(), response.version());
350        assert_eq!(expected.headers(), response.headers());
351    }
352
353    #[tokio::test]
354    async fn doesnt_change_http_method() {
355        let svc = tower::service_fn(|request: http::Request<()>| async move {
356            assert_eq!(request.method(), http::Method::OPTIONS);
357
358            Ok::<_, hyper::Error>(hyper::Response::new(()))
359        });
360
361        let svc = InterceptedService::new(svc, Ok);
362
363        let request = http::Request::builder()
364            .method(http::Method::OPTIONS)
365            .body(())
366            .unwrap();
367
368        svc.oneshot(request).await.unwrap();
369    }
370}
tonic/service/interceptor.rs

tonic/service/
interceptor.rs
