fix metadata interop, add testing-only content view

Author: mhils

mitmproxy-contentviews/src/lib.rs

@@ -2,15 +2,19 @@ mod hex_dump;
 mod hex_stream;
 mod msgpack;
 mod protobuf;
+mod test_inspect_metadata;
 
 pub use hex_dump::HexDump;
 pub use hex_stream::HexStream;
 pub use msgpack::MsgPack;
 pub use protobuf::Protobuf;
 pub use protobuf::GRPC;
+pub use test_inspect_metadata::TestInspectMetadata;
 
 use anyhow::Result;
 use mitmproxy_highlight::Language;
+
+use serde::Serialize;
 use std::path::Path;
 
 pub trait Metadata {
@@ -66,10 +70,9 @@ pub trait Reencode: Send + Sync {
 
 // no cfg(test) gate because it's used in benchmarks as well
 pub mod test {
-    use crate::Metadata;
-    use std::path::Path;
+    use super::*;
 
-    #[derive(Default)]
+    #[derive(Default, Serialize)]
     pub struct TestMetadata {
         pub content_type: Option<String>,
         pub headers: std::collections::HashMap<String, String>,

mitmproxy-contentviews/src/test_inspect_metadata.rs

@@ -0,0 +1,52 @@
+use crate::test::TestMetadata;
+use crate::{Metadata, Prettify};
+use anyhow::Context;
+use std::collections::HashMap;
+use std::path::Path;
+
+/// Contentview used for internal testing to ensure that the
+/// Python accessors in mitmproxy-rs all work properly.
+pub struct TestInspectMetadata;
+
+impl Prettify for TestInspectMetadata {
+    fn name(&self) -> &'static str {
+        "Inspect Metadata (test only)"
+    }
+
+    fn instance_name(&self) -> String {
+        "_test_inspect_metadata".to_string()
+    }
+
+    fn prettify(&self, _data: &[u8], metadata: &dyn Metadata) -> anyhow::Result<String> {
+        let mut headers = HashMap::new();
+        if let Some(host) = metadata.get_header("host") {
+            headers.insert("host".to_string(), host);
+        }
+        let meta = TestMetadata {
+            content_type: metadata.content_type().map(str::to_string),
+            headers,
+            path: metadata.get_path().map(str::to_string),
+            is_http_request: metadata.is_http_request(),
+            protobuf_definitions: metadata.protobuf_definitions().map(Path::to_path_buf),
+        };
+        // JSON would be nicer to consume on the Python side, 
+        // but let's not add dependencies for this.
+        serde_yaml::to_string(&meta).context("Failed to convert to YAML")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn prettify_simple() {
+        let result = TestInspectMetadata
+            .prettify(b"", &TestMetadata::default())
+            .unwrap();
+        assert_eq!(
+            result,
+            "content_type: null\nheaders: {}\nprotobuf_definitions: null\npath: null\nis_http_request: false\n"
+        );
+    }
+}

mitmproxy-rs/mitmproxy_rs/contentviews.pyi

@@ -18,6 +18,7 @@ class InteractiveContentview(Contentview):
     def reencode(self, data: str, metadata) -> bytes:
         pass
 
+_test_inspect_metadata: Contentview
 hex_dump: Contentview
 hex_stream: InteractiveContentview
 msgpack: InteractiveContentview
@@ -32,4 +33,5 @@ __all__ = [
     "msgpack",
     "protobuf",
     "grpc",
+    "_test_inspect_metadata",
 ]

mitmproxy-rs/src/contentviews.rs

@@ -1,7 +1,7 @@
 use mitmproxy_contentviews::{Metadata, Prettify, Reencode};
 use pyo3::{exceptions::PyValueError, prelude::*};
 use std::cell::OnceCell;
-use std::path::Path;
+use std::path::{Path, PathBuf};
 
 pub struct PythonMetadata<'py> {
     inner: Bound<'py, PyAny>,
@@ -56,9 +56,8 @@ impl Metadata for PythonMetadata<'_> {
                 self.inner
                     .getattr("protobuf_definitions")
                     .ok()?
-                    .extract::<String>()
+                    .extract::<PathBuf>()
                     .ok()
-                    .map(std::path::PathBuf::from)
             })
             .as_deref()
     }
@@ -67,16 +66,13 @@ impl Metadata for PythonMetadata<'_> {
         let Ok(http_message) = self.inner.getattr("http_message") else {
             return false;
         };
-        let Ok(flow) = self
+        let Ok(request) = self
             .inner
             .getattr("flow")
             .and_then(|flow| flow.getattr("request"))
         else {
             return false;
         };
-        let Ok(request) = flow.getattr("request") else {
-            return false;
-        };
         http_message.is(&request)
     }
 }

mitmproxy-rs/src/lib.rs

@@ -91,7 +91,9 @@ mod mitmproxy_rs {
         use crate::contentviews::Contentview;
         #[pymodule_export]
         use crate::contentviews::InteractiveContentview;
-        use mitmproxy_contentviews::{HexDump, HexStream, MsgPack, Protobuf, GRPC};
+        use mitmproxy_contentviews::{
+            HexDump, HexStream, MsgPack, Protobuf, TestInspectMetadata, GRPC,
+        };
 
         #[pymodule_init]
         fn init(m: &Bound<'_, PyModule>) -> PyResult<()> {
@@ -100,6 +102,7 @@ mod mitmproxy_rs {
             m.add_interactive_contentview(&MsgPack)?;
             m.add_interactive_contentview(&Protobuf)?;
             m.add_interactive_contentview(&GRPC)?;
+            m.add_contentview(&TestInspectMetadata)?;
             Ok(())
         }
     }