Updates reflection server implementation

This change updates jen20's work in hyperium#340 to work with the file descriptor set changes that landed in prost 0.7. * Updates branch with new changes from master * Updates server to use tokio 1.0 symantics * `tonic_build::Builder::include_file_descriptor_set` -> `file_descriptor_set_path`, which matchesprost. * The reflection server now returns UNIMPLEMENTED rather than NOT_FOUND when unsupported methods are called.
sgg · Jan 16, 2021 · 8811f0c · 8811f0c
1 parent 59514ec
commit 8811f0c
Show file tree

Hide file tree

Showing 10 changed files with 76 additions and 43 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -19,8 +19,3 @@ members = [
     "tests/extern_path/my_application",
     "tests/integration_tests"
 ]
-
-[replace]
-'prost:0.6.1' = { git = "https://github.com/jen20/prost", branch = "jen20/file-descriptor-set" }
-'prost-build:0.6.1' = { git = "https://github.com/jen20/prost", branch = "jen20/file-descriptor-set" }
-'prost-types:0.6.1' = { git = "https://github.com/jen20/prost", branch = "jen20/file-descriptor-set" }
diff --git a/examples/README.md b/examples/README.md
@@ -94,6 +94,14 @@ $ cargo run --bin tls-server
 $ cargo run --bin health-server
 ```
 
+## Reflection Service
+
+### Server
+
+```bash
+$ cargo run --bin reflection-server
+```
+
 ## Tower Middleware
 
 ### Server

diff --git a/examples/build.rs b/examples/build.rs
@@ -1,10 +1,20 @@
+use std::env;
+use std::path::PathBuf;
+
 fn main() {
+    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
+
     tonic_build::configure()
         .type_attribute("routeguide.Point", "#[derive(Hash)]")
         .compile(&["proto/routeguide/route_guide.proto"], &["proto"])
         .unwrap();
 
-    tonic_build::compile_protos("proto/helloworld/helloworld.proto").unwrap();
+    tonic_build::configure()
+        .file_descriptor_set_path(out_dir.join("helloworld_descriptor.bin"))
+        .compile(&["proto/helloworld/helloworld.proto"], &["proto"])
+        .unwrap();
+
     tonic_build::compile_protos("proto/echo/echo.proto").unwrap();
+
     tonic_build::compile_protos("proto/google/pubsub/pubsub.proto").unwrap();
 }
diff --git a/examples/src/reflection/server.rs b/examples/src/reflection/server.rs
@@ -4,7 +4,8 @@ use tonic::{Request, Response, Status};
 mod proto {
     tonic::include_proto!("helloworld");
 
-    pub(crate) const FILE_DESCRIPTOR_SET: &'static [u8] = tonic::include_file_descriptor_set!();
+    pub(crate) const HELLO_WORLD_DESCRIPTOR_SET: &'static [u8] =
+        include_bytes!(concat!(env!("OUT_DIR"), "/helloworld_descriptor.bin"));
 }
 
 #[derive(Default)]
@@ -27,16 +28,16 @@ impl proto::greeter_server::Greeter for MyGreeter {
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let service = tonic_reflection::server::Builder::configure()
-        .register_encoded_file_descriptor_set(proto::FILE_DESCRIPTOR_SET)
+    let reflection_service = tonic_reflection::server::Builder::configure()
+        .register_encoded_file_descriptor_set(proto::HELLO_WORLD_DESCRIPTOR_SET)
         .build()
         .unwrap();
 
     let addr = "[::1]:50052".parse().unwrap();
     let greeter = MyGreeter::default();
 
     Server::builder()
-        .add_service(service)
+        .add_service(reflection_service)
         .add_service(proto::greeter_server::GreeterServer::new(greeter))
         .serve(addr)
         .await?;

diff --git a/tonic-build/src/prost.rs b/tonic-build/src/prost.rs
@@ -12,7 +12,7 @@ pub fn configure() -> Builder {
     Builder {
         build_client: true,
         build_server: true,
-        include_file_descriptor_set: false,
+        file_descriptor_set_path: None,
         out_dir: None,
         extern_path: Vec::new(),
         field_attributes: Vec::new(),
@@ -181,7 +181,7 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
 pub struct Builder {
     pub(crate) build_client: bool,
     pub(crate) build_server: bool,
-    pub(crate) include_file_descriptor_set: bool,
+    pub(crate) file_descriptor_set_path: Option<PathBuf>,
     pub(crate) extern_path: Vec<(String, String)>,
     pub(crate) field_attributes: Vec<(String, String)>,
     pub(crate) type_attributes: Vec<(String, String)>,
@@ -207,8 +207,8 @@ impl Builder {
 
     /// Generate a file containing the encoded `prost_types::FileDescriptorSet` for protocol buffers
     /// modules. This is required for implementing gRPC Server Reflection.
-    pub fn include_file_descriptor_set(mut self, include: bool) -> Self {
-        self.include_file_descriptor_set = include;
+    pub fn file_descriptor_set_path(mut self, path: impl AsRef<Path>) -> Self {
+        self.file_descriptor_set_path = Some(path.as_ref().to_path_buf());
         self
     }
 
@@ -296,9 +296,11 @@ impl Builder {
         let format = self.format;
 
         config.out_dir(out_dir.clone());
-        if self.include_file_descriptor_set {
-            config.include_file_descriptor_set();
+
+        if let Some(path) = self.file_descriptor_set_path.as_ref() {
+            config.file_descriptor_set_path(path);
         }
+
         for (proto_path, rust_path) in self.extern_path.iter() {
             config.extern_path(proto_path, rust_path);
         }

diff --git a/tonic-health/Cargo.toml b/tonic-health/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "tonic-health"
 version = "0.3.0"
-authors = ["James Nugent <[email protected]>"]
+authors = ["James Nugent <[email protected]>", "Samani G. Gikandi <[email protected]>"]
 edition = "2018"
 license = "MIT"
 repository = "https://github.com/hyperium/tonic"

diff --git a/tonic-reflection/Cargo.toml b/tonic-reflection/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "tonic-reflection"
-version = "0.1.0"
+version = "0.4.0"
 authors = ["James Nugent <[email protected]>"]
 edition = "2018"
 license = "MIT"
@@ -14,11 +14,12 @@ categories = ["network-programming", "asynchronous"]
 keywords = ["rpc", "grpc", "async", "reflection"]
 
 [dependencies]
-bytes = "0.5"
-prost = "0.6"
-prost-types = "0.6"
-tokio = { version = "0.2", features = ["sync", "stream"] }
-tonic = { version = "0.2", path = "../tonic", features = ["codegen", "prost"] }
+bytes = "1.0"
+prost = "0.7"
+prost-types = "0.7"
+tokio = { version = "1.0", features = ["sync"] }
+tonic = { version = "0.4", path = "../tonic", features = ["codegen", "prost"] }
+tokio-stream = "0.1"
 
 [build-dependencies]
-tonic-build = { version = "0.2", path = "../tonic-build" }
+tonic-build = { version = "0.4", path = "../tonic-build" }
diff --git a/tonic-reflection/build.rs b/tonic-reflection/build.rs
@@ -1,6 +1,12 @@
+use std::env;
+use std::path::PathBuf;
+
 fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let reflection_descriptors =
+        PathBuf::from(env::var("OUT_DIR").unwrap()).join("reflection_descriptor.bin");
+
     tonic_build::configure()
-        .include_file_descriptor_set(true)
+        .file_descriptor_set_path(&reflection_descriptors)
         .build_server(true)
         .build_client(false)
         .format(true)

diff --git a/tonic-reflection/src/lib.rs b/tonic-reflection/src/lib.rs
@@ -18,7 +18,8 @@ mod proto {
     #![allow(unreachable_pub)]
     tonic::include_proto!("grpc.reflection.v1alpha");
 
-    pub(crate) const FILE_DESCRIPTOR_SET: &'static [u8] = tonic::include_file_descriptor_set!();
+    pub(crate) const REFLECTION_DESCRIPTOR_SET: &'static [u8] =
+        include_bytes!(concat!(env!("OUT_DIR"), "/reflection_descriptor.bin"));
 }
 
 /// Implementation of the server component of gRPC Server Reflection.

diff --git a/tonic-reflection/src/server.rs b/tonic-reflection/src/server.rs
@@ -1,22 +1,26 @@
-use crate::proto::server_reflection_request::MessageRequest;
-use crate::proto::server_reflection_response::MessageResponse;
-use crate::proto::server_reflection_server::{ServerReflection, ServerReflectionServer};
-use crate::proto::{
-    FileDescriptorResponse, ListServiceResponse, ServerReflectionRequest, ServerReflectionResponse,
-    ServiceResponse,
-};
+use std::collections::HashMap;
+use std::fmt::{Display, Formatter};
+use std::pin::Pin;
+use std::sync::Arc;
+
 use prost::{DecodeError, Message};
 use prost_types::{
     DescriptorProto, EnumDescriptorProto, FieldDescriptorProto, FileDescriptorProto,
     FileDescriptorSet,
 };
-use std::collections::HashMap;
-use std::fmt::{Display, Formatter};
-use std::sync::Arc;
-use tokio::stream::StreamExt;
 use tokio::sync::mpsc;
+use tokio_stream::wrappers::ReceiverStream;
+use tokio_stream::{Stream, StreamExt};
 use tonic::{Request, Response, Status, Streaming};
 
+use crate::proto::server_reflection_request::MessageRequest;
+use crate::proto::server_reflection_response::MessageResponse;
+use crate::proto::server_reflection_server::{ServerReflection, ServerReflectionServer};
+use crate::proto::{
+    FileDescriptorResponse, ListServiceResponse, ServerReflectionRequest, ServerReflectionResponse,
+    ServiceResponse,
+};
+
 /// Represents an error in the construction of a gRPC Reflection Service.
 #[derive(Debug)]
 pub enum Error {
@@ -78,6 +82,8 @@ impl<'b> Builder<'b> {
 
     /// Registers a byte slice containing an encoded `prost_types::FileDescriptorSet` with
     /// the gRPC Reflection Service builder.
+    ///
+    /// This can be called multiple times to append new descriptors to the serivce.
     pub fn register_encoded_file_descriptor_set(
         mut self,
         encoded_file_descriptor_set: &'b [u8],
@@ -97,7 +103,8 @@ impl<'b> Builder<'b> {
     /// Build a gRPC Reflection Service to be served via Tonic.
     pub fn build(mut self) -> Result<ServerReflectionServer<impl ServerReflection>, Error> {
         if self.include_reflection_service {
-            self = self.register_encoded_file_descriptor_set(crate::proto::FILE_DESCRIPTOR_SET);
+            self =
+                self.register_encoded_file_descriptor_set(crate::proto::REFLECTION_DESCRIPTOR_SET);
         }
 
         for encoded in &self.encoded_file_descriptor_sets {
@@ -304,14 +311,16 @@ struct ReflectionService {
 
 #[tonic::async_trait]
 impl ServerReflection for ReflectionService {
-    type ServerReflectionInfoStream = mpsc::Receiver<Result<ServerReflectionResponse, Status>>;
+    type ServerReflectionInfoStream = Pin<
+        Box<dyn Stream<Item = Result<ServerReflectionResponse, Status>> + Send + Sync + 'static>,
+    >;
 
     async fn server_reflection_info(
         &self,
         req: Request<Streaming<ServerReflectionRequest>>,
     ) -> Result<Response<Self::ServerReflectionInfoStream>, Status> {
         let mut req_rx = req.into_inner();
-        let (mut resp_tx, resp_rx) = mpsc::channel::<Result<ServerReflectionResponse, Status>>(1);
+        let (resp_tx, resp_rx) = mpsc::channel::<Result<ServerReflectionResponse, Status>>(1);
 
         let state = self.state.clone();
 
@@ -330,10 +339,10 @@ impl ServerReflection for ReflectionService {
                         MessageRequest::FileByFilename(s) => state.file_by_filename(&s),
                         MessageRequest::FileContainingSymbol(s) => state.symbol_by_name(&s),
                         MessageRequest::FileContainingExtension(_) => {
-                            Err(Status::not_found("extensions are not supported"))
+                            Err(Status::unimplemented("extensions are not supported"))
                         }
                         MessageRequest::AllExtensionNumbersOfType(_) => {
-                            Err(Status::not_found("extensions are not supported"))
+                            Err(Status::unimplemented("extensions are not supported"))
                         }
                         MessageRequest::ListServices(_) => Ok(state.list_services()),
                     },
@@ -356,6 +365,6 @@ impl ServerReflection for ReflectionService {
             }
         });
 
-        Ok(Response::new(resp_rx))
+        Ok(Response::new(Box::pin(ReceiverStream::new(resp_rx))))
     }
 }