main
use std::fs;
use std::path::{Path, PathBuf};
const JSON_PACKAGES: [&str; 5] = [
"v1",
"common.v1",
"admin.v1",
"cluster.v1",
"cluster.internal.v1",
];
const SERDE_HELPERS: &str = "crate::serde_helpers";
const DESCRIPTOR_TMP: &str = ".descriptor.bin";
fn main() -> Result<(), Box<dyn std::error::Error>> {
let manifest_dir = Path::new(env!("CARGO_MANIFEST_DIR"));
let out_dir = PathBuf::from(std::env::var("OUT_DIR")?);
build_protos(manifest_dir, &out_dir)
}
fn build_protos(manifest_dir: &Path, out_dir: &Path) -> Result<(), Box<dyn std::error::Error>> {
let proto_dir = PathBuf::from("proto");
let proto_root = manifest_dir.join(&proto_dir);
let mut files = Vec::new();
collect_protos(&proto_root, &mut files)?;
files.sort();
if files.is_empty() {
return Err(format!("no .proto found under {}", proto_root.display()).into());
}
let relative: Vec<PathBuf> = files
.iter()
.map(|file| proto_dir.join(file.strip_prefix(&proto_root).expect("scanned under root")))
.collect();
println!("cargo:rerun-if-changed={}", proto_root.display());
for file in &files {
println!("cargo:rerun-if-changed={}", file.display());
}
let mut builder = tonic_prost_build::configure()
.build_client(true)
.build_server(true)
.file_descriptor_set_path(DESCRIPTOR_TMP)
.enum_attribute(".", "#[allow(clippy::large_enum_variant)]")
.include_file("_include.rs");
for package in JSON_PACKAGES {
let prefix = format!(".{package}");
builder = builder
.type_attribute(&prefix, "#[derive(serde::Serialize, serde::Deserialize)]")
.message_attribute(&prefix, "#[serde(default)]")
.enum_attribute(&prefix, "#[serde(rename_all = \"SCREAMING_SNAKE_CASE\")]");
}
let bytes_helper = format!(
"#[serde(with = \"{SERDE_HELPERS}::bytes\")] #[schema(value_type = String)] #[doc = \"Hex encoded bytes (two hex digits per byte).\"]"
);
let bytes_list_helper = format!(
"#[serde(with = \"{SERDE_HELPERS}::bytes_list\")] #[schema(value_type = Vec<String>)]"
);
let timestamp_helper = format!(
"#[serde(with = \"{SERDE_HELPERS}::timestamp_opt\")] #[schema(value_type = Option<String>, default = json!(null))]"
);
for file in &files {
let text = fs::read_to_string(file)?;
for decl in scan_proto(&text) {
builder = match decl {
Decl::Schema(path) => builder.type_attribute(path, "#[derive(utoipa::ToSchema)]"),
Decl::Oneof(holder) => builder
.field_attribute(holder, "#[schema(value_type = Option<serde_json::Value>)]"),
Decl::Field(path, FieldKind::Bytes) => builder.field_attribute(path, &bytes_helper),
Decl::Field(path, FieldKind::BytesList) => {
builder.field_attribute(path, &bytes_list_helper)
}
Decl::Field(path, FieldKind::Timestamp) => {
builder.field_attribute(path, ×tamp_helper)
}
};
}
}
builder.compile_protos(&relative, std::slice::from_ref(&proto_dir))?;
fs::copy(DESCRIPTOR_TMP, out_dir.join("descriptor.bin"))?;
let _ = fs::remove_file(DESCRIPTOR_TMP);
Ok(())
}
fn collect_protos(dir: &Path, out: &mut Vec<PathBuf>) -> std::io::Result<()> {
for entry in fs::read_dir(dir)? {
let path = entry?.path();
if path.is_dir() {
collect_protos(&path, out)?;
} else if path.extension().is_some_and(|ext| ext == "proto") {
out.push(path);
}
}
Ok(())
}
enum Decl {
Schema(String),
Oneof(String),
Field(String, FieldKind),
}
enum FieldKind {
Bytes,
BytesList,
Timestamp,
}
enum Block {
Message(String),
Other,
}
fn scan_proto(text: &str) -> Vec<Decl> {
let text = strip_comments(text);
let Some(package) = package_of(&text) else {
return Vec::new();
};
if !JSON_PACKAGES.contains(&package.as_str()) {
return Vec::new();
}
let mut decls = Vec::new();
let mut stack: Vec<Block> = Vec::new();
for raw in text.lines() {
let line = raw.trim();
if line.is_empty() {
continue;
}
if let Some(head) = line.strip_suffix('{') {
let head = head.trim();
let mut tokens = head.split_whitespace();
let keyword = tokens.next().unwrap_or_default();
let name = tokens.next().unwrap_or_default().to_string();
match keyword {
"message" => {
decls.push(Decl::Schema(path_of(&package, &stack, &name)));
stack.push(Block::Message(name));
}
"enum" => {
decls.push(Decl::Schema(path_of(&package, &stack, &name)));
stack.push(Block::Other);
}
"oneof" => {
decls.push(Decl::Oneof(path_of(&package, &stack, &name)));
stack.push(Block::Other);
}
_ => stack.push(Block::Other),
}
continue;
}
if let Some(head) = line.strip_suffix("{}") {
let head = head.trim();
let mut tokens = head.split_whitespace();
let keyword = tokens.next().unwrap_or_default();
let name = tokens.next().unwrap_or_default();
if matches!(keyword, "message" | "enum") && !name.is_empty() {
decls.push(Decl::Schema(path_of(&package, &stack, name)));
}
continue;
}
if line == "}" {
stack.pop();
continue;
}
if !matches!(stack.last(), Some(Block::Message(_))) {
continue;
}
let Some((ty, name, repeated)) = parse_field(line) else {
continue;
};
let kind = match (ty.as_str(), repeated) {
("google.protobuf.Timestamp", _) => FieldKind::Timestamp,
("bytes", true) => FieldKind::BytesList,
("bytes", false) => FieldKind::Bytes,
_ => continue,
};
decls.push(Decl::Field(path_of(&package, &stack, &name), kind));
}
decls
}
fn package_of(text: &str) -> Option<String> {
text.lines()
.find_map(|line| line.trim().strip_prefix("package "))
.map(|rest| rest.trim().trim_end_matches(';').trim().to_string())
}
fn path_of(package: &str, stack: &[Block], name: &str) -> String {
let mut path = format!(".{package}");
for block in stack {
if let Block::Message(outer) = block {
path.push('.');
path.push_str(outer);
}
}
path.push('.');
path.push_str(name);
path
}
fn parse_field(line: &str) -> Option<(String, String, bool)> {
let body = line.strip_suffix(';')?;
let tokens: Vec<&str> = body.split_whitespace().collect();
let eq = tokens.iter().position(|token| *token == "=")?;
if eq < 2 {
return None;
}
let name = tokens[eq - 1];
let ty = tokens[eq - 2];
if ty.contains('<') || !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
return None;
}
let repeated = eq >= 3 && tokens[eq - 3] == "repeated";
Some((ty.to_string(), name.to_string(), repeated))
}
fn strip_comments(text: &str) -> String {
let mut out = String::with_capacity(text.len());
let mut chars = text.chars().peekable();
let mut line_comment = false;
let mut block_comment = false;
while let Some(current) = chars.next() {
if line_comment {
if current == '\n' {
line_comment = false;
out.push('\n');
}
continue;
}
if block_comment {
if current == '*' && chars.peek() == Some(&'/') {
chars.next();
block_comment = false;
}
continue;
}
if current == '/' && chars.peek() == Some(&'/') {
chars.next();
line_comment = true;
continue;
}
if current == '/' && chars.peek() == Some(&'*') {
chars.next();
block_comment = true;
continue;
}
out.push(current);
}
out
}