add gRPC contentview

Author: mhils

mitmproxy-rs/mitmproxy_rs/contentviews.pyi

@@ -14,7 +14,8 @@ class InteractiveContentview(Contentview):
 hex_dump: Contentview
 hex_stream: InteractiveContentview
 msgpack: InteractiveContentview
-protobuf: Contentview
+protobuf: InteractiveContentview
+grpc: InteractiveContentview
 
 __all__ = [
     "Contentview",
@@ -23,4 +24,5 @@ __all__ = [
     "hex_stream",
     "msgpack",
     "protobuf",
+    "grpc",
 ]

mitmproxy-rs/src/contentview.rs

@@ -1,17 +1,21 @@
 use mitmproxy::contentviews::{Metadata, Prettify, Reencode};
 use pyo3::{exceptions::PyValueError, prelude::*};
 
-struct PythonMetadata(PyObject);
+pub struct PythonMetadata<'py>(Bound<'py, PyAny>);
 
-impl Metadata for PythonMetadata {
+impl Metadata for PythonMetadata<'_> {
     fn content_type(&self) -> Option<String> {
-        Python::with_gil(|py| {
-            self.0
-                .getattr(py, "content_type")
-                .ok()?
-                .extract::<String>(py)
-                .ok()
-        })
+        self.0
+            .getattr("content_type")
+            .ok()?
+            .extract::<String>()
+            .ok()
+    }
+}
+
+impl<'py> FromPyObject<'py> for PythonMetadata<'py> {
+    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
+        Ok(PythonMetadata(ob.clone()))
     }
 }
 
@@ -36,17 +40,14 @@ impl Contentview {
     }
 
     /// Pretty-print an (encoded) message.
-    pub fn prettify(&self, data: Vec<u8>, metadata: PyObject) -> PyResult<String> {
-        let metadata = PythonMetadata(metadata);
-
+    pub fn prettify(&self, data: Vec<u8>, metadata: PythonMetadata) -> PyResult<String> {
         self.0
             .prettify(&data, &metadata)
             .map_err(|e| PyValueError::new_err(e.to_string()))
     }
 
     /// Return the priority of this view for rendering data.
-    pub fn render_priority(&self, data: Vec<u8>, metadata: PyObject) -> PyResult<f64> {
-        let metadata = PythonMetadata(metadata);
+    pub fn render_priority(&self, data: Vec<u8>, metadata: PythonMetadata) -> PyResult<f64> {
         Ok(self.0.render_priority(&data, &metadata))
     }
 
@@ -85,9 +86,7 @@ impl InteractiveContentview {
 
 #[pymethods]
 impl InteractiveContentview {
-    pub fn reencode(&self, data: &str, metadata: PyObject) -> PyResult<Vec<u8>> {
-        let metadata = PythonMetadata(metadata);
-
+    pub fn reencode(&self, data: &str, metadata: PythonMetadata) -> PyResult<Vec<u8>> {
         self.0
             .reencode(data, &metadata)
             .map_err(|e| PyValueError::new_err(e.to_string()))

mitmproxy-rs/src/lib.rs

@@ -88,18 +88,19 @@ mod mitmproxy_rs {
     #[pymodule]
     mod contentviews {
         use super::*;
-        //#[pymodule_export]
-        //use crate::contentview::Contentview;
-        //#[pymodule_export]
-        //use crate::contentview::InteractiveContentview;
-        use mitmproxy::contentviews::{HexDump, HexStream, MsgPack, Protobuf};
+        #[pymodule_export]
+        use crate::contentview::Contentview;
+        #[pymodule_export]
+        use crate::contentview::InteractiveContentview;
+        use mitmproxy::contentviews::{HexDump, HexStream, MsgPack, Protobuf, GRPC};
 
         #[pymodule_init]
         fn init(m: &Bound<'_, PyModule>) -> PyResult<()> {
             m.add_contentview(&HexDump)?;
             m.add_interactive_contentview(&HexStream)?;
             m.add_interactive_contentview(&MsgPack)?;
-            m.add_contentview(&Protobuf)?;
+            m.add_interactive_contentview(&Protobuf)?;
+            m.add_interactive_contentview(&GRPC)?;
             Ok(())
         }
     }

src/contentviews/grpc.rs

@@ -0,0 +1,85 @@
+use crate::contentviews::{Metadata, Prettify, Protobuf, Reencode};
+use crate::syntax_highlight::Language;
+use anyhow::{bail, Context, Result};
+use serde::Deserialize;
+use serde_yaml::Value;
+
+pub struct GRPC;
+
+impl Prettify for GRPC {
+    fn name(&self) -> &'static str {
+        "gRPC"
+    }
+
+    fn syntax_highlight(&self) -> Language {
+        Language::Yaml
+    }
+
+    fn prettify(&self, mut data: &[u8], metadata: &dyn Metadata) -> Result<String> {
+        let mut protos = vec![];
+
+        while !data.is_empty() {
+            let compressed = match data[0] {
+                0 => false,
+                1 => true,
+                _ => bail!("invalid gRPC: first byte is not a boolean"),
+            };
+            let len = match data.get(1..5) {
+                Some(x) => u32::from_be_bytes(x.try_into()?) as usize,
+                _ => bail!("invalid gRPC: first byte is not a boolean"),
+            };
+            let Some(proto) = data.get(5..5 + len) else {
+                bail!("Invald gRPC: not enough data")
+            };
+            if compressed {
+                todo!();
+            }
+            protos.push(proto);
+            data = &data[5 + len..];
+        }
+
+        let prettified = protos
+            .into_iter()
+            .map(|proto| Protobuf.prettify(proto, metadata))
+            .collect::<Result<Vec<String>>>()?;
+        Ok(prettified.join("\n---\n\n"))
+    }
+
+    fn render_priority(&self, _data: &[u8], metadata: &dyn Metadata) -> f64 {
+        let Some(ct) = metadata.content_type() else {
+            return 0.0;
+        };
+        match ct.as_str() {
+            "application/grpc" => 2.0,
+            "application/grpc+proto" => 2.0,
+            "application/prpc" => 2.0,
+            _ => 0.0,
+        }
+    }
+}
+
+impl Reencode for GRPC {
+    fn reencode(&self, data: &str, metadata: &dyn Metadata) -> Result<Vec<u8>> {
+        let mut ret = vec![];
+        for document in serde_yaml::Deserializer::from_str(data) {
+            let value = Value::deserialize(document).context("Invalid YAML")?;
+            let proto = Protobuf::reencode_yaml(value, metadata)?;
+            ret.push(0); // compressed
+            ret.extend(u32::to_be_bytes(proto.len() as u32));
+            ret.extend(proto);
+        }
+        Ok(ret)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::contentviews::TestMetadata;
+
+    #[test]
+    fn test_grpc() {
+        let result = GRPC.prettify(b"foo", &TestMetadata::default()).unwrap();
+        assert_eq!(result, "666f6f");
+    }
+}

src/contentviews/mod.rs

@@ -1,3 +1,4 @@
+mod grpc;
 mod hex_dump;
 mod hex_stream;
 mod msgpack;
@@ -6,6 +7,7 @@ mod protobuf;
 use anyhow::Result;
 
 use crate::syntax_highlight;
+pub use grpc::GRPC;
 pub use hex_dump::HexDump;
 pub use hex_stream::HexStream;
 pub use msgpack::MsgPack;

src/contentviews/protobuf.rs

@@ -1,4 +1,5 @@
 use crate::contentviews::{Metadata, Prettify, Reencode};
+use crate::syntax_highlight::Language;
 use anyhow::{bail, Context, Result};
 use protobuf::descriptor::field_descriptor_proto::Label::LABEL_REPEATED;
 use protobuf::descriptor::field_descriptor_proto::Type;
@@ -54,10 +55,14 @@ impl Prettify for Protobuf {
         "Protocol Buffer"
     }
 
+    fn syntax_highlight(&self) -> Language {
+        Language::Yaml
+    }
+
     fn prettify(&self, data: &[u8], _metadata: &dyn Metadata) -> Result<String> {
         // Check if data is empty first
         if data.is_empty() {
-            bail!("Empty protobuf data");
+            return Ok("{}  # empty protobuf message".to_string());
         }
 
         let existing = Empty::descriptor();
@@ -79,13 +84,9 @@ impl Prettify for Protobuf {
 }
 
 impl Reencode for Protobuf {
-    fn reencode(&self, data: &str, _metadata: &dyn Metadata) -> Result<Vec<u8>> {
-        let descriptor = Empty::descriptor();
-        let message = descriptor.new_instance();
-
+    fn reencode(&self, data: &str, metadata: &dyn Metadata) -> Result<Vec<u8>> {
         let value: Value = serde_yaml::from_str(data).context("Invalid YAML")?;
-
-        Self::merge_yaml_into_message(value, message)
+        Self::reencode_yaml(value, metadata)
     }
 }
 
@@ -143,6 +144,12 @@ fn int_value(n: Number, field: Option<&FieldDescriptor>) -> UnknownValue {
 }
 
 impl Protobuf {
+    pub(super) fn reencode_yaml(value: Value, _metadata: &dyn Metadata) -> Result<Vec<u8>> {
+        let descriptor = Empty::descriptor();
+        let message = descriptor.new_instance();
+        Self::merge_yaml_into_message(value, message)
+    }
+
     fn merge_yaml_into_message(value: Value, mut message: Box<dyn MessageDyn>) -> Result<Vec<u8>> {
         let Value::Mapping(mapping) = value else {
             bail!("YAML is not a mapping");