From 5a183881dd23bdd56a13fae22b19b1545db13d4f Mon Sep 17 00:00:00 2001
From: Will Keeling <william_keeling@mentor.com>
Date: Sat, 5 Nov 2022 15:35:58 +0000
Subject: [PATCH] Take into account SOCKS proxies for no_proxy matches

---
 seleniumwire/thirdparty/mitmproxy/connections.py         | 6 +++++-
 .../thirdparty/mitmproxy/server/protocol/http.py         | 6 ++++--
 tests/end2end/test_end2end.py                            | 9 +++++++++
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/seleniumwire/thirdparty/mitmproxy/connections.py b/seleniumwire/thirdparty/mitmproxy/connections.py
index 00d0d72..f2be0e8 100644
--- a/seleniumwire/thirdparty/mitmproxy/connections.py
+++ b/seleniumwire/thirdparty/mitmproxy/connections.py
@@ -310,13 +310,17 @@ class SocksServerConnection(ServerConnection):
     def __init__(self, socks_config, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.socks_config = socks_config
+        self.use_socks = True
 
     def getaddrinfo(self, host, port, *args, **kwargs):
-        if self.socks_config.scheme == "socks5h":
+        if self.use_socks and self.socks_config.scheme == "socks5h":
             return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", (host, port))]
         return super().getaddrinfo(host, port, *args, **kwargs)
 
     def makesocket(self, family, type, proto):
+        if not self.use_socks:
+            return super().makesocket(family, type, proto)
+
         try:
             socks_type = dict(
                 socks4=socks.PROXY_TYPE_SOCKS4,
diff --git a/seleniumwire/thirdparty/mitmproxy/server/protocol/http.py b/seleniumwire/thirdparty/mitmproxy/server/protocol/http.py
index 7a5a391..7e8f1d9 100644
--- a/seleniumwire/thirdparty/mitmproxy/server/protocol/http.py
+++ b/seleniumwire/thirdparty/mitmproxy/server/protocol/http.py
@@ -267,9 +267,11 @@ def _process_flow(self, f):
 
             f.request = request
 
-            if self.mode is HTTPMode.upstream and self.should_bypass_upstream_proxy(f.request):
+            if (self.mode is HTTPMode.upstream or "socks" in self.config.options.mode) and \
+                    self.matches_no_proxy(f.request):
                 self.set_server((f.request.host, f.request.port))
                 self.mode = HTTPMode.regular
+                self.server_conn.use_socks = False
 
             if request.first_line_format == "authority":
                 # The standards are silent on what we should do with a CONNECT
@@ -538,7 +540,7 @@ def establish_server_connection(self, host: str, port: int, scheme: str):
             if tls:
                 raise exceptions.HttpProtocolException("Cannot change scheme in upstream mitmproxy mode.")
 
-    def should_bypass_upstream_proxy(self, request):
+    def matches_no_proxy(self, request):
         """Whether we should bypass any upstream proxy.
 
         This checks whether the request address is in the no_proxy list.
diff --git a/tests/end2end/test_end2end.py b/tests/end2end/test_end2end.py
index 2b175e7..1e95544 100644
--- a/tests/end2end/test_end2end.py
+++ b/tests/end2end/test_end2end.py
@@ -352,6 +352,15 @@ def test_upstream_socks_proxy(driver_path, chrome_options, httpbin, socksproxy):
         assert 'This passed through a socks proxy' in driver.page_source
 
 
+def test_bypass_upstream_socks_proxy(driver_path, chrome_options, httpbin, socksproxy):
+    sw_options = {'proxy': {'https': f'{socksproxy}', 'no_proxy': 'localhost:8085'}}
+
+    with create_driver(driver_path, chrome_options, sw_options) as driver:
+        driver.get(f'{httpbin}/html')
+
+        assert 'This passed through a socks proxy' not in driver.page_source
+
+
 def test_bypass_upstream_proxy_when_target_http(driver_path, chrome_options, httpproxy):
     sw_options = {'proxy': {'https': f'{httpproxy}', 'no_proxy': 'localhost:9091'}}