upstream-zulip/zerver/tests/test_outgoing_http.py

import os
from typing import Any
from unittest import mock

import botocore.utils
import requests
import responses
from django.test import override_settings
from requests.adapters import HTTPAdapter
from typing_extensions import override
from urllib3.util import Retry

from zerver.apps import skip_smokescreen_proxy_for_boto3
from zerver.lib.outgoing_http import OutgoingSession
from zerver.lib.test_classes import ZulipTestCase


class RequestMockWithProxySupport(responses.RequestsMock):
    @override
    def _on_request(
        self,
        adapter: requests.adapters.HTTPAdapter,
        request: requests.PreparedRequest,
        **kwargs: Any,
    ) -> requests.Response:
        if "proxies" in kwargs and request.url:
            proxy_url = requests.utils.select_proxy(request.url, kwargs["proxies"])
            if proxy_url is not None:
                request = requests.Request(
                    method="GET",
                    url=f"{proxy_url}/",
                    headers=adapter.proxy_headers(proxy_url),
                ).prepare()
        return super()._on_request(adapter, request, **kwargs)


class RequestMockWithTimeoutAsHeader(responses.RequestsMock):
    @override
    def _on_request(
        self,
        adapter: requests.adapters.HTTPAdapter,
        request: requests.PreparedRequest,
        **kwargs: Any,
    ) -> requests.Response:
        if kwargs.get("timeout") is not None:
            request.headers["X-Timeout"] = kwargs["timeout"]
        return super()._on_request(adapter, request, **kwargs)


# The general forcing of requests through Smokescreen only happens
# when DEVELOPMENT=False
@override_settings(DEVELOPMENT=False)
class TestOutgoingHttp(ZulipTestCase):
    def test_headers(self) -> None:
        with RequestMockWithProxySupport() as mock_requests:
            mock_requests.add(responses.GET, "http://example.com/")
            OutgoingSession(role="testing", timeout=1, headers={"X-Foo": "bar"}, proxies={}).get(
                "http://example.com/"
            )
            self.assert_length(mock_requests.calls, 1)
            headers = mock_requests.calls[0].request.headers
            # We don't see a proxy header with no proxy set
            self.assertFalse("X-Smokescreen-Role" in headers)
            self.assertEqual(headers["X-Foo"], "bar")

    def test_proxy_headers_config(self) -> None:
        with (
            RequestMockWithProxySupport() as mock_requests,
            mock.patch("zerver.lib.outgoing_http.get_config") as get_config,
        ):
            get_config.side_effect = ("localhost", "4242")
            mock_requests.add(responses.GET, "http://localhost:4242/")
            OutgoingSession(role="testing", timeout=1, headers={"X-Foo": "bar"}).get(
                "http://example.com/"
            )
            self.assert_length(mock_requests.calls, 1)
            headers = mock_requests.calls[0].request.headers
            self.assertEqual(headers["X-Smokescreen-Role"], "testing")

    def test_proxy_headers_no_config(self) -> None:
        with (
            RequestMockWithProxySupport() as mock_requests,
            mock.patch("zerver.lib.outgoing_http.get_config") as get_config,
        ):
            get_config.side_effect = ("", "")
            mock_requests.add(responses.GET, "http://example.com/")
            OutgoingSession(role="testing", timeout=1).get("http://example.com/")

    def test_proxy_headers_proxy_override(self) -> None:
        with (
            RequestMockWithProxySupport() as mock_requests,
            mock.patch("zerver.lib.outgoing_http.get_config") as get_config,
        ):
            get_config.side_effect = ("localhost", "4242")
            mock_requests.add(responses.GET, "http://proxy-host:4343/")
            OutgoingSession(
                role="testing", timeout=1, proxies={"http": "http://proxy-host:4343"}
            ).get("http://example.com/")

    @mock.patch.dict(os.environ, {"http_proxy": "http://localhost:4242"})
    def test_proxy_headers_env(self) -> None:
        with RequestMockWithProxySupport() as mock_requests:
            mock_requests.add(responses.GET, "http://localhost:4242/")
            OutgoingSession(role="testing", timeout=1, headers={"X-Foo": "bar"}).get(
                "http://example.com/"
            )
            self.assert_length(mock_requests.calls, 1)
            headers = mock_requests.calls[0].request.headers
            self.assertEqual(headers["X-Smokescreen-Role"], "testing")

            # We don't see the request-level headers in the proxy
            # request.  This isn't a _true_ test because we're
            # fiddling the headers above, instead of urllib3 actually
            # setting them.
            self.assertFalse("X-Foo" in headers)

    def test_timeouts(self) -> None:
        with RequestMockWithTimeoutAsHeader() as mock_requests:
            mock_requests.add(responses.GET, "http://example.com/")
            OutgoingSession(role="testing", timeout=17).get("http://example.com/")
            self.assert_length(mock_requests.calls, 1)
            self.assertEqual(mock_requests.calls[0].request.headers["X-Timeout"], 17)

        with RequestMockWithTimeoutAsHeader() as mock_requests:
            mock_requests.add(responses.GET, "http://example.com/")
            OutgoingSession(role="testing", timeout=17).get("http://example.com/", timeout=42)
            self.assert_length(mock_requests.calls, 1)
            self.assertEqual(mock_requests.calls[0].request.headers["X-Timeout"], 42)

    def test_retries(self) -> None:
        # Responses doesn't support testing the low-level retry
        # functionality, so we can't test the retry itself easily. :(
        # https://github.com/getsentry/responses/issues/135

        # Defaults to no retries
        session = requests.Session()
        assert isinstance(session.adapters["http://"], HTTPAdapter)
        self.assertEqual(session.adapters["http://"].max_retries.total, 0)
        assert isinstance(session.adapters["https://"], HTTPAdapter)
        self.assertEqual(session.adapters["https://"].max_retries.total, 0)

        session = OutgoingSession(role="testing", timeout=1)
        assert isinstance(session.adapters["http://"], HTTPAdapter)
        self.assertEqual(session.adapters["http://"].max_retries.total, 0)
        assert isinstance(session.adapters["https://"], HTTPAdapter)
        self.assertEqual(session.adapters["https://"].max_retries.total, 0)

        session = OutgoingSession(role="testing", timeout=1, max_retries=2)
        assert isinstance(session.adapters["http://"], HTTPAdapter)
        self.assertEqual(session.adapters["http://"].max_retries.total, 2)
        assert isinstance(session.adapters["https://"], HTTPAdapter)
        self.assertEqual(session.adapters["https://"].max_retries.total, 2)

        session = OutgoingSession(role="testing", timeout=1, max_retries=Retry(total=5))
        assert isinstance(session.adapters["http://"], HTTPAdapter)
        self.assertEqual(session.adapters["http://"].max_retries.total, 5)
        assert isinstance(session.adapters["https://"], HTTPAdapter)
        self.assertEqual(session.adapters["https://"].max_retries.total, 5)


class TestBoto3SmokescreenBypass(ZulipTestCase):
    # boto3 is set up at startup to bypass the Smokescreen proxy; see
    # skip_smokescreen_proxy_for_boto3 for why.
    @mock.patch.dict(os.environ, {"http_proxy": "http://smokescreen.example:4750"})
    def test_skip_smokescreen_proxy_for_boto3(self) -> None:
        # Restore botocore's global state however the test exits, so the
        # bypass doesn't leak into other tests in this process.
        original = botocore.utils.should_bypass_proxies
        self.addCleanup(setattr, botocore.utils, "should_bypass_proxies", original)

        imds_url = "http://169.254.169.254/latest/api/token"

        # Without the setup done by skip_smokescreen_proxy_for_boto3, boto3
        # would route IMDS requests through Smokescreen.
        self.assertEqual(
            botocore.utils.get_environ_proxies(imds_url),
            {"http": "http://smokescreen.example:4750"},
        )

        skip_smokescreen_proxy_for_boto3()

        # After the patch, boto3 bypasses Smokescreen and connects directly.
        self.assertEqual(botocore.utils.get_environ_proxies(imds_url), {})