From bd8da93aaf819c8a835a709e82d6a7f9c914241b Mon Sep 17 00:00:00 2001
From: James Nugent <james@jen20.com>
Date: Mon, 11 May 2020 11:55:50 -0500
Subject: [PATCH] Add `include_file_descriptor_set` to `prost-build`

This commit adds support for writing an encoded version of the
appropriate FileDescriptorSet to a file named `file_descriptor_set.bin`
in the output directory. This can be used with the `include_bytes` macro
by consuming applications or libraries.

The rationale for this pull request is to support the gRPC Server
Reflection Protocol, which operates in terms of FileDescriptorProto, in
Tonic. FileDescriptorProto instances are contained inside a
FileDescriptorSet.
---
 README.md              | 10 +++++++++-
 prost-build/src/lib.rs | 19 ++++++++++++++++++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 9c641c2e7..036ae2d30 100644
--- a/README.md
+++ b/README.md
@@ -22,7 +22,9 @@ Compared to other Protocol Buffers implementations, `prost`
 * Respects the Protobuf `package` specifier when organizing generated code
   into Rust modules.
 * Preserves unknown enum values during deserialization.
-* Does not include support for runtime reflection or message descriptors.
+* Does not include support for runtime reflection or message descriptors, but
+  allows encoded `FileDescriptorSet` messages to be saved alongside generated
+  code for use of consumers.
 
 ## Using `prost` in a Cargo Project
 
@@ -59,6 +61,12 @@ package foo.bar;
 
 All Rust types generated from the file will be in the `foo::bar` module.
 
+If the `include_file_descriptor_set` option is set on the `prost_build::Config`
+for the invocation of `prost_build`, an additional file named `file_descriptor_set.bin`
+will be written to the output directory. This can be used in conjunction with the
+`include_bytes` macro, and decoded using the `FileDescriptorSet` type from the
+`prost-types` crate in applications or libraries using Prost.
+
 ### Messages
 
 Given a simple message declaration:
diff --git a/prost-build/src/lib.rs b/prost-build/src/lib.rs
index 521b2d2a2..3e81be3ae 100644
--- a/prost-build/src/lib.rs
+++ b/prost-build/src/lib.rs
@@ -118,7 +118,7 @@ use std::env;
 use std::ffi::{OsStr, OsString};
 use std::fmt;
 use std::fs;
-use std::io::{Error, ErrorKind, Result};
+use std::io::{Error, ErrorKind, Result, Write};
 use std::path::{Path, PathBuf};
 use std::process::Command;
 
@@ -183,6 +183,7 @@ pub trait ServiceGenerator {
 ///
 /// This configuration builder can be used to set non-default code generation options.
 pub struct Config {
+    include_file_descriptor_set: bool,
     service_generator: Option<Box<dyn ServiceGenerator>>,
     btree_map: Vec<String>,
     bytes: Vec<String>,
@@ -522,6 +523,15 @@ impl Config {
         self
     }
 
+    /// Configures the code generator to output the encoded bytes of the `FileDescriptorSet`
+    /// for this `protoc` invocation to a file named `file_descriptor_set.bin` in the configured
+    /// output directory. This can be used in conjunction with the `include_bytes!` macro and
+    /// the types in the `prost-types` crate for implementing some reflection capabilities.
+    pub fn include_file_descriptor_set(&mut self) -> &mut Self {
+        self.include_file_descriptor_set = true;
+        self
+    }
+
     /// Configures the code generator to not strip the enum name from variant names.
     ///
     /// Protobuf enum definitions commonly include the enum name as a prefix of every variant name.
@@ -644,6 +654,12 @@ impl Config {
             )
         })?;
 
+        if self.include_file_descriptor_set {
+            let filename = target.join("file_descriptor_set.bin");
+            let mut file = std::fs::File::create(filename)?;
+            file.write_all(&buf)?;
+        }
+
         let modules = self.generate(descriptor_set.file)?;
         for (module, content) in modules {
             let mut filename = module.join(".");
@@ -710,6 +726,7 @@ impl Config {
 impl default::Default for Config {
     fn default() -> Config {
         Config {
+            include_file_descriptor_set: false,
             service_generator: None,
             btree_map: Vec::new(),
             bytes: Vec::new(),