fix content_type handling, add render_priority implementations

Author: mhils

mitmproxy-contentviews/src/grpc.rs

@@ -46,13 +46,10 @@ impl Prettify for GRPC {
     }
 
     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,
+        match metadata.content_type() {
+            Some("application/grpc") => 1.0,
+            Some("application/grpc+proto") => 1.0,
+            Some("application/prpc") => 1.0,
             _ => 0.0,
         }
     }

mitmproxy-contentviews/src/hex_dump.rs

@@ -26,7 +26,7 @@ impl Prettify for HexDump {
 
     fn render_priority(&self, data: &[u8], _metadata: &dyn Metadata) -> f64 {
         if is_binary(data) {
-            0.95
+            0.5
         } else {
             0.0
         }

mitmproxy-contentviews/src/hex_stream.rs

@@ -28,7 +28,7 @@ impl Prettify for HexStream {
 
     fn render_priority(&self, data: &[u8], _metadata: &dyn Metadata) -> f64 {
         if is_binary(data) {
-            0.95
+            0.4
         } else {
             0.0
         }

mitmproxy-contentviews/src/lib.rs

@@ -4,17 +4,17 @@ mod hex_stream;
 mod msgpack;
 mod protobuf;
 
-use anyhow::Result;
-
 pub use grpc::GRPC;
 pub use hex_dump::HexDump;
 pub use hex_stream::HexStream;
-use mitmproxy_highlight::Language;
 pub use msgpack::MsgPack;
 pub use protobuf::Protobuf;
 
+use anyhow::Result;
+use mitmproxy_highlight::Language;
+
 pub trait Metadata {
-    fn content_type(&self) -> Option<String>;
+    fn content_type(&self) -> Option<&str>;
 }
 
 pub trait Prettify: Send + Sync {
@@ -30,7 +30,8 @@ pub trait Prettify: Send + Sync {
 
     fn prettify(&self, data: &[u8], metadata: &dyn Metadata) -> Result<String>;
 
-    fn render_priority(&self, _data: &[u8], _metadata: &dyn Metadata) -> f64 {
+    #[allow(unused_variables)]
+    fn render_priority(&self, data: &[u8], metadata: &dyn Metadata) -> f64 {
         0.0
     }
 }
@@ -45,7 +46,7 @@ pub struct TestMetadata {
 }
 
 impl Metadata for TestMetadata {
-    fn content_type(&self) -> Option<String> {
-        self.content_type.clone()
+    fn content_type(&self) -> Option<&str> {
+        self.content_type.as_deref()
     }
 }

mitmproxy-contentviews/src/msgpack.rs

@@ -22,6 +22,14 @@ impl Prettify for MsgPack {
         // Convert the Value to prettified YAML
         serde_yaml::to_string(&value).context("Failed to convert to YAML")
     }
+
+    fn render_priority(&self, _data: &[u8], metadata: &dyn Metadata) -> f64 {
+        match metadata.content_type() {
+            Some("application/msgpack") => 1.0,
+            Some("application/x-msgpack") => 1.0,
+            _ => 0.0,
+        }
+    }
 }
 
 impl Reencode for MsgPack {

mitmproxy-contentviews/src/protobuf.rs

@@ -76,6 +76,14 @@ impl Prettify for Protobuf {
         let yaml_str = serde_yaml::to_string(&yaml_value).context("Failed to convert to YAML")?;
         yaml_to_pretty::apply_replacements(&yaml_str)
     }
+
+    fn render_priority(&self, _data: &[u8], metadata: &dyn Metadata) -> f64 {
+        match metadata.content_type() {
+            Some("application/x-protobuf") => 1.0,
+            Some("application/x-protobuffer") => 1.0,
+            _ => 0.0,
+        }
+    }
 }
 
 impl Reencode for Protobuf {

mitmproxy-rs/src/contentview.rs

@@ -1,21 +1,39 @@
 use mitmproxy_contentviews::{Metadata, Prettify, Reencode};
 use pyo3::{exceptions::PyValueError, prelude::*};
+use std::cell::OnceCell;
 
-pub struct PythonMetadata<'py>(Bound<'py, PyAny>);
+pub struct PythonMetadata<'py> {
+    inner: Bound<'py, PyAny>,
+    content_type: OnceCell<Option<String>>,
+}
+
+impl<'py> PythonMetadata<'py> {
+    pub fn new(inner: Bound<'py, PyAny>) -> Self {
+        PythonMetadata {
+            inner,
+            content_type: OnceCell::new(),
+        }
+    }
+}
 
 impl Metadata for PythonMetadata<'_> {
-    fn content_type(&self) -> Option<String> {
-        self.0
-            .getattr("content_type")
-            .ok()?
-            .extract::<String>()
-            .ok()
+    fn content_type(&self) -> Option<&str> {
+        self.content_type
+            .get_or_init(|| {
+                self.inner
+                    .getattr("content_type")
+                    .ok()?
+                    .extract::<String>()
+                    .ok()
+            })
+            .as_ref()
+            .map(|ct| ct.as_str())
     }
 }
 
 impl<'py> FromPyObject<'py> for PythonMetadata<'py> {
     fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
-        Ok(PythonMetadata(ob.clone()))
+        Ok(PythonMetadata::new(ob.clone()))
     }
 }
 

src/processes/mod.rs

@@ -1,6 +1,5 @@
 pub use image;
 use std::path::PathBuf;
-use std::sync::LazyLock;
 
 #[cfg(any(target_os = "linux", target_os = "macos"))]
 mod nix_list;