protocheck-build 0.1.2

Build-time entrypoint for protocheck.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

#![cfg_attr(docsrs, feature(doc_auto_cfg))]

//! Build-time entrypoint for [`protocheck`](https://docs.rs/protocheck/0.1.0/protocheck/)

use std::{
  env,
  error::Error,
  io::Read,
  path::{Path, PathBuf},
};

use prost_build::Config;
use prost_reflect::{prost::Message, prost_types::FileDescriptorSet};

#[cfg(not(feature = "cel"))]
fn enable_cel() -> bool {
  false
}
#[cfg(feature = "cel")]
fn enable_cel() -> bool {
  true
}

pub fn compile_protos_with_validators(
  config: &mut Config,
  proto_files: &[impl AsRef<Path>],
  proto_include_paths: &[impl AsRef<Path>],
  packages: &[&str],
) -> Result<(), Box<dyn Error>> {
  let out_dir = env::var("OUT_DIR")
    .map(PathBuf::from)
    .unwrap_or_else(|_| env::temp_dir());

  let temp_descriptor_path = out_dir.join("temp_file_descriptor_set_for_protocheck.bin");
  {
    let mut temp_config = prost_build::Config::new();
    temp_config.file_descriptor_set_path(temp_descriptor_path.clone());
    temp_config.compile_protos(proto_files, proto_include_paths)?;
  }

  let mut fds_file = std::fs::File::open(&temp_descriptor_path)?;
  let mut fds_bytes = Vec::new();
  fds_file.read_to_end(&mut fds_bytes)?;
  let fds = FileDescriptorSet::decode(fds_bytes.as_slice())?;
  let pool = prost_reflect::DescriptorPool::from_file_descriptor_set(fds)?;

  for message_desc in pool.all_messages() {
    let message_name = message_desc.full_name();
    if packages.contains(&message_desc.package_name()) {
      let attribute_str = format!(
        r#"#[::protocheck::macros::protobuf_validate("{}")]"#,
        message_name
      );
      config.message_attribute(message_name, &attribute_str);

      if enable_cel() {
        config.message_attribute(
          message_name,
          "#[derive(::protocheck::macros::TryIntoCelValue)]",
        );
      }

      for oneof in message_desc.oneofs() {
        let oneof_name = oneof.full_name();
        config.type_attribute(
          oneof_name,
          format!(
            r#"#[::protocheck::macros::protobuf_validate_oneof("{}")]"#,
            oneof_name
          ),
        );

        config.type_attribute(oneof_name, r#"#[derive(::protocheck::macros::Oneof)]"#);

        if enable_cel() {
          config.type_attribute(
            oneof_name,
            r#"#[derive(::protocheck::macros::OneofTryIntoCelValue)]"#,
          );
        }

        for field in oneof.fields() {
          config.field_attribute(
            format!("{}.{}", oneof_name, field.name()),
            format!(r#"#[protocheck(proto_name = "{}")]"#, field.name()),
          );
        }
      }
    }
  }

  config.extern_path(".buf.validate", "::protocheck::types::protovalidate");
  config
    .extern_path(".google.protobuf", "::protocheck::types")
    .compile_well_known_types();

  std::fs::remove_file(&temp_descriptor_path)?;

  Ok(())
}