use crate::{generate_doc_comment, generate_doc_comments, naive_snake_case}; use proc_macro2::{Span, TokenStream}; use prost_build::{Method, Service}; use quote::quote; use syn::{Ident, Lit, LitStr}; pub(crate) fn generate(service: &Service, proto_path: &str) -> TokenStream { let methods = generate_methods(&service, proto_path); let server_service = quote::format_ident!("{}Server", service.name); let server_trait = quote::format_ident!("{}", service.name); let server_mod = quote::format_ident!("{}_server", naive_snake_case(&service.name)); let generated_trait = generate_trait(service, proto_path, server_trait.clone()); let service_doc = generate_doc_comments(&service.comments.leading); // Transport based implementations let path = format!("{}.{}", service.package, service.proto_name); let transport = generate_transport(&server_service, &server_trait, &path); quote! { /// Generated server implementations. pub mod #server_mod { #![allow(unused_variables, dead_code, missing_docs)] use tonic::codegen::*; #generated_trait #service_doc #[derive(Debug)] #[doc(hidden)] pub struct #server_service { inner: _Inner, } struct _Inner(Arc, Option); impl #server_service { pub fn new(inner: T) -> Self { let inner = Arc::new(inner); let inner = _Inner(inner, None); Self { inner } } pub fn with_interceptor(inner: T, interceptor: impl Into) -> Self { let inner = Arc::new(inner); let inner = _Inner(inner, Some(interceptor.into())); Self { inner } } } impl Service> for #server_service { type Response = http::Response; type Error = Never; type Future = BoxFuture; fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn call(&mut self, req: http::Request) -> Self::Future { let inner = self.inner.clone(); match req.uri().path() { #methods _ => Box::pin(async move { Ok(http::Response::builder() .status(200) .header("grpc-status", "12") .body(tonic::body::BoxBody::empty()) .unwrap()) }), } } } impl Clone for #server_service { fn clone(&self) -> Self { let inner = self.inner.clone(); Self { inner } } } impl Clone for _Inner { fn clone(&self) -> Self { Self(self.0.clone(), self.1.clone()) } } impl std::fmt::Debug for _Inner { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{:?}", self.0) } } #transport } } } fn generate_trait(service: &Service, proto_path: &str, server_trait: Ident) -> TokenStream { let methods = generate_trait_methods(service, proto_path); let trait_doc = generate_doc_comment(&format!( "Generated trait containing gRPC methods that should be implemented for use with {}Server.", service.name )); quote! { #trait_doc #[async_trait] pub trait #server_trait : Send + Sync + 'static { #methods } } } fn generate_trait_methods(service: &Service, proto_path: &str) -> TokenStream { let mut stream = TokenStream::new(); for method in &service.methods { let name = quote::format_ident!("{}", method.name); let (req_message, res_message) = crate::replace_wellknown(proto_path, &method); let method_doc = generate_doc_comments(&method.comments.leading); let method = match (method.client_streaming, method.server_streaming) { (false, false) => { quote! { #method_doc async fn #name(&self, request: tonic::Request<#req_message>) -> Result, tonic::Status>; } } (true, false) => { quote! { #method_doc async fn #name(&self, request: tonic::Request>) -> Result, tonic::Status>; } } (false, true) => { let stream = quote::format_ident!("{}Stream", method.proto_name); let stream_doc = generate_doc_comment(&format!( "Server streaming response type for the {} method.", method.proto_name )); quote! { #stream_doc type #stream: Stream> + Send + Sync + 'static; #method_doc async fn #name(&self, request: tonic::Request<#req_message>) -> Result, tonic::Status>; } } (true, true) => { let stream = quote::format_ident!("{}Stream", method.proto_name); let stream_doc = generate_doc_comment(&format!( "Server streaming response type for the {} method.", method.proto_name )); quote! { #stream_doc type #stream: Stream> + Send + Sync + 'static; #method_doc async fn #name(&self, request: tonic::Request>) -> Result, tonic::Status>; } } }; stream.extend(method); } stream } #[cfg(feature = "transport")] fn generate_transport( server_service: &syn::Ident, server_trait: &syn::Ident, service_name: &str, ) -> TokenStream { let service_name = syn::LitStr::new(service_name, proc_macro2::Span::call_site()); quote! { impl tonic::transport::NamedService for #server_service { const NAME: &'static str = #service_name; } } } #[cfg(not(feature = "transport"))] fn generate_transport( _server_service: &syn::Ident, _server_trait: &syn::Ident, _service_name: &str, ) -> TokenStream { TokenStream::new() } fn generate_methods(service: &Service, proto_path: &str) -> TokenStream { let mut stream = TokenStream::new(); for method in &service.methods { let path = format!( "/{}.{}/{}", service.package, service.proto_name, method.proto_name ); let method_path = Lit::Str(LitStr::new(&path, Span::call_site())); let ident = quote::format_ident!("{}", method.name); let server_trait = quote::format_ident!("{}", service.name); let method_stream = match (method.client_streaming, method.server_streaming) { (false, false) => generate_unary(method, ident, proto_path, server_trait), (false, true) => { generate_server_streaming(method, ident.clone(), proto_path, server_trait) } (true, false) => { generate_client_streaming(method, ident.clone(), proto_path, server_trait) } (true, true) => generate_streaming(method, ident.clone(), proto_path, server_trait), }; let method = quote! { #method_path => { #method_stream } }; stream.extend(method); } stream } fn generate_unary( method: &Method, method_ident: Ident, proto_path: &str, server_trait: Ident, ) -> TokenStream { let service_ident = quote::format_ident!("{}Svc", method.proto_name); let (request, response) = crate::replace_wellknown(proto_path, &method); quote! { struct #service_ident(pub Arc); impl tonic::server::UnaryService<#request> for #service_ident { type Response = #response; type Future = BoxFuture, tonic::Status>; fn call(&mut self, request: tonic::Request<#request>) -> Self::Future { let inner = self.0.clone(); let fut = async move { inner.#method_ident(request).await }; Box::pin(fut) } } let inner = self.inner.clone(); let fut = async move { let interceptor = inner.1.clone(); let inner = inner.0; let method = #service_ident(inner); let codec = tonic::codec::ProstCodec::default(); let mut grpc = if let Some(interceptor) = interceptor { tonic::server::Grpc::with_interceptor(codec, interceptor) } else { tonic::server::Grpc::new(codec) }; let res = grpc.unary(method, req).await; Ok(res) }; Box::pin(fut) } } fn generate_server_streaming( method: &Method, method_ident: Ident, proto_path: &str, server_trait: Ident, ) -> TokenStream { let service_ident = quote::format_ident!("{}Svc", method.proto_name); let (request, response) = crate::replace_wellknown(proto_path, &method); let response_stream = quote::format_ident!("{}Stream", method.proto_name); quote! { struct #service_ident(pub Arc); impl tonic::server::ServerStreamingService<#request> for #service_ident { type Response = #response; type ResponseStream = T::#response_stream; type Future = BoxFuture, tonic::Status>; fn call(&mut self, request: tonic::Request<#request>) -> Self::Future { let inner = self.0.clone(); let fut = async move { inner.#method_ident(request).await }; Box::pin(fut) } } let inner = self.inner.clone(); let fut = async move { let interceptor = inner.1; let inner = inner.0; let method = #service_ident(inner); let codec = tonic::codec::ProstCodec::default(); let mut grpc = if let Some(interceptor) = interceptor { tonic::server::Grpc::with_interceptor(codec, interceptor) } else { tonic::server::Grpc::new(codec) }; let res = grpc.server_streaming(method, req).await; Ok(res) }; Box::pin(fut) } } fn generate_client_streaming( method: &Method, method_ident: Ident, proto_path: &str, server_trait: Ident, ) -> TokenStream { let service_ident = quote::format_ident!("{}Svc", method.proto_name); let (request, response) = crate::replace_wellknown(proto_path, &method); quote! { struct #service_ident(pub Arc); impl tonic::server::ClientStreamingService<#request> for #service_ident { type Response = #response; type Future = BoxFuture, tonic::Status>; fn call(&mut self, request: tonic::Request>) -> Self::Future { let inner = self.0.clone(); let fut = async move { inner.#method_ident(request).await }; Box::pin(fut) } } let inner = self.inner.clone(); let fut = async move { let interceptor = inner.1; let inner = inner.0; let method = #service_ident(inner); let codec = tonic::codec::ProstCodec::default(); let mut grpc = if let Some(interceptor) = interceptor { tonic::server::Grpc::with_interceptor(codec, interceptor) } else { tonic::server::Grpc::new(codec) }; let res = grpc.client_streaming(method, req).await; Ok(res) }; Box::pin(fut) } } fn generate_streaming( method: &Method, method_ident: Ident, proto_path: &str, server_trait: Ident, ) -> TokenStream { let service_ident = quote::format_ident!("{}Svc", method.proto_name); let (request, response) = crate::replace_wellknown(proto_path, &method); let response_stream = quote::format_ident!("{}Stream", method.proto_name); quote! { struct #service_ident(pub Arc); impl tonic::server::StreamingService<#request> for #service_ident { type Response = #response; type ResponseStream = T::#response_stream; type Future = BoxFuture, tonic::Status>; fn call(&mut self, request: tonic::Request>) -> Self::Future { let inner = self.0.clone(); let fut = async move { inner.#method_ident(request).await }; Box::pin(fut) } } let inner = self.inner.clone(); let fut = async move { let interceptor = inner.1; let inner = inner.0; let method = #service_ident(inner); let codec = tonic::codec::ProstCodec::default(); let mut grpc = if let Some(interceptor) = interceptor { tonic::server::Grpc::with_interceptor(codec, interceptor) } else { tonic::server::Grpc::new(codec) }; let res = grpc.streaming(method, req).await; Ok(res) }; Box::pin(fut) } }