use crate::{generate_doc_comment, generate_doc_comments}; 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_make_service = quote::format_ident!("{}Server", service.name); let server_service = quote::format_ident!("{}ServerSvc", service.name); let server_trait = quote::format_ident!("{}", service.name); let generated_trait = generate_trait(service, proto_path, server_trait.clone()); let service_doc = generate_doc_comments(&service.comments.leading); let server_new_doc = generate_doc_comment(&format!( "Create a new {} from a type that implements {}.", server_make_service, server_trait )); quote! { #generated_trait #service_doc #[derive(Clone, Debug)] pub struct #server_make_service { inner: Arc, } #[derive(Clone, Debug)] #[doc(hidden)] pub struct #server_service { inner: Arc, } impl #server_make_service { #server_new_doc pub fn new(inner: T) -> Self { let inner = Arc::new(inner); Self::from_shared(inner) } pub fn from_shared(inner: Arc) -> Self { Self { inner } } } impl #server_service { pub fn new(inner: Arc) -> Self { Self { inner } } } impl Service for #server_make_service { type Response = #server_service; type Error = Never; type Future = Ready>; fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn call(&mut self, _: R) -> Self::Future { ok(#server_service::new(self.inner.clone())) } } 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()) }), } } } } } 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 = crate::replace_wellknown(proto_path, &method.input_type); let res_message = crate::replace_wellknown(proto_path, &method.output_type); 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 + '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 + 'static; #method_doc async fn #name(&self, request: tonic::Request>) -> Result, tonic::Status>; } } }; stream.extend(method); } stream } 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 = Ident::new(&method.proto_name, Span::call_site()); let request = crate::replace_wellknown(proto_path, &method.input_type); let response = crate::replace_wellknown(proto_path, &method.output_type); 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 method = #service_ident(inner); let codec = tonic::codec::ProstCodec::new(); let mut grpc = 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 = Ident::new(&method.proto_name, Span::call_site()); let request = crate::replace_wellknown(proto_path, &method.input_type); let response = crate::replace_wellknown(proto_path, &method.output_type); 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 method = #service_ident(inner); let codec = tonic::codec::ProstCodec::new(); let mut grpc = 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 = Ident::new(&method.proto_name, Span::call_site()); let request = crate::replace_wellknown(proto_path, &method.input_type); let response = crate::replace_wellknown(proto_path, &method.output_type); 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 method = #service_ident(inner); let codec = tonic::codec::ProstCodec::new(); let mut grpc = 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 = Ident::new(&method.proto_name, Span::call_site()); let request = crate::replace_wellknown(proto_path, &method.input_type); let response = crate::replace_wellknown(proto_path, &method.output_type); 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 method = #service_ident(inner); let codec = tonic::codec::ProstCodec::new(); let mut grpc = tonic::server::Grpc::new(codec); let res = grpc.streaming(method, req).await; Ok(res) }; Box::pin(fut) } }