From b02b4b238bfee96b886609396b957e2592477ecb Mon Sep 17 00:00:00 2001
From: James Steele <steele@users.noreply.github.com>
Date: Sat, 9 Nov 2019 15:45:22 +0000
Subject: [PATCH] fix(build): Prevent duplicated client/server generated code
 (#121)

* fix(build): Prevent duplicated client/server generated code

The tonic-build process collects up RPC services as they are provided by
prost, before writing them out as part of the finalization step for a
given protocol buffer package.

In the case of imported protocol buffer packages, there may be RPC
services included by import in addition to those in the top-level
package. Therefore it is necessary to make sure each set of
client/server services gathered by tonic-build is cleared after the
finalization process for a given protocol buffer package, otherwise they
will be incorrectly aggregated as the generation process proceeds
through the subsequent packages.

* Test case for duplicated client/server generated code

A simple test case that will fail to build without a fix to prevent
RPC services being duplicated into inappropriate modules (that related
to particular protocol buffer packages).

* Additional test case for included_service

Introduces an additional case that captures making sure services defined
before including a package with additional services doesn't
incidentially clear such precursor services from the including package.

* Fix unnecessary newline to keep `cargo fmt` happy
---
 Cargo.toml                                  |  1 +
 tests/included_service/Cargo.toml           | 16 ++++++++++++++++
 tests/included_service/build.rs             |  3 +++
 tests/included_service/proto/includee.proto | 11 +++++++++++
 tests/included_service/proto/includer.proto | 11 +++++++++++
 tests/included_service/src/lib.rs           |  8 ++++++++
 tonic-build/src/lib.rs                      |  4 ++++
 7 files changed, 54 insertions(+)
 create mode 100644 tests/included_service/Cargo.toml
 create mode 100644 tests/included_service/build.rs
 create mode 100644 tests/included_service/proto/includee.proto
 create mode 100644 tests/included_service/proto/includer.proto
 create mode 100644 tests/included_service/src/lib.rs

diff --git a/Cargo.toml b/Cargo.toml
index a0febc4..a34e834 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,6 +5,7 @@ members = [
     "tonic-examples",
     "tonic-interop",
 
+    "tests/included_service",
     "tests/same_name",
     "tests/wellknown",
 ]
diff --git a/tests/included_service/Cargo.toml b/tests/included_service/Cargo.toml
new file mode 100644
index 0000000..c70a3eb
--- /dev/null
+++ b/tests/included_service/Cargo.toml
@@ -0,0 +1,16 @@
+[package]
+name = "included_service"
+version = "0.1.0"
+authors = ["Lucio Franco <luciofranco14@gmail.com>"]
+edition = "2018"
+publish = false
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+tonic = { path = "../../tonic" }
+bytes = "0.4"
+prost = "0.5"
+
+[build-dependencies]
+tonic-build = { path = "../../tonic-build" }
diff --git a/tests/included_service/build.rs b/tests/included_service/build.rs
new file mode 100644
index 0000000..1c29065
--- /dev/null
+++ b/tests/included_service/build.rs
@@ -0,0 +1,3 @@
+fn main() {
+    tonic_build::compile_protos("proto/includer.proto").unwrap();
+}
diff --git a/tests/included_service/proto/includee.proto b/tests/included_service/proto/includee.proto
new file mode 100644
index 0000000..c2aebc0
--- /dev/null
+++ b/tests/included_service/proto/includee.proto
@@ -0,0 +1,11 @@
+syntax = "proto3";
+
+package includee;
+
+service Included {
+  rpc SomeMethod(SomeRequest) returns (SomeResponse) {}
+}
+
+message SomeRequest {}
+
+message SomeResponse {}
diff --git a/tests/included_service/proto/includer.proto b/tests/included_service/proto/includer.proto
new file mode 100644
index 0000000..8f283d6
--- /dev/null
+++ b/tests/included_service/proto/includer.proto
@@ -0,0 +1,11 @@
+syntax = "proto3";
+
+package includer;
+
+message TopMessage {}
+
+service TopService {
+  rpc TopMethod(TopMessage) returns (TopMessage) {}
+}
+
+import "includee.proto";
diff --git a/tests/included_service/src/lib.rs b/tests/included_service/src/lib.rs
new file mode 100644
index 0000000..d8e6f10
--- /dev/null
+++ b/tests/included_service/src/lib.rs
@@ -0,0 +1,8 @@
+pub mod pb {
+    tonic::include_proto!("includer");
+}
+
+// Ensure that an RPC service, defined before including a file that defines
+// another service in a different protocol buffer package, is not incorrectly
+// cleared from the context of its package.
+type _Test = dyn pb::server::TopService;
diff --git a/tonic-build/src/lib.rs b/tonic-build/src/lib.rs
index b6e28d8..6b9cac5 100644
--- a/tonic-build/src/lib.rs
+++ b/tonic-build/src/lib.rs
@@ -262,6 +262,8 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
 
             let code = format!("{}", client_service);
             buf.push_str(&code);
+
+            self.clients = TokenStream::default();
         }
 
         if self.builder.build_server && !self.servers.is_empty() {
@@ -279,6 +281,8 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
 
             let code = format!("{}", server_service);
             buf.push_str(&code);
+
+            self.servers = TokenStream::default();
         }
     }
 }