From ff6a690cec9daca33984cabea66f9d370ac63462 Mon Sep 17 00:00:00 2001
From: Tom Dyas <tdyas@toolchain.com>
Date: Fri, 22 Oct 2021 08:44:27 -0400
Subject: [PATCH] fix(build): Correctly convert `Empty` to `()` (#734)

---
 tests/wellknown-compiled/Cargo.toml       |  1 +
 tests/wellknown-compiled/build.rs         |  9 ++++-
 tests/wellknown-compiled/proto/test.proto | 11 ++++++
 tests/wellknown-compiled/src/lib.rs       | 17 +++++++--
 tonic-build/src/prost.rs                  | 46 ++++++++++-------------
 5 files changed, 53 insertions(+), 31 deletions(-)
 create mode 100644 tests/wellknown-compiled/proto/test.proto

diff --git a/tests/wellknown-compiled/Cargo.toml b/tests/wellknown-compiled/Cargo.toml
index 5a5c837..075369b 100644
--- a/tests/wellknown-compiled/Cargo.toml
+++ b/tests/wellknown-compiled/Cargo.toml
@@ -16,4 +16,5 @@ tonic = { path = "../../tonic" }
 prost = "0.9"
 
 [build-dependencies]
+prost-build = "0.9"
 tonic-build = { path = "../../tonic-build" }
diff --git a/tests/wellknown-compiled/build.rs b/tests/wellknown-compiled/build.rs
index 01301f5..cb07fca 100644
--- a/tests/wellknown-compiled/build.rs
+++ b/tests/wellknown-compiled/build.rs
@@ -1,6 +1,13 @@
 fn main() {
+    let mut config = prost_build::Config::new();
+    config.extern_path(".google.protobuf.Empty", "()");
+
     tonic_build::configure()
         .compile_well_known_types(true)
-        .compile(&["proto/google.proto"], &["proto"])
+        .compile_with_config(
+            config,
+            &["proto/google.proto", "proto/test.proto"],
+            &["proto"],
+        )
         .unwrap();
 }
diff --git a/tests/wellknown-compiled/proto/test.proto b/tests/wellknown-compiled/proto/test.proto
new file mode 100644
index 0000000..2e4d3b4
--- /dev/null
+++ b/tests/wellknown-compiled/proto/test.proto
@@ -0,0 +1,11 @@
+syntax = "proto3";
+package test;
+
+import "google/protobuf/empty.proto";
+
+message Input {
+}
+
+service Service {
+  rpc Call(Input) returns (google.protobuf.Empty);
+}
diff --git a/tests/wellknown-compiled/src/lib.rs b/tests/wellknown-compiled/src/lib.rs
index 2d287e7..60ab404 100644
--- a/tests/wellknown-compiled/src/lib.rs
+++ b/tests/wellknown-compiled/src/lib.rs
@@ -1,9 +1,18 @@
-pub mod google {
-    pub mod protobuf {
-        tonic::include_proto!("google.protobuf");
+pub mod gen {
+    pub mod google {
+        pub mod protobuf {
+            tonic::include_proto!("google.protobuf");
+        }
+    }
+
+    pub mod test {
+        tonic::include_proto!("test");
     }
 }
 
 pub fn grok() {
-    let _empty = crate::google::protobuf::Empty {};
+    let _any = crate::gen::google::protobuf::Any {
+        type_url: "foo".to_owned(),
+        value: Vec::new(),
+    };
 }
diff --git a/tonic-build/src/prost.rs b/tonic-build/src/prost.rs
index 108cf80..3b0f514 100644
--- a/tonic-build/src/prost.rs
+++ b/tonic-build/src/prost.rs
@@ -49,6 +49,9 @@ pub fn compile_protos(proto: impl AsRef<Path>) -> io::Result<()> {
 
 const PROST_CODEC_PATH: &str = "tonic::codec::ProstCodec";
 
+/// Non-path Rust types allowed for request/response types.
+const NON_PATH_TYPE_ALLOWLIST: &[&str] = &["()"];
+
 impl crate::Service for Service {
     const CODEC_PATH: &'static str = PROST_CODEC_PATH;
 
@@ -105,34 +108,25 @@ impl crate::Method for Method {
         proto_path: &str,
         compile_well_known_types: bool,
     ) -> (TokenStream, TokenStream) {
-        let request = if (is_google_type(&self.input_proto_type) && !compile_well_known_types)
-            || self.input_type.starts_with("::")
-        {
-            self.input_type.parse::<TokenStream>().unwrap()
-        } else if self.input_type.starts_with("crate::") {
-            syn::parse_str::<syn::Path>(&self.input_type)
-                .unwrap()
-                .to_token_stream()
-        } else {
-            syn::parse_str::<syn::Path>(&format!("{}::{}", proto_path, self.input_type))
-                .unwrap()
-                .to_token_stream()
-        };
-
-        let response = if (is_google_type(&self.output_proto_type) && !compile_well_known_types)
-            || self.output_type.starts_with("::")
-        {
-            self.output_type.parse::<TokenStream>().unwrap()
-        } else if self.output_type.starts_with("crate::") {
-            syn::parse_str::<syn::Path>(&self.output_type)
-                .unwrap()
-                .to_token_stream()
-        } else {
-            syn::parse_str::<syn::Path>(&format!("{}::{}", proto_path, self.output_type))
-                .unwrap()
-                .to_token_stream()
+        let convert_type = |proto_type: &str, rust_type: &str| -> TokenStream {
+            if (is_google_type(proto_type) && !compile_well_known_types)
+                || rust_type.starts_with("::")
+                || NON_PATH_TYPE_ALLOWLIST.iter().any(|ty| *ty == rust_type)
+            {
+                rust_type.parse::<TokenStream>().unwrap()
+            } else if rust_type.starts_with("crate::") {
+                syn::parse_str::<syn::Path>(rust_type)
+                    .unwrap()
+                    .to_token_stream()
+            } else {
+                syn::parse_str::<syn::Path>(&format!("{}::{}", proto_path, rust_type))
+                    .unwrap()
+                    .to_token_stream()
+            }
         };
 
+        let request = convert_type(&self.input_proto_type, &self.input_type);
+        let response = convert_type(&self.output_proto_type, &self.output_type);
         (request, response)
     }
 }