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), {})