1 use super::{Method, Service}; 2 use crate::{generate_doc_comment, generate_doc_comments, naive_snake_case}; 3 use proc_macro2::{Span, TokenStream}; 4 use quote::quote; 5 use syn::{Ident, Lit, LitStr}; 6 7 /// Generate service for Server. 8 /// 9 /// This takes some `Service` and will generate a `TokenStream` that contains 10 /// a public module containing the server service and handler trait. 11 pub fn generate<T: Service>(service: &T, proto_path: &str) -> TokenStream { 12 let methods = generate_methods(service, proto_path); 13 14 let server_service = quote::format_ident!("{}Server", service.name()); 15 let server_trait = quote::format_ident!("{}", service.name()); 16 let server_mod = quote::format_ident!("{}_server", naive_snake_case(&service.name())); 17 let generated_trait = generate_trait(service, proto_path, server_trait.clone()); 18 let service_doc = generate_doc_comments(service.comment()); 19 20 // Transport based implementations 21 let path = format!("{}.{}", service.package(), service.identifier()); 22 let transport = generate_transport(&server_service, &server_trait, &path); 23 24 quote! { 25 /// Generated server implementations. 26 pub mod #server_mod { 27 #![allow(unused_variables, dead_code, missing_docs)] 28 use tonic::codegen::*; 29 30 #generated_trait 31 32 #service_doc 33 #[derive(Debug)] 34 pub struct #server_service<T: #server_trait> { 35 inner: _Inner<T>, 36 } 37 38 struct _Inner<T>(Arc<T>, Option<tonic::Interceptor>); 39 40 impl<T: #server_trait> #server_service<T> { 41 pub fn new(inner: T) -> Self { 42 let inner = Arc::new(inner); 43 let inner = _Inner(inner, None); 44 Self { inner } 45 } 46 47 pub fn with_interceptor(inner: T, interceptor: impl Into<tonic::Interceptor>) -> Self { 48 let inner = Arc::new(inner); 49 let inner = _Inner(inner, Some(interceptor.into())); 50 Self { inner } 51 } 52 } 53 54 impl<T, B> Service<http::Request<B>> for #server_service<T> 55 where 56 T: #server_trait, 57 B: HttpBody + Send + Sync + 'static, 58 B::Error: Into<StdError> + Send + 'static, 59 { 60 type Response = http::Response<tonic::body::BoxBody>; 61 type Error = Never; 62 type Future = BoxFuture<Self::Response, Self::Error>; 63 64 fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { 65 Poll::Ready(Ok(())) 66 } 67 68 fn call(&mut self, req: http::Request<B>) -> Self::Future { 69 let inner = self.inner.clone(); 70 71 match req.uri().path() { 72 #methods 73 74 _ => Box::pin(async move { 75 Ok(http::Response::builder() 76 .status(200) 77 .header("grpc-status", "12") 78 .body(tonic::body::BoxBody::empty()) 79 .unwrap()) 80 }), 81 } 82 } 83 } 84 85 impl<T: #server_trait> Clone for #server_service<T> { 86 fn clone(&self) -> Self { 87 let inner = self.inner.clone(); 88 Self { inner } 89 } 90 } 91 92 impl<T: #server_trait> Clone for _Inner<T> { 93 fn clone(&self) -> Self { 94 Self(self.0.clone(), self.1.clone()) 95 } 96 } 97 98 impl<T: std::fmt::Debug> std::fmt::Debug for _Inner<T> { 99 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 100 write!(f, "{:?}", self.0) 101 } 102 } 103 104 #transport 105 } 106 } 107 } 108 109 fn generate_trait<T: Service>(service: &T, proto_path: &str, server_trait: Ident) -> TokenStream { 110 let methods = generate_trait_methods(service, proto_path); 111 let trait_doc = generate_doc_comment(&format!( 112 "Generated trait containing gRPC methods that should be implemented for use with {}Server.", 113 service.name() 114 )); 115 116 quote! { 117 #trait_doc 118 #[async_trait] 119 pub trait #server_trait : Send + Sync + 'static { 120 #methods 121 } 122 } 123 } 124 125 fn generate_trait_methods<T: Service>(service: &T, proto_path: &str) -> TokenStream { 126 let mut stream = TokenStream::new(); 127 128 for method in service.methods() { 129 let name = quote::format_ident!("{}", method.name()); 130 131 let (req_message, res_message) = method.request_response_name(proto_path); 132 133 let method_doc = generate_doc_comments(method.comment()); 134 135 let method = match (method.client_streaming(), method.server_streaming()) { 136 (false, false) => { 137 quote! { 138 #method_doc 139 async fn #name(&self, request: tonic::Request<#req_message>) 140 -> Result<tonic::Response<#res_message>, tonic::Status>; 141 } 142 } 143 (true, false) => { 144 quote! { 145 #method_doc 146 async fn #name(&self, request: tonic::Request<tonic::Streaming<#req_message>>) 147 -> Result<tonic::Response<#res_message>, tonic::Status>; 148 } 149 } 150 (false, true) => { 151 let stream = quote::format_ident!("{}Stream", method.identifier()); 152 let stream_doc = generate_doc_comment(&format!( 153 "Server streaming response type for the {} method.", 154 method.identifier() 155 )); 156 157 quote! { 158 #stream_doc 159 type #stream: Stream<Item = Result<#res_message, tonic::Status>> + Send + Sync + 'static; 160 161 #method_doc 162 async fn #name(&self, request: tonic::Request<#req_message>) 163 -> Result<tonic::Response<Self::#stream>, tonic::Status>; 164 } 165 } 166 (true, true) => { 167 let stream = quote::format_ident!("{}Stream", method.identifier()); 168 let stream_doc = generate_doc_comment(&format!( 169 "Server streaming response type for the {} method.", 170 method.identifier() 171 )); 172 173 quote! { 174 #stream_doc 175 type #stream: Stream<Item = Result<#res_message, tonic::Status>> + Send + Sync + 'static; 176 177 #method_doc 178 async fn #name(&self, request: tonic::Request<tonic::Streaming<#req_message>>) 179 -> Result<tonic::Response<Self::#stream>, tonic::Status>; 180 } 181 } 182 }; 183 184 stream.extend(method); 185 } 186 187 stream 188 } 189 190 #[cfg(feature = "transport")] 191 fn generate_transport( 192 server_service: &syn::Ident, 193 server_trait: &syn::Ident, 194 service_name: &str, 195 ) -> TokenStream { 196 let service_name = syn::LitStr::new(service_name, proc_macro2::Span::call_site()); 197 198 quote! { 199 impl<T: #server_trait> tonic::transport::NamedService for #server_service<T> { 200 const NAME: &'static str = #service_name; 201 } 202 } 203 } 204 205 #[cfg(not(feature = "transport"))] 206 fn generate_transport( 207 _server_service: &syn::Ident, 208 _server_trait: &syn::Ident, 209 _service_name: &str, 210 ) -> TokenStream { 211 TokenStream::new() 212 } 213 214 fn generate_methods<T: Service>(service: &T, proto_path: &str) -> TokenStream { 215 let mut stream = TokenStream::new(); 216 217 for method in service.methods() { 218 let path = format!( 219 "/{}{}{}/{}", 220 service.package(), 221 if service.package().is_empty() { 222 "" 223 } else { 224 "." 225 }, 226 service.identifier(), 227 method.identifier() 228 ); 229 let method_path = Lit::Str(LitStr::new(&path, Span::call_site())); 230 let ident = quote::format_ident!("{}", method.name()); 231 let server_trait = quote::format_ident!("{}", service.name()); 232 233 let method_stream = match (method.client_streaming(), method.server_streaming()) { 234 (false, false) => generate_unary(method, proto_path, ident, server_trait), 235 236 (false, true) => { 237 generate_server_streaming(method, proto_path, ident.clone(), server_trait) 238 } 239 (true, false) => { 240 generate_client_streaming(method, proto_path, ident.clone(), server_trait) 241 } 242 243 (true, true) => generate_streaming(method, proto_path, ident.clone(), server_trait), 244 }; 245 246 let method = quote! { 247 #method_path => { 248 #method_stream 249 } 250 }; 251 stream.extend(method); 252 } 253 254 stream 255 } 256 257 fn generate_unary<T: Method>( 258 method: &T, 259 proto_path: &str, 260 method_ident: Ident, 261 server_trait: Ident, 262 ) -> TokenStream { 263 let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap(); 264 265 let service_ident = quote::format_ident!("{}Svc", method.identifier()); 266 267 let (request, response) = method.request_response_name(proto_path); 268 269 quote! { 270 #[allow(non_camel_case_types)] 271 struct #service_ident<T: #server_trait >(pub Arc<T>); 272 273 impl<T: #server_trait> tonic::server::UnaryService<#request> for #service_ident<T> { 274 type Response = #response; 275 type Future = BoxFuture<tonic::Response<Self::Response>, tonic::Status>; 276 277 fn call(&mut self, request: tonic::Request<#request>) -> Self::Future { 278 let inner = self.0.clone(); 279 let fut = async move { 280 (*inner).#method_ident(request).await 281 }; 282 Box::pin(fut) 283 } 284 } 285 286 let inner = self.inner.clone(); 287 let fut = async move { 288 let interceptor = inner.1.clone(); 289 let inner = inner.0; 290 let method = #service_ident(inner); 291 let codec = #codec_name::default(); 292 293 let mut grpc = if let Some(interceptor) = interceptor { 294 tonic::server::Grpc::with_interceptor(codec, interceptor) 295 } else { 296 tonic::server::Grpc::new(codec) 297 }; 298 299 let res = grpc.unary(method, req).await; 300 Ok(res) 301 }; 302 303 Box::pin(fut) 304 } 305 } 306 307 fn generate_server_streaming<T: Method>( 308 method: &T, 309 proto_path: &str, 310 method_ident: Ident, 311 server_trait: Ident, 312 ) -> TokenStream { 313 let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap(); 314 315 let service_ident = quote::format_ident!("{}Svc", method.identifier()); 316 317 let (request, response) = method.request_response_name(proto_path); 318 319 let response_stream = quote::format_ident!("{}Stream", method.identifier()); 320 321 quote! { 322 #[allow(non_camel_case_types)] 323 struct #service_ident<T: #server_trait >(pub Arc<T>); 324 325 impl<T: #server_trait> tonic::server::ServerStreamingService<#request> for #service_ident<T> { 326 type Response = #response; 327 type ResponseStream = T::#response_stream; 328 type Future = BoxFuture<tonic::Response<Self::ResponseStream>, tonic::Status>; 329 330 fn call(&mut self, request: tonic::Request<#request>) -> Self::Future { 331 let inner = self.0.clone(); 332 let fut = async move { 333 (*inner).#method_ident(request).await 334 335 }; 336 Box::pin(fut) 337 } 338 } 339 340 let inner = self.inner.clone(); 341 let fut = async move { 342 let interceptor = inner.1; 343 let inner = inner.0; 344 let method = #service_ident(inner); 345 let codec = #codec_name::default(); 346 347 let mut grpc = if let Some(interceptor) = interceptor { 348 tonic::server::Grpc::with_interceptor(codec, interceptor) 349 } else { 350 tonic::server::Grpc::new(codec) 351 }; 352 353 let res = grpc.server_streaming(method, req).await; 354 Ok(res) 355 }; 356 357 Box::pin(fut) 358 } 359 } 360 361 fn generate_client_streaming<T: Method>( 362 method: &T, 363 proto_path: &str, 364 method_ident: Ident, 365 server_trait: Ident, 366 ) -> TokenStream { 367 let service_ident = quote::format_ident!("{}Svc", method.identifier()); 368 369 let (request, response) = method.request_response_name(proto_path); 370 let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap(); 371 372 quote! { 373 #[allow(non_camel_case_types)] 374 struct #service_ident<T: #server_trait >(pub Arc<T>); 375 376 impl<T: #server_trait> tonic::server::ClientStreamingService<#request> for #service_ident<T> 377 { 378 type Response = #response; 379 type Future = BoxFuture<tonic::Response<Self::Response>, tonic::Status>; 380 381 fn call(&mut self, request: tonic::Request<tonic::Streaming<#request>>) -> Self::Future { 382 let inner = self.0.clone(); 383 let fut = async move { 384 (*inner).#method_ident(request).await 385 386 }; 387 Box::pin(fut) 388 } 389 } 390 391 let inner = self.inner.clone(); 392 let fut = async move { 393 let interceptor = inner.1; 394 let inner = inner.0; 395 let method = #service_ident(inner); 396 let codec = #codec_name::default(); 397 398 let mut grpc = if let Some(interceptor) = interceptor { 399 tonic::server::Grpc::with_interceptor(codec, interceptor) 400 } else { 401 tonic::server::Grpc::new(codec) 402 }; 403 404 let res = grpc.client_streaming(method, req).await; 405 Ok(res) 406 }; 407 408 Box::pin(fut) 409 } 410 } 411 412 fn generate_streaming<T: Method>( 413 method: &T, 414 proto_path: &str, 415 method_ident: Ident, 416 server_trait: Ident, 417 ) -> TokenStream { 418 let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap(); 419 420 let service_ident = quote::format_ident!("{}Svc", method.identifier()); 421 422 let (request, response) = method.request_response_name(proto_path); 423 424 let response_stream = quote::format_ident!("{}Stream", method.identifier()); 425 426 quote! { 427 #[allow(non_camel_case_types)] 428 struct #service_ident<T: #server_trait>(pub Arc<T>); 429 430 impl<T: #server_trait> tonic::server::StreamingService<#request> for #service_ident<T> 431 { 432 type Response = #response; 433 type ResponseStream = T::#response_stream; 434 type Future = BoxFuture<tonic::Response<Self::ResponseStream>, tonic::Status>; 435 436 fn call(&mut self, request: tonic::Request<tonic::Streaming<#request>>) -> Self::Future { 437 let inner = self.0.clone(); 438 let fut = async move { 439 (*inner).#method_ident(request).await 440 }; 441 Box::pin(fut) 442 } 443 } 444 445 let inner = self.inner.clone(); 446 let fut = async move { 447 let interceptor = inner.1; 448 let inner = inner.0; 449 let method = #service_ident(inner); 450 let codec = #codec_name::default(); 451 452 let mut grpc = if let Some(interceptor) = interceptor { 453 tonic::server::Grpc::with_interceptor(codec, interceptor) 454 } else { 455 tonic::server::Grpc::new(codec) 456 }; 457 458 let res = grpc.streaming(method, req).await; 459 Ok(res) 460 }; 461 462 Box::pin(fut) 463 } 464 } 465