From 232d47446c7d7d63f5eaec78ec44df362b37d94d Mon Sep 17 00:00:00 2001
From: Noah Gold <nkgold@google.com>
Date: Thu, 10 Nov 2022 13:15:53 -0800
Subject: [PATCH] proto_build_tools: add crate for proto build tools.

In addition to DRYing out our proto code, we're switching to using
PathBufs to reduce the potential for cross platform errors.

BUG=b:256951877
TEST=builds

Change-Id: Ib7588de231afe67853c099e4f81683731b9439de
Reviewed-on: https://chromium-review.googlesource.com/c/crosvm/crosvm/+/4021590
Reviewed-by: Vikram Auradkar <auradkar@google.com>
---
 Cargo.lock                   |   7 +++
 Cargo.toml                   |   1 +
 proto_build_tools/Cargo.toml |  11 ++++
 proto_build_tools/src/lib.rs | 102 +++++++++++++++++++++++++++++++++++
 4 files changed, 121 insertions(+)
 create mode 100644 proto_build_tools/Cargo.toml
 create mode 100644 proto_build_tools/src/lib.rs

diff --git a/Cargo.lock b/Cargo.lock
index c497ac9783..0c43982280 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1550,6 +1550,13 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "proto_build_tools"
+version = "0.1.0"
+dependencies = [
+ "protoc-rust",
+]
+
 [[package]]
 name = "protobuf"
 version = "2.27.1"
diff --git a/Cargo.toml b/Cargo.toml
index 1a0db55259..a826ed6349 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -85,6 +85,7 @@ members = [
     "power_monitor",
     "prebuilts",
     "protos",
+    "proto_build_tools",
     "qcow_utils",
     "resources",
     "rutabaga_gfx",
diff --git a/proto_build_tools/Cargo.toml b/proto_build_tools/Cargo.toml
new file mode 100644
index 0000000000..a6e84697c4
--- /dev/null
+++ b/proto_build_tools/Cargo.toml
@@ -0,0 +1,11 @@
+[package]
+name = "proto_build_tools"
+authors = ["The Chromium OS Authors"]
+version = "0.1.0"
+edition = "2021"
+
+[features]
+
+[dependencies]
+protoc-rust = "2.22"
+
diff --git a/proto_build_tools/src/lib.rs b/proto_build_tools/src/lib.rs
new file mode 100644
index 0000000000..c55a3e2635
--- /dev/null
+++ b/proto_build_tools/src/lib.rs
@@ -0,0 +1,102 @@
+// Copyright 2022 The ChromiumOS Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+//! Contains utilities to build protos & make them available to Rust.
+
+use std::collections::HashSet;
+use std::fs;
+use std::fs::File;
+use std::io::Write;
+use std::path::PathBuf;
+
+use protoc_rust::Customize;
+
+/// Builds a set of Rust protos based on the provided proto files. The individual protos will be
+/// dumped into `out_dir` (will be created if needed), along with a file that wraps them
+/// `out_dir/generated.rs`. The wrapper file can then be included using a pattern like:
+/// ```ignore
+/// pub mod protos {
+///    // Suppose the `out_dir` supplied to `build_protos` was
+///    // format!("{}/my_crate_protos", env!("OUT_DIR"))
+///    include!(concat!(env!("OUT_DIR"), "/my_crate_protos/generated.rs"));
+/// }
+/// // Protos are available at protos::proto_file_name.
+/// ```
+pub fn build_protos(out_dir: &PathBuf, proto_paths: &[PathBuf]) {
+    build_protos_explicit(
+        out_dir,
+        proto_paths,
+        proto_paths,
+        to_includes(proto_paths).as_slice(),
+    )
+}
+
+/// Allows for more control than build_protos (useful when the proto build is more complex).
+pub fn build_protos_explicit(
+    out_dir: &PathBuf,
+    proto_paths: &[PathBuf],
+    rebuild_if_changed_paths: &[PathBuf],
+    includes: &[PathBuf],
+) {
+    for file in rebuild_if_changed_paths {
+        // Triggers rebuild if the file has newer mtime.
+        println!(
+            "cargo:rerun-if-changed={}",
+            file.to_str().expect("proto path must be UTF-8")
+        );
+    }
+    fs::create_dir_all(&out_dir).unwrap();
+    gen_protos(out_dir, proto_paths, includes);
+    create_gen_file(out_dir, proto_paths);
+}
+
+/// Given a list of proto files, extract the set of include directories needed to pass to the proto
+/// compiler.
+fn to_includes(proto_paths: &[PathBuf]) -> Vec<PathBuf> {
+    let mut include_paths = HashSet::new();
+
+    for proto in proto_paths {
+        include_paths.insert(
+            proto
+                .parent()
+                .expect("protos must be files in a directory")
+                .to_owned(),
+        );
+    }
+
+    include_paths.drain().collect::<Vec<PathBuf>>()
+}
+
+fn gen_protos(out_dir: &PathBuf, proto_paths: &[PathBuf], includes: &[PathBuf]) {
+    if let Err(e) = protoc_rust::Codegen::new()
+        .out_dir(out_dir)
+        .inputs(proto_paths)
+        .includes(includes)
+        .customize(Customize {
+            serde_derive: Some(true),
+            ..Default::default()
+        })
+        .run()
+    {
+        println!("failed to build Rust protos: {}", e);
+        println!("protos in failed build: {:?}", proto_paths);
+        println!("includes in failed build: {:?}", includes);
+    }
+}
+
+fn create_gen_file(out_dir: &PathBuf, proto_files: &[PathBuf]) {
+    let generated = PathBuf::from(&out_dir).join("generated.rs");
+    let out = File::create(generated).expect("Failed to create generated file.");
+
+    for proto_path in proto_files {
+        let file_stem = proto_path.file_stem().unwrap().to_str().unwrap();
+        let out_dir = out_dir
+            .to_str()
+            .expect("path must be UTF-8")
+            .replace('\\', "/");
+        writeln!(&out, "#[path = \"{}/{}.rs\"]", out_dir, file_stem)
+            .expect("failed to write to generated.");
+        writeln!(&out, "pub mod {};", file_stem).expect("failed to write to generated.");
+    }
+}