build.rs

// Copyright (c) 2020 Ant Group
//
// SPDX-License-Identifier: Apache-2.0
//

use std::fs::File;
use std::io::{BufRead, BufReader, Read, Write};
use std::path::Path;
use std::process::exit;
use ttrpc_codegen::{Codegen, Customize, ProtobufCustomize};

fn replace_text_in_file(file_name: &str, from: &str, to: &str) -> Result<(), std::io::Error> {
    let mut src = File::open(file_name)?;
    let mut contents = String::new();
    src.read_to_string(&mut contents).unwrap();
    drop(src);

    let new_contents = contents.replace(from, to);

    let mut dst = File::create(&file_name)?;
    dst.write_all(new_contents.as_bytes())?;

    Ok(())
}

fn use_serde(protos: &[&str], out_dir: &Path) -> Result<(), std::io::Error> {
    protos
        .iter()
        .try_for_each(|f: &&str| -> Result<(), std::io::Error> {
            let out_file = Path::new(f)
                .file_name()
                .and_then(|s| s.to_str())
                .ok_or(format!("failed to get proto file name for {:?}", f))
                .map(|s| {
                    let t = s.replace(".proto", ".rs");
                    out_dir.join(t)
                })
                .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?
                .to_str()
                .ok_or(format!("cannot convert {:?} path to string", f))
                .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?
                .to_string();

            replace_text_in_file(
                &out_file,
                "derive(Serialize, Deserialize)",
                "derive(serde::Serialize, serde::Deserialize)",
            )
        })
}

fn handle_file(autogen_comment: &str, rust_filename: &str) -> Result<(), std::io::Error> {
    let mut new_contents = Vec::new();

    let file = File::open(rust_filename)?;

    let reader = BufReader::new(file);

    // Guard the code since it is only needed for the agent-ctl tool,
    // not the agent itself.
    let serde_default_code = r#"#[cfg_attr(feature = "with-serde", serde(default))]"#;

    for line in reader.lines() {
        let line = line?;

        new_contents.push(line.clone());

        let pattern = "//! Generated file from";

        if line.starts_with(&pattern) {
            new_contents.push(autogen_comment.into());
        }

        let struct_pattern = "pub struct ";

        // Although we've requested serde support via `Customize`, to
        // allow the `kata-agent-ctl` tool to partially deserialise structures
        // specified in JSON, we need this bit of additional magic.
        if line.starts_with(&struct_pattern) {
            new_contents.insert(new_contents.len() - 1, serde_default_code.trim().into());
        }
    }

    let data = new_contents.join("\n");

    let mut dst = File::create(&rust_filename)?;

    dst.write_all(data.as_bytes())?;

    Ok(())
}

fn real_main() -> Result<(), std::io::Error> {
    let autogen_comment = format!("\n//! Generated by {:?} ({:?})", file!(), module_path!());

    let protos = vec![
        "protos/agent.proto",
        "protos/csi.proto",
        "protos/google/protobuf/empty.proto",
        "protos/health.proto",
        "protos/oci.proto",
        "protos/types.proto",
        "protos/image.proto",
    ];

    // Tell Cargo that if the .proto files changed, to rerun this build script.
    protos
        .iter()
        .for_each(|p| println!("cargo:rerun-if-changed={}", &p));

    let ttrpc_options = Customize {
        async_server: true,
        ..Default::default()
    };

    let protobuf_options = ProtobufCustomize {
        serde_derive: Some(true),
        ..Default::default()
    };

    let out_dir = Path::new("src");

    Codegen::new()
        .out_dir(out_dir)
        .inputs(&protos)
        .include("protos")
        .customize(ttrpc_options)
        .rust_protobuf()
        .rust_protobuf_customize(protobuf_options)
        .run()?;

    for file in protos.iter() {
        let proto_filename = Path::new(file).file_name().unwrap();

        let generated_file = proto_filename
            .to_str()
            .ok_or("failed")
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?
            .replace(".proto", ".rs");

        let out_file = out_dir.join(generated_file);

        let out_file_str = out_file
            .to_str()
            .ok_or("failed")
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;

        handle_file(&autogen_comment, out_file_str)?;
    }

    // There is a message named 'Box' in oci.proto
    // so there is a struct named 'Box', we should replace Box<Self> to ::std::boxed::Box<Self>
    // to avoid the conflict.
    replace_text_in_file(
        "src/oci.rs",
        "self: Box<Self>",
        "self: ::std::boxed::Box<Self>",
    )?;

    use_serde(&protos, out_dir)?;

    Ok(())
}

fn main() {
    if let Err(e) = real_main() {
        eprintln!("ERROR: {}", e);
        exit(1);
    }
}