Skip to content

Commit d6703c0

Browse files
mhilsmhils
committed
fix review nits
1 parent 9a9cca8 commit d6703c0

7 files changed

Lines changed: 123 additions & 102 deletions

File tree

‎mitmproxy-contentviews/src/protobuf/existing_proto_definitions.rs‎

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,49 @@
1+
use crate::protobuf::raw_to_proto::new_empty_descriptor;
12
use crate::Metadata;
23
use anyhow::Context;
34
use protobuf::reflect::{FileDescriptor, MessageDescriptor};
45
use protobuf_parse::Parser;
56
use std::path::Path;
67

8+
pub(super) struct DescriptorWithDeps {
9+
pub descriptor: MessageDescriptor,
10+
pub dependencies: Vec<FileDescriptor>,
11+
}
12+
13+
impl Default for DescriptorWithDeps {
14+
fn default() -> Self {
15+
Self {
16+
descriptor: new_empty_descriptor(None, "Unknown"),
17+
dependencies: vec![],
18+
}
19+
}
20+
}
21+
722
pub(super) fn find_best_match(
823
metadata: &dyn Metadata,
9-
) -> anyhow::Result<Option<(MessageDescriptor, Vec<FileDescriptor>)>> {
24+
) -> anyhow::Result<Option<DescriptorWithDeps>> {
1025
// Parse existing protobuf definitions if available
11-
let file_descriptors = metadata
26+
let Some(file_descriptors) = metadata
1227
.protobuf_definitions()
1328
.map(parse_file_descriptor_set)
1429
.transpose()
15-
.context("failed to parse proto file(s)")?;
16-
let Some(file_descriptors) = file_descriptors else {
30+
.context("failed to parse proto file(s)")?
31+
else {
1732
return Ok(None);
1833
};
1934

2035
// Find MessageDescriptor for the RPC.
2136
let rpc_info = RpcInfo::from_metadata(metadata);
22-
let Some(message) = find_best_message(&file_descriptors, rpc_info, metadata.is_http_request())
37+
let Some(descriptor) =
38+
find_best_message(&file_descriptors, rpc_info, metadata.is_http_request())
2339
else {
2440
return Ok(None);
2541
};
2642

27-
Ok(Some((message, file_descriptors)))
43+
Ok(Some(DescriptorWithDeps {
44+
descriptor,
45+
dependencies: file_descriptors,
46+
}))
2847
}
2948

3049
fn find_best_message(

‎mitmproxy-contentviews/src/protobuf/proto_to_yaml.rs‎

Lines changed: 11 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use crate::protobuf::view_protobuf::tags;
12
/// Parsed protobuf message => YAML value
23
use protobuf::descriptor::field_descriptor_proto::Type;
34
use protobuf::descriptor::field_descriptor_proto::Type::{
@@ -7,10 +8,9 @@ use protobuf::reflect::{ReflectFieldRef, ReflectValueRef};
78
use protobuf::MessageDyn;
89
use serde_yaml::value::TaggedValue;
910
use serde_yaml::{Mapping, Number, Value};
10-
use std::fmt::Write;
1111
use std::ops::Deref;
1212

13-
pub(crate) fn message_to_yaml(message: &dyn MessageDyn) -> Value {
13+
pub(super) fn message_to_yaml(message: &dyn MessageDyn) -> Value {
1414
let mut ret = Mapping::new();
1515

1616
for field in message.descriptor_dyn().fields() {
@@ -28,7 +28,7 @@ pub(crate) fn message_to_yaml(message: &dyn MessageDyn) -> Value {
2828
let value = match field.get_reflect(message) {
2929
ReflectFieldRef::Optional(x) => {
3030
if let Some(x) = x.value() {
31-
primitive_type_to_yaml(x, field_type)
31+
value_to_yaml(x, field_type)
3232
} else {
3333
continue;
3434
}
@@ -39,7 +39,7 @@ pub(crate) fn message_to_yaml(message: &dyn MessageDyn) -> Value {
3939
}
4040
Value::Sequence(
4141
x.into_iter()
42-
.map(|x| primitive_type_to_yaml(x, field_type))
42+
.map(|x| value_to_yaml(x, field_type))
4343
.collect(),
4444
)
4545
}
@@ -49,12 +49,7 @@ pub(crate) fn message_to_yaml(message: &dyn MessageDyn) -> Value {
4949
}
5050
Value::Mapping(
5151
x.into_iter()
52-
.map(|(k, v)| {
53-
(
54-
primitive_type_to_yaml(k, field_type),
55-
primitive_type_to_yaml(v, field_type),
56-
)
57-
})
52+
.map(|(k, v)| (value_to_yaml(k, field_type), value_to_yaml(v, field_type)))
5853
.collect(),
5954
)
6055
}
@@ -64,7 +59,7 @@ pub(crate) fn message_to_yaml(message: &dyn MessageDyn) -> Value {
6459
Value::Mapping(ret)
6560
}
6661

67-
fn primitive_type_to_yaml(x: ReflectValueRef, field_type: Type) -> Value {
62+
fn value_to_yaml(x: ReflectValueRef, field_type: Type) -> Value {
6863
match x {
6964
ReflectValueRef::U32(x) => tag_number(Value::Number(Number::from(x)), field_type),
7065
ReflectValueRef::U64(x) => tag_number(Value::Number(Number::from(x)), field_type),
@@ -75,8 +70,8 @@ fn primitive_type_to_yaml(x: ReflectValueRef, field_type: Type) -> Value {
7570
ReflectValueRef::Bool(x) => Value::from(x),
7671
ReflectValueRef::String(x) => Value::from(x),
7772
ReflectValueRef::Bytes(x) => Value::Tagged(Box::new(TaggedValue {
78-
tag: crate::protobuf::view_protobuf::tags::BINARY.clone(),
79-
value: Value::String(bytes_to_hex_string(x)),
73+
tag: tags::BINARY.clone(),
74+
value: Value::String(data_encoding::HEXLOWER.encode(x)),
8075
})),
8176
ReflectValueRef::Enum(descriptor, i) => descriptor
8277
.value_by_number(i)
@@ -89,26 +84,17 @@ fn primitive_type_to_yaml(x: ReflectValueRef, field_type: Type) -> Value {
8984
fn tag_number(value: Value, field_type: Type) -> Value {
9085
match field_type {
9186
TYPE_UINT64 => Value::Tagged(Box::new(TaggedValue {
92-
tag: crate::protobuf::view_protobuf::tags::VARINT.clone(),
87+
tag: tags::VARINT.clone(),
9388
value,
9489
})),
9590
TYPE_FIXED64 => Value::Tagged(Box::new(TaggedValue {
96-
tag: crate::protobuf::view_protobuf::tags::FIXED64.clone(),
91+
tag: tags::FIXED64.clone(),
9792
value,
9893
})),
9994
TYPE_FIXED32 => Value::Tagged(Box::new(TaggedValue {
100-
tag: crate::protobuf::view_protobuf::tags::FIXED32.clone(),
95+
tag: tags::FIXED32.clone(),
10196
value,
10297
})),
10398
_ => value,
10499
}
105100
}
106-
107-
// Convert length-delimited protobuf data to a hex string
108-
fn bytes_to_hex_string(bytes: &[u8]) -> String {
109-
let mut result = String::with_capacity(bytes.len() * 2);
110-
for b in bytes {
111-
let _ = write!(result, "{:02x}", b);
112-
}
113-
result
114-
}

‎mitmproxy-contentviews/src/protobuf/raw_to_proto.rs‎

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use crate::protobuf::existing_proto_definitions::DescriptorWithDeps;
12
use anyhow::Context;
23
use protobuf::descriptor::field_descriptor_proto::Label::LABEL_REPEATED;
34
use protobuf::descriptor::field_descriptor_proto::Type;
@@ -17,21 +18,20 @@ enum GuessedFieldType {
1718
}
1819

1920
/// Create a "merged" MessageDescriptor. Mostly a wrapper around `create_descriptor_proto`.
20-
pub(crate) fn merge_proto_and_descriptor(
21+
pub(super) fn merge_proto_and_descriptor(
2122
data: &[u8],
22-
existing: &MessageDescriptor,
23-
dependencies: &[FileDescriptor],
23+
desc: &DescriptorWithDeps,
2424
) -> anyhow::Result<MessageDescriptor> {
25-
let new_proto = create_descriptor_proto(data, existing)?;
25+
let new_proto = create_descriptor_proto(data, &desc.descriptor)?;
2626

2727
let descriptor = {
28-
let mut file_descriptor_proto = existing.file_descriptor_proto().clone();
28+
let mut file_descriptor_proto = desc.descriptor.file_descriptor_proto().clone();
2929

3030
let message_idx = file_descriptor_proto
3131
.message_type
3232
.iter()
3333
.enumerate()
34-
.filter_map(|(i, d)| (d.name() == existing.name_to_package()).then_some(i))
34+
.filter_map(|(i, d)| (d.name() == desc.descriptor.name_to_package()).then_some(i))
3535
.next()
3636
.context("failed to find existing message descriptor index")?;
3737
file_descriptor_proto.message_type[message_idx] = new_proto;
@@ -45,17 +45,25 @@ pub(crate) fn merge_proto_and_descriptor(
4545
.collect::<Vec<_>>();
4646
*/
4747

48-
FileDescriptor::new_dynamic(file_descriptor_proto, dependencies)
48+
FileDescriptor::new_dynamic(file_descriptor_proto, &desc.dependencies)
4949
.context("failed to create new dynamic file descriptor")?
50-
.message_by_package_relative_name(existing.name_to_package())
51-
.with_context(|| format!("did not find {} in descriptor", existing.name_to_package()))?
50+
.message_by_package_relative_name(desc.descriptor.name_to_package())
51+
.with_context(|| {
52+
format!(
53+
"did not find {} in descriptor",
54+
desc.descriptor.name_to_package()
55+
)
56+
})?
5257
};
5358

5459
Ok(descriptor)
5560
}
5661

5762
/// Create a new (empty) MessageDescriptor for the given package and name.
5863
pub(super) fn new_empty_descriptor(package: Option<String>, name: &str) -> MessageDescriptor {
64+
// Create nested descriptor protos. For example, if the name is OuterMessage.InnerMessage,
65+
// we create a descriptor for InnerMessage and set it as a nested type of OuterMessage.
66+
// This is a bit of a hack, but the best way to get type_name right.
5967
let mut parts = name.rsplit(".");
6068
let mut head = {
6169
let mut descriptor = DescriptorProto::new();

‎mitmproxy-contentviews/src/protobuf/reencode.rs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,18 @@
11
use crate::protobuf::view_protobuf::tags;
2-
use crate::Metadata;
32
use anyhow::{bail, Context};
43
use protobuf::descriptor::field_descriptor_proto::Type;
54
use protobuf::descriptor::field_descriptor_proto::Type::{TYPE_FIXED32, TYPE_FIXED64};
6-
use protobuf::reflect::{FieldDescriptor, RuntimeFieldType, RuntimeType};
5+
use protobuf::reflect::{FieldDescriptor, MessageDescriptor, RuntimeFieldType, RuntimeType};
76
use protobuf::well_known_types::empty::Empty;
87
use protobuf::{MessageDyn, MessageFull, UnknownValue};
98
use serde_yaml::{Number, Value};
109
use std::num::ParseIntError;
1110
use std::str::FromStr;
1211

13-
pub(crate) fn reencode_yaml(value: Value, _metadata: &dyn Metadata) -> anyhow::Result<Vec<u8>> {
14-
let descriptor = Empty::descriptor();
12+
pub(super) fn reencode_yaml(
13+
value: Value,
14+
descriptor: &MessageDescriptor,
15+
) -> anyhow::Result<Vec<u8>> {
1516
let message = descriptor.new_instance();
1617
merge_yaml_into_message(value, message)
1718
}

‎mitmproxy-contentviews/src/protobuf/view_grpc.rs‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
use super::{existing_proto_definitions, reencode};
2-
use crate::protobuf::raw_to_proto::new_empty_descriptor;
32
use crate::{Metadata, Prettify, Protobuf, Reencode};
43
use anyhow::{bail, Context, Result};
54
use flate2::read::{DeflateDecoder, GzDecoder};
@@ -22,8 +21,7 @@ impl Prettify for GRPC {
2221
fn prettify(&self, mut data: &[u8], metadata: &dyn Metadata) -> Result<String> {
2322
let mut protos = vec![];
2423

25-
let (descriptor, dependencies) = existing_proto_definitions::find_best_match(metadata)?
26-
.unwrap_or_else(|| (new_empty_descriptor(None, "Unknown"), vec![]));
24+
let descriptor = existing_proto_definitions::find_best_match(metadata)?.unwrap_or_default();
2725

2826
while !data.is_empty() {
2927
let compressed = match data[0] {
@@ -59,7 +57,7 @@ impl Prettify for GRPC {
5957
} else {
6058
proto
6159
};
62-
protos.push(Protobuf.prettify_with_descriptor(proto, &descriptor, &dependencies)?);
60+
protos.push(Protobuf.prettify_with_descriptor(proto, &descriptor)?);
6361
data = &data[5 + len..];
6462
}
6563

@@ -78,10 +76,13 @@ impl Prettify for GRPC {
7876

7977
impl Reencode for GRPC {
8078
fn reencode(&self, data: &str, metadata: &dyn Metadata) -> Result<Vec<u8>> {
79+
let descriptor = existing_proto_definitions::find_best_match(metadata)?
80+
.unwrap_or_default()
81+
.descriptor;
8182
let mut ret = vec![];
8283
for document in serde_yaml::Deserializer::from_str(data) {
8384
let value = Value::deserialize(document).context("Invalid YAML")?;
84-
let proto = reencode::reencode_yaml(value, metadata)?;
85+
let proto = reencode::reencode_yaml(value, &descriptor)?;
8586
ret.push(0); // uncompressed
8687
ret.extend(u32::to_be_bytes(proto.len() as u32));
8788
ret.extend(proto);

0 commit comments

Comments
 (0)