From f8826b29a24b51ef8fe6b9dbf574477db007592b Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Wed, 25 Aug 2021 17:23:40 +0200
Subject: [PATCH] cache socks5 auth for entire connection

---
 mitmproxy/addons/proxyauth.py           | 8 ++++----
 mitmproxy/proxy/layers/modes.py         | 5 +++--
 test/mitmproxy/addons/test_proxyauth.py | 2 +-
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/mitmproxy/addons/proxyauth.py b/mitmproxy/addons/proxyauth.py
index 8dd6a1786..e80295036 100644
--- a/mitmproxy/addons/proxyauth.py
+++ b/mitmproxy/addons/proxyauth.py
@@ -62,12 +62,12 @@ class ProxyAuth:
     def socks5_auth(self, data: modes.Socks5AuthData) -> None:
         if self.validator and self.validator(data.username, data.password):
             data.valid = True
+            self.authenticated[data.client_conn] = data.username, data.password
 
     def http_connect(self, f: http.HTTPFlow) -> None:
-        if self.validator:
-            if self.authenticate_http(f):
-                # Make a note that all further requests over this connection are ok.
-                self.authenticated[f.client_conn] = f.metadata["proxyauth"]
+        if self.validator and self.authenticate_http(f):
+            # Make a note that all further requests over this connection are ok.
+            self.authenticated[f.client_conn] = f.metadata["proxyauth"]
 
     def requestheaders(self, f: http.HTTPFlow) -> None:
         if self.validator:
diff --git a/mitmproxy/proxy/layers/modes.py b/mitmproxy/proxy/layers/modes.py
index 5c7f4797b..426c2e8c0 100644
--- a/mitmproxy/proxy/layers/modes.py
+++ b/mitmproxy/proxy/layers/modes.py
@@ -4,7 +4,7 @@ from abc import ABCMeta
 from dataclasses import dataclass
 from typing import Optional
 
-from mitmproxy import platform
+from mitmproxy import connection, platform
 from mitmproxy.net import server_spec
 from mitmproxy.proxy import commands, events, layer
 from mitmproxy.proxy.commands import StartHook
@@ -92,6 +92,7 @@ SOCKS5_REP_ADDRESS_TYPE_NOT_SUPPORTED = 0x08
 
 @dataclass
 class Socks5AuthData:
+    client_conn: connection.Client
     username: str
     password: str
     valid: bool = False
@@ -188,7 +189,7 @@ class Socks5Proxy(DestinationKnown):
         user = self.buf[2:(2 + user_len)].decode("utf-8", "backslashreplace")
         password = self.buf[(3 + user_len):(3 + user_len + pass_len)].decode("utf-8", "backslashreplace")
 
-        data = Socks5AuthData(user, password)
+        data = Socks5AuthData(self.context.client, user, password)
         yield Socks5AuthHook(data)
         if not data.valid:
             yield from self.socks_err("authentication failed", 0x01)
diff --git a/test/mitmproxy/addons/test_proxyauth.py b/test/mitmproxy/addons/test_proxyauth.py
index eaa489509..6182cdf99 100644
--- a/test/mitmproxy/addons/test_proxyauth.py
+++ b/test/mitmproxy/addons/test_proxyauth.py
@@ -78,7 +78,7 @@ class TestProxyAuth:
         pa = proxyauth.ProxyAuth()
         with taddons.context(pa, loadcore=False) as ctx:
             ctx.configure(pa, proxyauth="foo:bar", mode="regular")
-            data = modes.Socks5AuthData("foo", "baz")
+            data = modes.Socks5AuthData(tflow.tclient_conn(), "foo", "baz")
             pa.socks5_auth(data)
             assert not data.valid
             data.password = "bar"