feat: X-Proxy
This commit is contained in:
+380
-52
@@ -9,6 +9,8 @@ import tarfile
|
||||
import typing
|
||||
import unittest.mock
|
||||
|
||||
import requests
|
||||
|
||||
from src.handlers import BaseHandler, RequestHandler, UpgradeHandler
|
||||
from src.page import Page
|
||||
from src.params import Parameters
|
||||
@@ -22,6 +24,7 @@ class BaseHandlerTestCase(BaseTestCase, abc.ABC):
|
||||
self,
|
||||
path: str = "/",
|
||||
headers: dict[str, str | None] | None = None,
|
||||
method: str = "GET",
|
||||
rfile: io.BufferedIOBase | None = None,
|
||||
) -> BaseHandler:
|
||||
pass
|
||||
@@ -70,7 +73,7 @@ class BaseHandlerTestCase(BaseTestCase, abc.ABC):
|
||||
send_header_mock.assert_has_calls(
|
||||
[
|
||||
unittest.mock.call("Content-Length", str(len(body.encode()))),
|
||||
unittest.mock.call("Content-type", f"{content_type}; charset=UTF-8"),
|
||||
unittest.mock.call("Content-Type", f"{content_type}; charset=UTF-8"),
|
||||
]
|
||||
+ [unittest.mock.call(header, value) for header, value in headers.items()],
|
||||
any_order=True,
|
||||
@@ -129,6 +132,7 @@ class TestRequestHandler(BaseHandlerTestCase):
|
||||
self,
|
||||
path: str = "/",
|
||||
headers: dict[str, str | None] | None = None,
|
||||
method: str = "GET",
|
||||
rfile: io.BufferedIOBase | None = None,
|
||||
) -> RequestHandler:
|
||||
if headers is None:
|
||||
@@ -146,9 +150,11 @@ class TestRequestHandler(BaseHandlerTestCase):
|
||||
token_manager=self.token_manager,
|
||||
)
|
||||
handler.address_string = lambda: "127.0.0.1" # ty:ignore[invalid-assignment]
|
||||
handler.requestline = "GET /"
|
||||
handler.requestline = f"{method} {path}"
|
||||
handler.path = path
|
||||
handler.command = method
|
||||
handler.request_version = "HTTP/0.9"
|
||||
handler.client_address = ("127.0.0.1", 12345)
|
||||
handler.headers = collections.defaultdict(lambda: None, headers) # ty:ignore[invalid-assignment]
|
||||
handler.rfile = rfile if rfile is not None else io.BytesIO()
|
||||
handler.wfile = io.BytesIO()
|
||||
@@ -156,25 +162,7 @@ class TestRequestHandler(BaseHandlerTestCase):
|
||||
handler.data_dir = self.data_dir
|
||||
return handler
|
||||
|
||||
def test_do_head_redirect(self) -> None:
|
||||
handler = self._get_handler("/path")
|
||||
with (
|
||||
self.mock_call(
|
||||
self.registry.get_from_path,
|
||||
["path"],
|
||||
Page("path", redirect="https://example.com"),
|
||||
),
|
||||
self.expects_status_only(
|
||||
handler,
|
||||
http.HTTPStatus.MOVED_PERMANENTLY,
|
||||
headers={"Location": "https://example.com"},
|
||||
),
|
||||
self.patch("http.server.SimpleHTTPRequestHandler.do_HEAD", count=0),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
handler.do_HEAD()
|
||||
|
||||
def test_do_head_proxy(self) -> None:
|
||||
def test_do_head_forward(self) -> None:
|
||||
handler = self._get_handler()
|
||||
with (
|
||||
self.patch("http.server.SimpleHTTPRequestHandler.do_HEAD"),
|
||||
@@ -191,37 +179,16 @@ class TestRequestHandler(BaseHandlerTestCase):
|
||||
):
|
||||
handler.do_GET()
|
||||
|
||||
def test_do_get_redirect(self) -> None:
|
||||
handler = self._get_handler("/path")
|
||||
with (
|
||||
self.mock_call(
|
||||
self.registry.get_from_path,
|
||||
["path"],
|
||||
Page("path", redirect="https://example.com"),
|
||||
),
|
||||
self.expects_status_only(
|
||||
handler,
|
||||
http.HTTPStatus.MOVED_PERMANENTLY,
|
||||
headers={"Location": "https://example.com"},
|
||||
),
|
||||
self.patch("http.server.SimpleHTTPRequestHandler.do_GET", count=0),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
handler.do_GET()
|
||||
|
||||
def test_do_get_proxy_on_other_path(self) -> None:
|
||||
def test_do_get_forward_on_other_path(self) -> None:
|
||||
handler = self._get_handler("/file")
|
||||
with (
|
||||
self.mock_call(
|
||||
self.registry.get_from_path,
|
||||
["file"],
|
||||
),
|
||||
self.mock_call(self.registry.get_from_path, ["file"], Page("file")),
|
||||
self.patch("http.server.SimpleHTTPRequestHandler.do_GET"),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
handler.do_GET()
|
||||
|
||||
def test_do_get_proxy_on_other_host(self) -> None:
|
||||
def test_do_get_forward_on_other_host(self) -> None:
|
||||
handler = self._get_handler("/", {"Host": "other_host"})
|
||||
with (
|
||||
self.mock_call(
|
||||
@@ -236,6 +203,7 @@ class TestRequestHandler(BaseHandlerTestCase):
|
||||
def test_do_put_no_token(self) -> None:
|
||||
handler = self._get_handler("/path")
|
||||
with (
|
||||
self.mock_call(self.registry.get_from_path, ["path"]),
|
||||
self.expects_error(
|
||||
handler, http.HTTPStatus.BAD_REQUEST, "No X-Token header in request"
|
||||
),
|
||||
@@ -246,6 +214,7 @@ class TestRequestHandler(BaseHandlerTestCase):
|
||||
def test_do_post_is_do_put(self) -> None:
|
||||
handler = self._get_handler("/path")
|
||||
with (
|
||||
self.mock_call(self.registry.get_from_path, ["path"]),
|
||||
self.expects_error(
|
||||
handler, http.HTTPStatus.BAD_REQUEST, "No X-Token header in request"
|
||||
),
|
||||
@@ -256,6 +225,7 @@ class TestRequestHandler(BaseHandlerTestCase):
|
||||
def test_do_patch_is_do_put(self) -> None:
|
||||
handler = self._get_handler("/path")
|
||||
with (
|
||||
self.mock_call(self.registry.get_from_path, ["path"]),
|
||||
self.expects_error(
|
||||
handler, http.HTTPStatus.BAD_REQUEST, "No X-Token header in request"
|
||||
),
|
||||
@@ -525,10 +495,8 @@ class TestRequestHandler(BaseHandlerTestCase):
|
||||
["secret", "path"],
|
||||
True, # noqa: FBT003
|
||||
),
|
||||
self.mock_call(self.data_dir.empty, ["path"]),
|
||||
self.mock_call(self.registry.add, ["path"]),
|
||||
self.mock_call(self.token_manager.set_token, ["path", "secret"]),
|
||||
self.mock_call(self.registry.set_redirect, ["path", "https://example.com"]),
|
||||
self.mock_call(self.token_manager.set_token, ["path", "secret"]),
|
||||
self.expects_status_only(
|
||||
handler, http.HTTPStatus.CREATED, "Resource /path/ updated"
|
||||
),
|
||||
@@ -553,10 +521,8 @@ class TestRequestHandler(BaseHandlerTestCase):
|
||||
True, # noqa: FBT003
|
||||
),
|
||||
self.mock_call(self.registry.get_from_host, ["example.com"], Page("path")),
|
||||
self.mock_call(self.data_dir.empty, ["path"]),
|
||||
self.mock_call(self.registry.add, ["path"]),
|
||||
self.mock_call(self.token_manager.set_token, ["path", "secret"]),
|
||||
self.mock_call(self.registry.set_redirect, ["path", "https://example.com"]),
|
||||
self.mock_call(self.token_manager.set_token, ["path", "secret"]),
|
||||
self.mock_call(self.cert_manager.create_or_update, ["example.com"], True), # noqa: FBT003
|
||||
self.mock_call(self.registry.set_host, ["path", "example.com"]),
|
||||
self.expects_status_only(
|
||||
@@ -566,9 +532,114 @@ class TestRequestHandler(BaseHandlerTestCase):
|
||||
):
|
||||
handler.do_PUT()
|
||||
|
||||
def test_do_put_proxy_with_content(self) -> None:
|
||||
handler = self._get_handler(
|
||||
"/path",
|
||||
{
|
||||
"X-Token": "secret",
|
||||
"X-Proxy": "https://example.com",
|
||||
"Content-Length": "1",
|
||||
},
|
||||
)
|
||||
with (
|
||||
self.mock_call(self.token_manager.is_valid, ["secret"], True), # noqa: FBT003
|
||||
self.mock_call(
|
||||
self.token_manager.is_valid_for_path,
|
||||
["secret", "path"],
|
||||
True, # noqa: FBT003
|
||||
),
|
||||
self.expects_error(
|
||||
handler,
|
||||
http.HTTPStatus.BAD_REQUEST,
|
||||
"No content must be sent with X-Proxy",
|
||||
),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
handler.do_PUT()
|
||||
|
||||
def test_do_put_proxy_ok(self) -> None:
|
||||
handler = self._get_handler(
|
||||
"/path",
|
||||
{
|
||||
"X-Token": "secret",
|
||||
"X-Proxy": "https://example.com",
|
||||
},
|
||||
)
|
||||
with (
|
||||
self.mock_call(self.token_manager.is_valid, ["secret"], True), # noqa: FBT003
|
||||
self.mock_call(
|
||||
self.token_manager.is_valid_for_path,
|
||||
["secret", "path"],
|
||||
True, # noqa: FBT003
|
||||
),
|
||||
self.mock_call(self.registry.set_proxy, ["path", "https://example.com"]),
|
||||
self.mock_call(self.token_manager.set_token, ["path", "secret"]),
|
||||
self.expects_status_only(
|
||||
handler, http.HTTPStatus.CREATED, "Resource /path/ updated"
|
||||
),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
handler.do_PUT()
|
||||
|
||||
def test_do_put_proxy_with_host(self) -> None:
|
||||
handler = self._get_handler(
|
||||
"/path",
|
||||
{
|
||||
"X-Token": "secret",
|
||||
"X-Proxy": "https://example.com",
|
||||
"X-Host": "example.com",
|
||||
},
|
||||
)
|
||||
with (
|
||||
self.mock_call(self.token_manager.is_valid, ["secret"], True), # noqa: FBT003
|
||||
self.mock_call(
|
||||
self.token_manager.is_valid_for_path,
|
||||
["secret", "path"],
|
||||
True, # noqa: FBT003
|
||||
),
|
||||
self.mock_call(self.registry.get_from_host, ["example.com"], Page("path")),
|
||||
self.mock_call(self.registry.set_proxy, ["path", "https://example.com"]),
|
||||
self.mock_call(self.token_manager.set_token, ["path", "secret"]),
|
||||
self.mock_call(self.cert_manager.create_or_update, ["example.com"], True), # noqa: FBT003
|
||||
self.mock_call(self.registry.set_host, ["path", "example.com"]),
|
||||
self.expects_status_only(
|
||||
handler, http.HTTPStatus.CREATED, "Resource /path/ updated"
|
||||
),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
handler.do_PUT()
|
||||
|
||||
def test_do_put_proxy_and_redirect(self) -> None:
|
||||
handler = self._get_handler(
|
||||
"/path",
|
||||
{
|
||||
"X-Token": "secret",
|
||||
"X-Proxy": "https://example.com",
|
||||
"X-Redirect": "https://example.com",
|
||||
"X-Host": "example.com",
|
||||
},
|
||||
)
|
||||
with (
|
||||
self.mock_call(self.token_manager.is_valid, ["secret"], True), # noqa: FBT003
|
||||
self.mock_call(
|
||||
self.token_manager.is_valid_for_path,
|
||||
["secret", "path"],
|
||||
True, # noqa: FBT003
|
||||
),
|
||||
self.mock_call(self.registry.get_from_host, ["example.com"], Page("path")),
|
||||
self.expects_status_only(
|
||||
handler,
|
||||
http.HTTPStatus.BAD_REQUEST,
|
||||
"Cannot use X-Proxy with X-Redirect",
|
||||
),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
handler.do_PUT()
|
||||
|
||||
def test_do_delete_no_token(self) -> None:
|
||||
handler = self._get_handler("/path")
|
||||
with (
|
||||
self.mock_call(self.registry.get_from_path, ["path"]),
|
||||
self.expects_error(
|
||||
handler, http.HTTPStatus.BAD_REQUEST, "No X-Token header in request"
|
||||
),
|
||||
@@ -662,6 +733,261 @@ class TestRequestHandler(BaseHandlerTestCase):
|
||||
):
|
||||
handler.do_DELETE()
|
||||
|
||||
def test_do_post_proxy_no_body(self) -> None:
|
||||
handler = self._get_handler("/path", method="POST")
|
||||
response = requests.Response()
|
||||
response.status_code = 200
|
||||
response.reason = "OK"
|
||||
response.raw = io.BytesIO()
|
||||
with (
|
||||
self.mock_call(
|
||||
self.registry.get_from_path,
|
||||
["path"],
|
||||
Page("path", proxy="https://example.com"),
|
||||
),
|
||||
self.patch_call(
|
||||
"requests.request",
|
||||
[
|
||||
"POST",
|
||||
"https://example.com",
|
||||
],
|
||||
response,
|
||||
{
|
||||
"data": None,
|
||||
"headers": {
|
||||
"Host": "example.com",
|
||||
"X-Forwarded-For": "127.0.0.1",
|
||||
"X-Real-IP": "127.0.0.1",
|
||||
},
|
||||
"timeout": 240,
|
||||
},
|
||||
),
|
||||
self.expects_status_only(handler, 200, "OK"),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
handler.do_POST()
|
||||
|
||||
def test_do_post_proxy_with_request_body(self) -> None:
|
||||
handler = self._get_handler(
|
||||
"/path",
|
||||
method="POST",
|
||||
headers={"Content-Length": "5"},
|
||||
rfile=io.BytesIO(b"hello"),
|
||||
)
|
||||
response = requests.Response()
|
||||
response.status_code = 200
|
||||
response.reason = "OK"
|
||||
response.raw = io.BytesIO()
|
||||
with (
|
||||
self.mock_call(
|
||||
self.registry.get_from_path,
|
||||
["path"],
|
||||
Page("path", proxy="https://example.com"),
|
||||
),
|
||||
self.patch_call(
|
||||
"requests.request",
|
||||
[
|
||||
"POST",
|
||||
"https://example.com",
|
||||
],
|
||||
response,
|
||||
{
|
||||
"data": b"hello",
|
||||
"headers": {
|
||||
"Host": "example.com",
|
||||
"X-Forwarded-For": "127.0.0.1",
|
||||
"X-Real-IP": "127.0.0.1",
|
||||
"Content-Length": "5",
|
||||
},
|
||||
"timeout": 240,
|
||||
},
|
||||
),
|
||||
self.expects_status_only(handler, 200, "OK"),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
handler.do_POST()
|
||||
|
||||
def test_do_post_proxy_with_response_body(self) -> None:
|
||||
handler = self._get_handler(
|
||||
"/path",
|
||||
method="POST",
|
||||
)
|
||||
response = requests.Response()
|
||||
response.status_code = 200
|
||||
response.reason = "OK"
|
||||
response.headers["Content-Type"] = "text/plain; charset=UTF-8"
|
||||
response.raw = io.BytesIO(b"hello")
|
||||
with (
|
||||
self.mock_call(
|
||||
self.registry.get_from_path,
|
||||
["path"],
|
||||
Page("path", proxy="https://example.com"),
|
||||
),
|
||||
self.patch_call(
|
||||
"requests.request",
|
||||
[
|
||||
"POST",
|
||||
"https://example.com",
|
||||
],
|
||||
response,
|
||||
{
|
||||
"data": None,
|
||||
"headers": {
|
||||
"Host": "example.com",
|
||||
"X-Forwarded-For": "127.0.0.1",
|
||||
"X-Real-IP": "127.0.0.1",
|
||||
},
|
||||
"timeout": 240,
|
||||
},
|
||||
),
|
||||
self.expects_basic_body(handler, "hello", message="OK"),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
handler.do_POST()
|
||||
|
||||
def test_do_post_proxy_fail(self) -> None:
|
||||
handler = self._get_handler("/path", method="POST")
|
||||
with (
|
||||
self.mock_call(
|
||||
self.registry.get_from_path,
|
||||
["path"],
|
||||
Page("path", proxy="https://example.com"),
|
||||
),
|
||||
self.patch_call(
|
||||
"requests.request",
|
||||
[
|
||||
"POST",
|
||||
"https://example.com",
|
||||
],
|
||||
None,
|
||||
{
|
||||
"data": None,
|
||||
"headers": {
|
||||
"Host": "example.com",
|
||||
"X-Forwarded-For": "127.0.0.1",
|
||||
"X-Real-IP": "127.0.0.1",
|
||||
},
|
||||
"timeout": 240,
|
||||
},
|
||||
) as request_mock,
|
||||
self.expects_status_only(
|
||||
handler,
|
||||
http.HTTPStatus.BAD_GATEWAY,
|
||||
"Could not reach https://example.com",
|
||||
),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
request_mock.side_effect = Exception
|
||||
handler.do_POST()
|
||||
|
||||
def test_do_post_proxy_sub_path(self) -> None:
|
||||
handler = self._get_handler("/path/index.html", method="POST")
|
||||
response = requests.Response()
|
||||
response.status_code = 200
|
||||
response.reason = "OK"
|
||||
response.raw = io.BytesIO()
|
||||
with (
|
||||
self.mock_call(
|
||||
self.registry.get_from_path,
|
||||
["path"],
|
||||
Page("path", proxy="https://example.com"),
|
||||
),
|
||||
self.patch_call(
|
||||
"requests.request",
|
||||
[
|
||||
"POST",
|
||||
"https://example.com/index.html",
|
||||
],
|
||||
response,
|
||||
{
|
||||
"data": None,
|
||||
"headers": {
|
||||
"Host": "example.com",
|
||||
"X-Forwarded-For": "127.0.0.1",
|
||||
"X-Real-IP": "127.0.0.1",
|
||||
},
|
||||
"timeout": 240,
|
||||
},
|
||||
),
|
||||
self.expects_status_only(handler, 200, "OK"),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
handler.do_POST()
|
||||
|
||||
def test_do_post_proxy_sub_path_for_host(self) -> None:
|
||||
handler = self._get_handler(
|
||||
"/path/index.html", method="POST", headers={"Host": "host"}
|
||||
)
|
||||
response = requests.Response()
|
||||
response.status_code = 200
|
||||
response.reason = "OK"
|
||||
response.raw = io.BytesIO()
|
||||
with (
|
||||
self.mock_call(
|
||||
self.registry.get_from_host,
|
||||
["host"],
|
||||
Page("path", proxy="https://example.com"),
|
||||
),
|
||||
self.patch_call(
|
||||
"requests.request",
|
||||
[
|
||||
"POST",
|
||||
"https://example.com/path/index.html",
|
||||
],
|
||||
response,
|
||||
{
|
||||
"data": None,
|
||||
"headers": {
|
||||
"Host": "example.com",
|
||||
"X-Forwarded-For": "127.0.0.1",
|
||||
"X-Real-IP": "127.0.0.1",
|
||||
},
|
||||
"timeout": 240,
|
||||
},
|
||||
),
|
||||
self.expects_status_only(handler, 200, "OK"),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
handler.do_POST()
|
||||
|
||||
def test_do_method_not_supported(self) -> None:
|
||||
handler = self._get_handler("/path")
|
||||
for method in ["CONNECT", "TRACE", "OPTIONS"]:
|
||||
with (
|
||||
self.subTest(None, method=method),
|
||||
self.mock_call(
|
||||
self.registry.get_from_path,
|
||||
["path"],
|
||||
),
|
||||
self.expects_status_only(
|
||||
handler, http.HTTPStatus.METHOD_NOT_ALLOWED, "Method Not Allowed"
|
||||
),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
getattr(handler, f"do_{method}")()
|
||||
|
||||
def test_do_redirect(self) -> None:
|
||||
handler = self._get_handler("/path")
|
||||
for method in [method.value for method in http.HTTPMethod]:
|
||||
with (
|
||||
self.subTest(None, method=method),
|
||||
self.mock_call(
|
||||
self.registry.get_from_path,
|
||||
["path"],
|
||||
Page("path", redirect="https://example.com"),
|
||||
),
|
||||
self.expects_status_only(
|
||||
handler,
|
||||
http.HTTPStatus.MOVED_PERMANENTLY,
|
||||
headers={"Location": "https://example.com"},
|
||||
),
|
||||
self.patch(
|
||||
f"http.server.SimpleHTTPRequestHandler.do_{method}", count=0
|
||||
),
|
||||
self.seal_mocks(),
|
||||
):
|
||||
getattr(handler, f"do_{method}")()
|
||||
|
||||
def test_list_directory(self) -> None:
|
||||
handler = self._get_handler("/path/", {"Accept": "text/html"})
|
||||
with (
|
||||
@@ -780,6 +1106,7 @@ class TestUpgradeHandler(BaseHandlerTestCase):
|
||||
self,
|
||||
path: str = "/",
|
||||
headers: dict[str, str | None] | None = None,
|
||||
method: str = "GET",
|
||||
rfile: io.BufferedIOBase | None = None,
|
||||
) -> UpgradeHandler:
|
||||
if headers is None:
|
||||
@@ -792,8 +1119,9 @@ class TestUpgradeHandler(BaseHandlerTestCase):
|
||||
params=Parameters(),
|
||||
)
|
||||
handler.address_string = lambda: "127.0.0.1" # ty:ignore[invalid-assignment]
|
||||
handler.requestline = "GET /"
|
||||
handler.requestline = f"{method} {path}"
|
||||
handler.path = path
|
||||
handler.command = method
|
||||
handler.request_version = "HTTP/0.9"
|
||||
handler.headers = collections.defaultdict(lambda: None, headers) # ty:ignore[invalid-assignment]
|
||||
handler.rfile = rfile if rfile is not None else io.BytesIO()
|
||||
|
||||
Reference in New Issue
Block a user