Handle head request on onboarding app: Fix #2324 (#2325)

2 files changed, 43 insertions, 4 deletions

diff --git a/mitmproxy/addons/onboardingapp/app.py b/mitmproxy/addons/onboardingapp/app.py
index d418952c..0f09e32c 100644
--- a/mitmproxy/addons/onboardingapp/app.py
+++ b/mitmproxy/addons/onboardingapp/app.py
@@ -44,6 +44,18 @@ class PEM(tornado.web.RequestHandler):
     def filename(self):
         return config.CONF_BASENAME + "-ca-cert.pem"
 
+    def head(self):
+        p = os.path.join(self.request.master.options.cadir, self.filename)
+        p = os.path.expanduser(p)
+        content_length = os.path.getsize(p)
+
+        self.set_header("Content-Type", "application/x-x509-ca-cert")
+        self.set_header(
+            "Content-Disposition",
+            "inline; filename={}".format(
+                self.filename))
+        self.set_header("Content-Length", content_length)
+
     def get(self):
         p = os.path.join(self.request.master.options.cadir, self.filename)
         p = os.path.expanduser(p)
@@ -63,6 +75,19 @@ class P12(tornado.web.RequestHandler):
     def filename(self):
         return config.CONF_BASENAME + "-ca-cert.p12"
 
+    def head(self):
+        p = os.path.join(self.request.master.options.cadir, self.filename)
+        p = os.path.expanduser(p)
+        content_length = os.path.getsize(p)
+
+        self.set_header("Content-Type", "application/x-pkcs12")
+        self.set_header(
+            "Content-Disposition",
+            "inline; filename={}".format(
+                self.filename))
+
+        self.set_header("Content-Length", content_length)
+
     def get(self):
         p = os.path.join(self.request.master.options.cadir, self.filename)
         p = os.path.expanduser(p)
diff --git a/test/mitmproxy/addons/test_onboarding.py b/test/mitmproxy/addons/test_onboarding.py
index 42a3b574..474e6c3c 100644
--- a/test/mitmproxy/addons/test_onboarding.py
+++ b/test/mitmproxy/addons/test_onboarding.py
@@ -1,5 +1,8 @@
+import pytest
+
 from mitmproxy.addons import onboarding
 from mitmproxy.test import taddons
+from mitmproxy import options
 from .. import tservers
 
 
@@ -12,10 +15,21 @@ class TestApp(tservers.HTTPProxyTest):
             tctx.configure(self.addons()[0])
             assert self.app("/").status_code == 200
 
-    def test_cert(self):
+    @pytest.mark.parametrize("ext", ["pem", "p12"])
+    def test_cert(self, ext):
+        with taddons.context() as tctx:
+            tctx.configure(self.addons()[0])
+            resp = self.app("/cert/%s" % ext)
+            assert resp.status_code == 200
+            assert resp.content
+
+    @pytest.mark.parametrize("ext", ["pem", "p12"])
+    def test_head(self, ext):
         with taddons.context() as tctx:
             tctx.configure(self.addons()[0])
-            for ext in ["pem", "p12"]:
-                resp = self.app("/cert/%s" % ext)
+            p = self.pathoc()
+            with p.connect():
+                resp = p.request("head:'http://%s/cert/%s'" % (options.APP_HOST, ext))
                 assert resp.status_code == 200
-                assert resp.content
+                assert "Content-Length" in resp.headers
+                assert not resp.content