CODEWITHSUNDEEP
rustrust-tonictower
Teimuraz
Teimuraz

Reputation: 9325

Rust Tonic + Tower: How to apply middleware to specific services only?

I developed middleware using Tower and have been applying it to tonic services via the layer method.

However layer applies middleware to all services.

I need to apply my middleware to particular service. With my current implementation grpc returns 12 error code (not implemented) for service being wrapped into middleware .add_service(AuthMiddleware::new(grpc_project_service), this middleware and grpc_project_service is not invoked this way (only when using via layer).

Is there a way to achieve this selective middleware application? I've been unable to find relevant documentation on this.

Middleware

use std::task::{Context, Poll};

use tower::{Layer, Service};

#[derive(Debug, Clone, Default)]
pub struct AuthMiddlewareLayer;

impl<S> Layer<S> for AuthMiddlewareLayer {
    type Service = AuthMiddleware<S>;

    fn layer(&self, inner: S) -> Self::Service {
        AuthMiddleware::new(inner)
    }
}

#[derive(Debug, Clone)]
pub struct AuthMiddleware<S> {
    inner: S,
}

impl<S> AuthMiddleware<S> {
    pub fn new(inner: S) -> Self {
        AuthMiddleware {
            inner
        }
    }
}

impl<S, Req> Service<Req> for AuthMiddleware<S>
    where S: Service<Req>,

{
    type Response = S::Response;
    type Error = S::Error;
    type Future = S::Future;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        println!("AuthMiddleware poll_ready called");
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Req) -> Self::Future {
        print!("Im here!!!");
        self.inner.call(req)
    }
}

Middleware usage (main.rs)

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    
    ...

    let project_api = ProjectApi { project_service };
    let grpc_project_service = grpc_project::ProjectServiceServer::new(project_api);

    let auth_api = AuthApi { auth_service };
    let grpc_auth_service = grpc_auth::AuthServiceServer::new(auth_api);

    ...

    let reflection_service = tonic_reflection::server::Builder::configure()
        .register_encoded_file_descriptor_set(FILE_DESCRIPTOR_SET)
        .build()
        .unwrap();

    Server::builder()
        // Works, however applies to all services.
        // .layer(AuthMiddlewareLayer::default())
        .add_service(grpc_auth_service)
        // Doesn't work, grpc returns 12 error code (not implemented), AuthMiddleware and grpc_project_service are not invoked at all.
        .add_service(AuthMiddleware::new(grpc_project_service))
        .add_service(reflection_service)
        .serve(grpc_addr)
        .await
        .tap_err(|e| error!("Cannot start grpc server, error: {}", e))?;

    Ok(())
}

impl NamedService for AuthMiddleware<ProjectServiceServer<ProjectApi>> {
    const NAME: &'static str = "ProjectService";
}

Versions:

tonic = "0.11.0"
tower = "0.4.13"
hyper = "1.2.0"

Upvotes: 0

Views: 689

Answers (1)

Teimuraz
Teimuraz

Reputation: 9325

Found the answer: Turns out I had to implemented NamedService for middleware, since tonic uses name under the hood to route requests...

impl<S> NamedService for AuthMiddleware<S>
    where
        S: NamedService {

    const NAME: &'static str = S::NAME;
}

Update Create small lib that simplifies creation of async interceptors and middlewares (Full interception of service call, adding custom logic before and after service call), using tower. https://github.com/teimuraz/tonic-middleware

Upvotes: 0

Related Questions