use protobuf::descriptor::field_descriptor_proto::{Label, Type};
use protobuf::descriptor::{
DescriptorProto, EnumDescriptorProto, FieldDescriptorProto, MethodDescriptorProto,
ServiceDescriptorProto,
};
use protobuf::rt::WireType;
use protobuf::Enum;
use protobuf_parse::Parser;
use std::collections::HashMap;
use tracing::warn;
pub struct DescriptorProtoHolder {
pub descriptor_proto: DescriptorProto,
pub field_index_map: HashMap<u32, u32>,
}
pub struct ServiceApiDescriptor {
pub service_descriptor: ServiceDescriptorProto,
pub message_descriptor_holder: HashMap<String, DescriptorProtoHolder>,
pub method_descriptor_map: HashMap<String, MethodDescriptorProto>,
pub enum_descriptor_map: HashMap<String, EnumDescriptorProto>,
}
impl ServiceApiDescriptor {
pub fn service_name(&self) -> &str {
self.service_descriptor.name()
}
pub fn methods(&self) -> &Vec<MethodDescriptorProto> {
&self.service_descriptor.method
}
pub fn method(&self, method_name: &str) -> Option<&MethodDescriptorProto> {
self.method_descriptor_map.get(method_name)
}
pub fn get_message_descriptor(&self, message_name: &str) -> Option<&DescriptorProto> {
let holder = self.message_descriptor_holder.get(message_name);
if holder.is_none() {
return None;
}
let holder = holder.unwrap();
Some(&holder.descriptor_proto)
}
fn put_enum_descriptor_in_map(
scope: &str,
enum_types: &Vec<EnumDescriptorProto>,
enum_descriptor_map: &mut HashMap<String, EnumDescriptorProto>,
) {
for enum_type in enum_types {
let enum_name = enum_type.name();
let key_name = format!("{}{}", scope, enum_name);
enum_descriptor_map.insert(key_name, enum_type.clone());
}
}
fn put_message_descriptor_in_map(
scope: &str,
message_types: &Vec<DescriptorProto>,
service_api_descriptor: &mut ServiceApiDescriptor,
) {
for msg_type in message_types {
let msg_name = msg_type.name();
let key_name = format!("{}{}", scope, msg_name);
let new_scope = format!("{}.", key_name);
let mut field_index_map: HashMap<u32, u32> = HashMap::new();
let mut cur_index: u32 = 0;
for field in &msg_type.field {
let number = field.number.unwrap();
let wire_type = get_wire_type(field);
let tag = make_tag(number as u32, wire_type);
field_index_map.insert(tag, cur_index);
cur_index += 1;
}
service_api_descriptor.message_descriptor_holder.insert(
key_name,
DescriptorProtoHolder {
descriptor_proto: msg_type.clone(),
field_index_map,
},
);
Self::put_message_descriptor_in_map(
&new_scope,
&msg_type.nested_type,
service_api_descriptor,
);
Self::put_enum_descriptor_in_map(
&new_scope,
&msg_type.enum_type,
&mut service_api_descriptor.enum_descriptor_map,
);
}
}
pub fn load(proto_file_path: &str) -> Result<ServiceApiDescriptor, String> {
Self::load_with_include_path(proto_file_path, ".", true)
}
pub fn load_with_include_path(
proto_file_path: &str,
include_path: &str,
check_service: bool,
) -> Result<ServiceApiDescriptor, String> {
let mut ins = ServiceApiDescriptor {
message_descriptor_holder: HashMap::new(),
method_descriptor_map: HashMap::new(),
service_descriptor: ServiceDescriptorProto::default(),
enum_descriptor_map: HashMap::new(),
};
let mut parser = Parser::new();
parser.pure();
parser.include(include_path);
parser.inputs(&[proto_file_path]);
let file_desc_set = parser.file_descriptor_set();
if file_desc_set.is_err() {
warn!("加载描述文件：{} 失败", proto_file_path);
return Err(format!("解析描述文件失败:{}", file_desc_set.err().unwrap()));
}
let file_desc_set = file_desc_set.unwrap();
let file_descriptor = file_desc_set.file.get(0);
if file_descriptor.is_none() {
warn!("加载描述文件：{} 失败", proto_file_path);
return Err(format!("解析描述文件失败, 描述文件不存在！"));
}
let file_descriptor = file_descriptor.unwrap();
if check_service {
let service = &file_descriptor.service;
if service.len() != 1 {
warn!(
"描述文件{}错误：描述文件只能由一个 Service 描述",
proto_file_path
);
return Err("描述文件错误：描述文件只能由一个 Service 描述".into());
}
let service: &ServiceDescriptorProto = service.get(0).unwrap();
ins.service_descriptor = service.clone();
for method in ins.service_descriptor.method.iter() {
ins.method_descriptor_map
.insert(String::from(method.name()), method.clone());
}
}
Self::put_message_descriptor_in_map(".", &file_descriptor.message_type, &mut ins);
Self::put_enum_descriptor_in_map(
".",
&file_descriptor.enum_type,
&mut ins.enum_descriptor_map,
);
Ok(ins)
}
}
pub fn get_wire_type(field: &FieldDescriptorProto) -> u32 {
let field_type = field.type_.unwrap().unwrap();
match field_type {
Type::TYPE_INT64
| Type::TYPE_UINT64
| Type::TYPE_INT32
| Type::TYPE_UINT32
| Type::TYPE_ENUM
| Type::TYPE_SINT32
| Type::TYPE_SINT64
| Type::TYPE_BOOL => WireType::Varint as u32,
Type::TYPE_FIXED64 | Type::TYPE_SFIXED64 | Type::TYPE_DOUBLE => WireType::Fixed64 as u32,
Type::TYPE_FIXED32 | Type::TYPE_SFIXED32 | Type::TYPE_FLOAT => WireType::Fixed32 as u32,
Type::TYPE_STRING | Type::TYPE_MESSAGE => WireType::LengthDelimited as u32,
_ => {
panic!("无法 识别 {} 类型", field.name.as_ref().unwrap())
}
}
}
pub fn make_tag(number: u32, wire_type: u32) -> u32 {
(number << 3) | wire_type
}
pub fn tag_to_number_and_wire_type(tag: u32) -> (u32, u32) {
let wire_type = tag & 0b111; let number = tag >> 3; (number, wire_type)
}
pub fn is_optional(field: &FieldDescriptorProto) -> bool {
return if field.proto3_optional.is_some() {
field.proto3_optional.unwrap()
} else {
false
};
}
pub fn is_integer(field: &FieldDescriptorProto) -> bool {
if let Some(type_value) = field.type_ {
let type_value = type_value.value();
type_value == Type::TYPE_INT64.value()
|| type_value == Type::TYPE_UINT64.value()
|| type_value == Type::TYPE_INT32.value()
|| type_value == Type::TYPE_FIXED64.value()
|| type_value == Type::TYPE_FIXED32.value()
|| type_value == Type::TYPE_UINT32.value()
|| type_value == Type::TYPE_ENUM.value()
|| type_value == Type::TYPE_SFIXED32.value()
|| type_value == Type::TYPE_SFIXED64.value()
|| type_value == Type::TYPE_SINT32.value()
|| type_value == Type::TYPE_SINT64.value()
} else {
false
}
}
pub fn is_float(field: &FieldDescriptorProto) -> bool {
if let Some(type_value) = field.type_ {
let type_value = type_value.value();
type_value == Type::TYPE_DOUBLE.value() || type_value == Type::TYPE_FLOAT.value()
} else {
false
}
}
pub fn is_string(field: &FieldDescriptorProto) -> bool {
if let Some(type_value) = field.type_ {
let type_value = type_value.value();
type_value == Type::TYPE_STRING.value()
} else {
false
}
}
pub fn is_object(field: &FieldDescriptorProto) -> bool {
if let Some(type_value) = field.type_ {
type_value.value() == Type::TYPE_MESSAGE.value()
} else {
false
}
}
pub fn is_message_or_enum(field: &FieldDescriptorProto) -> bool {
let field_type = field.type_.unwrap().unwrap();
match field_type {
Type::TYPE_MESSAGE | Type::TYPE_ENUM => true,
_ => false,
}
}
pub fn is_array(field: &FieldDescriptorProto) -> bool {
if let Some(label) = field.label {
label.value() == Label::LABEL_REPEATED.value()
} else {
false
}
}
pub fn is_map(field: &FieldDescriptorProto, api_descriptor: &ServiceApiDescriptor) -> bool {
if let Some(type_name) = &field.type_name {
let message_descriptor = api_descriptor.get_message_descriptor(type_name);
if message_descriptor.is_none() {
return false;
}
let message_descriptor = message_descriptor.unwrap();
return message_descriptor.options.map_entry();
}
false
}