Fix proxy-authentication headers for Python 3.* and long basic credential encodings

impala/_thrift_api.py

@@ -37,6 +37,7 @@
 import sys
 
 from impala.error import HttpError
+from impala.util import get_basic_credentials_for_request_headers
 from impala.util import get_logger_and_init_null
 from impala.util import get_all_matching_cookies, get_all_cookies, get_cookie_expiry
 
@@ -177,10 +178,10 @@ def __init__(self, uri_or_host, port=None, path=None, cafile=None, cert_file=Non
   def basic_proxy_auth_header(proxy):
     if proxy is None or not proxy.username:
       return None
-    ap = "%s:%s" % (urllib.parse.unquote(proxy.username),
-                    urllib.parse.unquote(proxy.password))
-    cr = base64.b64encode(ap).strip()
-    return "Basic " + cr
+    return "Basic " + get_basic_credentials_for_request_headers(
+       user=urllib.parse.unquote(proxy.username),
+       password=urllib.parse.unquote(proxy.password),
+    )
 
   def using_proxy(self):
     return self.realhost is not None
@@ -462,11 +463,7 @@ def get_http_transport(host, port, http_path, timeout=None, use_ssl=False,
             log.debug('get_http_transport: password=fuggetaboutit')
         auth_mechanism = 'PLAIN'  # sasl doesn't know mechanism LDAP
         # Set the BASIC auth header
-        user_password = '%s:%s'.encode() % (user.encode(), password.encode())
-        try:
-            auth = base64.encodebytes(user_password).decode().strip('\n')
-        except AttributeError:
-            auth = base64.encodestring(user_password).decode().strip('\n')
+        auth = get_basic_credentials_for_request_headers(user, password)
 
         def get_custom_headers(cookie_header, has_auth_cookie):
             custom_headers = {}

impala/tests/test_http_connect.py

@@ -188,6 +188,23 @@ def test_duplicate_headers(self, http_proxy_server):
     assert count_tuples_with_key(headers, "key1") == 2
     assert count_tuples_with_key(headers, "key2") == 1
     assert count_tuples_with_key(headers, "key3") == 0
+
+  def test_basic_auth_headers(self, http_proxy_server):
+    con = connect(
+      "localhost",
+      http_proxy_server.PORT,
+      use_http_transport=True,
+      user="thisisaratherlongusername",
+      password="very!long!passwordthatcreatesalongbasic64encoding",
+      auth_mechanism="PLAIN"
+    )
+    cur = con.cursor()
+    cur.execute('select 1')
+    rows = cur.fetchall()
+    assert rows == [(1,)]
+
+    headers = http_proxy_server.get_headers()
+    assert ('Authorization', "Basic dGhpc2lzYXJhdGhlcmxvbmd1c2VybmFtZTp2ZXJ5IWxvbmchcGFzc3dvcmR0aGF0Y3JlYXRlc2Fsb25nYmFzaWM2NGVuY29kaW5n") in headers
 
 def get_user_custom_headers_func():
   """Insert some custom http headers, including a duplicate."""

impala/tests/test_thrift_api.py

@@ -0,0 +1,24 @@
+import os
+
+import pytest
+
+from impala._thrift_api import ImpalaHttpClient
+
+
+@pytest.fixture()
+def proxy_env():
+    reset_value = os.environ.get("HTTPS_PROXY")
+    os.environ["HTTPS_PROXY"] = "https://foo:%3F%40%3D@localhost"
+    yield "proxy_env"
+    if reset_value is None:
+        del os.environ["HTTPS_PROXY"]
+    else:
+        os.environ["HTTPS_PROXY"] = reset_value
+
+
+class TestHttpTransport(object):
+    def test_proxy_auth_header(self, proxy_env):
+        client = ImpalaHttpClient(
+            uri_or_host="https://localhost:443/cliservice",
+        )
+        assert client.proxy_auth == "Basic Zm9vOj9APQ=="

impala/tests/test_util.py

@@ -25,7 +25,7 @@
 
 import pytest
 from impala.util import (cookie_matches_path, get_cookie_expiry, get_all_cookies,
-                         get_all_matching_cookies)
+                         get_all_matching_cookies, get_basic_credentials_for_request_headers)
 
 
 class ImpalaUtilTests(unittest.TestCase):
@@ -211,6 +211,20 @@ def test_get_all_matching_cookies(self):
         assert len(cookies) == 1
         assert cookies[0].key == 'c_cookie' and cookies[0].value == 'c_value'
 
+    def test_get_basic_credentials_for_request_headers(self):
+        assert get_basic_credentials_for_request_headers(
+            user="foo",
+            password="bar"
+        ) == "Zm9vOmJhcg=="
+        assert get_basic_credentials_for_request_headers(
+            user="thisisaratherlongusername",
+            password="withanotherverylongpasswordresultinginanencodinglongerthan76chars"
+        ) == "dGhpc2lzYXJhdGhlcmxvbmd1c2VybmFtZTp3aXRoYW5vdGhlcnZlcnlsb25ncGFzc3dvcmRyZXN1bHRpbmdpbmFuZW5jb2Rpbmdsb25nZXJ0aGFuNzZjaGFycw=="
+        assert get_basic_credentials_for_request_headers(
+            user="?",
+            password="?"
+        ) == "Pzo/"
+
 
 def make_cookie_headers(cookie_vals):
     """Make an HTTPMessage containing Set-Cookie headers for Python 2 or Python 3"""

impala/util.py

@@ -14,6 +14,7 @@
 
 from __future__ import absolute_import
 
+import base64
 import sys
 import warnings
 import logging
@@ -252,3 +253,15 @@ def get_all_matching_cookies(cookie_names, path, resp_headers):
             if c and cookie_matches_path(c, path):
                 matching_cookies.append(c)
     return matching_cookies
+
+
+def get_basic_credentials_for_request_headers(user, password):
+    """Returns base64 encoded credentials for HTTP request headers
+
+    This function produces RFC 2617-compliant basic credentials:
+    - RFC 2045 encoding of username:password without limitations to 76 chars
+      per line (and without trailing newline)
+    - No translation of characters (+,/) for URL-safety
+    """
+    user_password = '%s:%s' % (user, password)
+    return base64.b64encode(user_password.encode()).decode()