unify server spec parsing

7 files changed, 118 insertions, 54 deletions

diff --git a/mitmproxy/addons/core.py b/mitmproxy/addons/core.py
index 3de1638c..4c6a5516 100644
--- a/mitmproxy/addons/core.py
+++ b/mitmproxy/addons/core.py
@@ -3,8 +3,8 @@
     checked by other addons.
 """
 from mitmproxy import exceptions
-from mitmproxy import options
 from mitmproxy import platform
+from mitmproxy.net import server_spec
 from mitmproxy.utils import human
 
 
@@ -30,11 +30,10 @@ class Core:
         if "mode" in updated:
             mode = opts.mode
             if mode.startswith("reverse:") or mode.startswith("upstream:"):
-                spec = options.get_mode_spec(mode)
-                if not spec:
-                    raise exceptions.OptionsError(
-                        "Invalid mode specification: %s" % mode
-                    )
+                try:
+                    server_spec.parse_with_mode(mode)
+                except ValueError as e:
+                    raise exceptions.OptionsError(str(e)) from e
             elif mode == "transparent":
                 if not platform.original_addr:
                     raise exceptions.OptionsError(
diff --git a/mitmproxy/net/check.py b/mitmproxy/net/check.py
index d30c1df6..aaea851f 100644
--- a/mitmproxy/net/check.py
+++ b/mitmproxy/net/check.py
@@ -29,5 +29,5 @@ def is_valid_host(host: bytes) -> bool:
         return False
 
 
-def is_valid_port(port):
+def is_valid_port(port: int) -> bool:
     return 0 <= port <= 65535
diff --git a/mitmproxy/net/server_spec.py b/mitmproxy/net/server_spec.py
new file mode 100644
index 00000000..efbf1012
--- /dev/null
+++ b/mitmproxy/net/server_spec.py
@@ -0,0 +1,76 @@
+"""
+Parse scheme, host and port from a string.
+"""
+import collections
+import re
+from typing import Tuple
+
+from mitmproxy.net import check
+
+ServerSpec = collections.namedtuple("ServerSpec", ["scheme", "address"])
+
+server_spec_re = re.compile(
+    r"""
+        ^
+        (?:(?P<scheme>\w+)://)?  # scheme is optional
+        (?P<host>[^:/]+|\[.+\])  # hostname can be DNS name, IPv4, or IPv6 address.
+        (?::(?P<port>\d+))?  #  port is optional
+        /?  #  we allow a trailing backslash, but no path
+        $
+        """,
+    re.VERBOSE
+)
+
+
+def parse(server_spec: str) -> ServerSpec:
+    """
+    Parses a server mode specification, e.g.:
+
+        - http://example.com/
+        - example.org
+        - example.com:443
+
+    Raises:
+        ValueError, if the server specification is invalid.
+    """
+    m = server_spec_re.match(server_spec)
+    if not m:
+        raise ValueError("Invalid server specification: {}".format(server_spec))
+
+    # defaulting to https/port 443 may annoy some folks, but it's secure-by-default.
+    scheme = m.group("scheme") or "https"
+    if scheme not in ("http", "https"):
+        raise ValueError("Invalid server scheme: {}".format(scheme))
+
+    host = m.group("host")
+    # IPv6 brackets
+    if host.startswith("[") and host.endswith("]"):
+        host = host[1:-1]
+    if not check.is_valid_host(host.encode("idna")):
+        raise ValueError("Invalid hostname: {}".format(host))
+
+    if m.group("port"):
+        port = int(m.group("port"))
+    else:
+        port = {
+            "http": 80,
+            "https": 443
+        }[scheme]
+    if not check.is_valid_port(port):
+        raise ValueError("Invalid port: {}".format(port))
+
+    return ServerSpec(scheme, (host, port))
+
+
+def parse_with_mode(mode: str) -> Tuple[str, ServerSpec]:
+    """
+    Parse a proxy mode specification, which is usually just (reverse|upstream):server-spec
+
+    Returns:
+        A (mode, server_spec) tuple.
+
+    Raises:
+        ValueError, if the specification is invalid.
+    """
+    mode, server_spec = mode.split(":", maxsplit=1)
+    return mode, parse(server_spec)
diff --git a/mitmproxy/options.py b/mitmproxy/options.py
index 08d97cad..1063bab9 100644
--- a/mitmproxy/options.py
+++ b/mitmproxy/options.py
@@ -22,11 +22,6 @@ view_orders = [
     "size",
 ]
 
-
-def get_mode_spec(m):
-    return m.split(":", maxsplit=1)[1]
-
-
 APP_HOST = "mitm.it"
 APP_PORT = 80
 CA_DIR = "~/.mitmproxy"
diff --git a/mitmproxy/proxy/config.py b/mitmproxy/proxy/config.py
index 9cf2b00f..8417ebad 100644
--- a/mitmproxy/proxy/config.py
+++ b/mitmproxy/proxy/config.py
@@ -1,4 +1,3 @@
-import collections
 import os
 import re
 from typing import Any
@@ -9,7 +8,7 @@ from mitmproxy import exceptions
 from mitmproxy import options as moptions
 from mitmproxy import certs
 from mitmproxy.net import tcp
-from mitmproxy.net.http import url
+from mitmproxy.net import server_spec
 
 CONF_BASENAME = "mitmproxy"
 
@@ -33,24 +32,6 @@ class HostMatcher:
         return bool(self.patterns)
 
 
-ServerSpec = collections.namedtuple("ServerSpec", "scheme address")
-
-
-def parse_server_spec(spec):
-    try:
-        p = url.parse(spec)
-        if p[0] not in (b"http", b"https"):
-            raise ValueError()
-    except ValueError:
-        raise exceptions.OptionsError(
-            "Invalid server specification: %s" % spec
-        )
-    host, port = p[1:3]
-    address = (host.decode("ascii"), port)
-    scheme = p[0].decode("ascii").lower()
-    return ServerSpec(scheme, address)
-
-
 class ProxyConfig:
 
     def __init__(self, options: moptions.Options) -> None:
@@ -123,5 +104,5 @@ class ProxyConfig:
                 )
         m = options.mode
         if m.startswith("upstream:") or m.startswith("reverse:"):
-            spec = moptions.get_mode_spec(options.mode)
-            self.upstream_server = parse_server_spec(spec)
+            _, spec = server_spec.parse_with_mode(options.mode)
+            self.upstream_server = spec
diff --git a/test/mitmproxy/net/test_server_spec.py b/test/mitmproxy/net/test_server_spec.py
new file mode 100644
index 00000000..095ad519
--- /dev/null
+++ b/test/mitmproxy/net/test_server_spec.py
@@ -0,0 +1,32 @@
+import pytest
+
+from mitmproxy.net import server_spec
+
+
+def test_parse():
+    assert server_spec.parse("example.com") == ("https", ("example.com", 443))
+    assert server_spec.parse("example.com") == ("https", ("example.com", 443))
+    assert server_spec.parse("http://example.com") == ("http", ("example.com", 80))
+    assert server_spec.parse("http://127.0.0.1") == ("http", ("127.0.0.1", 80))
+    assert server_spec.parse("http://[::1]") == ("http", ("::1", 80))
+    assert server_spec.parse("http://[::1]/") == ("http", ("::1", 80))
+    assert server_spec.parse("https://[::1]/") == ("https", ("::1", 443))
+    assert server_spec.parse("http://[::1]:8080") == ("http", ("::1", 8080))
+
+    with pytest.raises(ValueError, match="Invalid server specification"):
+        server_spec.parse(":")
+
+    with pytest.raises(ValueError, match="Invalid server scheme"):
+        server_spec.parse("ftp://example.com")
+
+    with pytest.raises(ValueError, match="Invalid hostname"):
+        server_spec.parse("$$$")
+
+    with pytest.raises(ValueError, match="Invalid port"):
+        server_spec.parse("example.com:999999")
+
+
+def test_parse_with_mode():
+    assert server_spec.parse_with_mode("m:example.com") == ("m", ("https", ("example.com", 443)))
+    with pytest.raises(ValueError):
+        server_spec.parse_with_mode("moo")
diff --git a/test/mitmproxy/proxy/test_config.py b/test/mitmproxy/proxy/test_config.py
index 4272d952..777ab4dd 100644
--- a/test/mitmproxy/proxy/test_config.py
+++ b/test/mitmproxy/proxy/test_config.py
@@ -1,20 +1 @@
-import pytest
-from mitmproxy.proxy import config
-
-
-def test_parse_server_spec():
-    with pytest.raises(Exception, match="Invalid server specification"):
-        config.parse_server_spec("")
-    assert config.parse_server_spec("http://foo.com:88") == (
-        "http", ("foo.com", 88)
-    )
-    assert config.parse_server_spec("http://foo.com") == (
-        "http", ("foo.com", 80)
-    )
-    assert config.parse_server_spec("https://foo.com") == (
-        "https", ("foo.com", 443)
-    )
-    with pytest.raises(Exception, match="Invalid server specification"):
-        config.parse_server_spec("foo.com")
-    with pytest.raises(Exception, match="Invalid server specification"):
-        config.parse_server_spec("http://")
+# TODO: write tests