Prechádzať zdrojové kódy

[raz][abfs] Improve url params encoding for RAZ ABFS client (#3666)

What changes were proposed in this pull request?

- Earlier we were encoding the url params only for small set of use-cases but it turns out we should always encode the params to cover all use-cases.
- To support the above observation, I'd to override the _make_url method of HttpClient class and change its urlencode method to use quote() method instead of default quote_plus(). This was required to fix scenarios of whitespace present in the path which regressed after the above change.

How was this patch tested?

- Tested E2E in live setup.
- Updating existing unit tests.
- Adding new unit tests.
Harsh Gupta 1 rok pred
rodič
commit
622a0e7b0c

+ 64 - 17
desktop/core/src/desktop/lib/rest/raz_http_client.py

@@ -15,25 +15,36 @@
 # limitations under the License.
 
 import logging
-import sys
+import posixpath
+
+from django.utils.translation import gettext as _
+from django.utils.encoding import iri_to_uri
 
 from desktop import conf
 from desktop.lib.raz.clients import AdlsRazClient
 from desktop.lib.rest.http_client import HttpClient
 from desktop.lib.exceptions_renderable import PopupException
 
-if sys.version_info[0] > 2:
-  from django.utils.translation import gettext as _
-  from urllib.parse import quote as lib_urlquote
-else:
-  from django.utils.translation import ugettext as _
-  from urllib import quote as lib_urlquote
+from urllib.parse import quote, urlencode
 
 
 LOG = logging.getLogger()
 
 
 class RazHttpClient(HttpClient):
+  """
+  A custom HTTP client that adds support for generating Shared Access Signature (SAS) tokens for ABFS.
+
+  This class extends :class:`desktop.lib.rest.http_client.HttpClient`. The main difference is the addition of the `execute` method,
+  which generates a SAS token based on the given parameters and appends it to the URL before making the actual request using the
+  parent class's `execute` method to ABFS for performing the user action.
+
+  Args:
+    username (str): The name of the user associated with the object URL. Used to generate the SAS token.
+    base_url (str): Base URL for the REST API endpoint. Must include protocol (HTTP or HTTPS) and hostname.
+    exc_class (type, optional): An exception type used by this instance when raising errors. Defaults to None.
+    logger (logging.Logger, optional): A logger instance. If not provided, uses the default logger. Defaults to None.
+  """
 
   def __init__(self, username, base_url, exc_class=None, logger=None):
     super(RazHttpClient, self).__init__(base_url, exc_class, logger)
@@ -42,18 +53,21 @@ class RazHttpClient(HttpClient):
   def execute(self, http_method, path, params=None, data=None, headers=None, allow_redirects=False, urlencode=True,
               files=None, stream=False, clear_cookies=False, timeout=conf.REST_CONN_TIMEOUT.get(), retry=1):
     """
-    From an object URL we get back the SAS token as a GET param string, e.g.:
-    https://{storageaccountname}.dfs.core.windows.net/{container}/{path}
-    -->
-    https://{storageaccountname}.dfs.core.windows.net/{container}/{path}?sv=2014-02-14&sr=b&
-    sig=pJL%2FWyed41tptiwBM5ymYre4qF8wzrO05tS5MCjkutc%3D&st=2015-01-02T01%3A40%3A51Z&se=2015-01-02T02%3A00%3A51Z&sp=r
-    """
-    if sys.version_info[0] < 3 and isinstance(path, unicode):
-      path = path.encode('utf-8')
+    Overrides the parent class's `execute` method. Before making the request, generates a SAS token based on the given parameters and
+    appends it to the URL. Then calls the parent class's `execute` method with the modified URL to send the user action request to ABFS.
+
+    Eg: https://{storageaccountname}.dfs.core.windows.net/{container}/{path}
+        -->
+        https://{storageaccountname}.dfs.core.windows.net/{container}/{path}?sv=2014-02-14&sr=b&
+        sig=pJL%2FWyed41tptiwBM5ymYre4qF8wzrO05tS5MCjkutc%3D&st=2015-01-02T01%3A40%3A51Z&se=2015-01-02T02%3A00%3A51Z&sp=r
 
-    do_urlencode = True if params and 'directory' in params and '%' in params['directory'] else False
+    Returns:
+      Any: The result returned by the parent class's `execute` method after performing the user action to ABFS.
 
-    url = self._make_url(lib_urlquote(path), params, do_urlencode=do_urlencode)
+    Raises:
+      PopupException: When no SAS token is present in the RAZ response.
+    """
+    url = self._make_url(quote(path), params)
 
     # For root stats, the root path needs to end with '/' before adding the query params.
     if params and 'action' in params and params['action'] == 'getAccessControl':
@@ -109,6 +123,18 @@ class RazHttpClient(HttpClient):
         raise e
 
   def get_sas_token(self, http_method, username, url, params=None, headers=None):
+    """
+    Request a SAS token from the RAZ service for the specified resource.
+
+    Calls the RAZ client's `get_url` method and returns the received token if available. Otherwise, raises
+    a PopupException indicating a missing token.
+
+    Returns:
+      str: The generated SAS token if successful; otherwise, raises an exception.
+
+    Raises:
+      PopupException: When no SAS token is present in the RAZ response.
+    """
     raz_client = AdlsRazClient(username=username)
     response = raz_client.get_url(action=http_method, path=url, headers=headers)
 
@@ -116,3 +142,24 @@ class RazHttpClient(HttpClient):
       return response.get('token')
     else:
       raise PopupException(_('No SAS token in RAZ response'), error_code=403)
+
+  def _make_url(self, path, params, do_urlencode=True):
+    """
+    Construct a complete URL with the given path and optional query parameters.
+
+    The method overrides parent class's `_make_url` method to change parameter normalization and ensures proper escaping
+    for RAZ by changing the default behaviour of `urlencode` helper method for special characters.
+
+    Returns:
+      str: A fully qualified URL including the scheme, domain, path, and optionally, query parameters.
+    """
+    res = self._base_url
+
+    if path:
+      res += posixpath.normpath('/' + path.lstrip('/'))
+
+    if params:
+      param_str = urlencode(params, safe='/', quote_via=quote) if do_urlencode else '&'.join(['%s=%s' % (k, v) for k, v in params.items()])
+      res += '?' + param_str
+
+    return iri_to_uri(res)

+ 53 - 14
desktop/core/src/desktop/lib/rest/raz_http_client_test.py

@@ -15,20 +15,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import sys
-
-from nose.tools import assert_equal, assert_false, assert_true, assert_raises
+from nose.tools import assert_equal, assert_raises
+from unittest.mock import patch, Mock
 
 from desktop.lib.rest.raz_http_client import RazHttpClient
 from desktop.lib.exceptions_renderable import PopupException
 
 from hadoop.fs.exceptions import WebHdfsException
 
-if sys.version_info[0] > 2:
-  from unittest.mock import patch, Mock
-else:
-  from mock import patch, Mock
-
 
 class TestRazHttpClient():
 
@@ -93,10 +87,7 @@ class TestRazHttpClient():
         client = RazHttpClient(username='test', base_url='https://gethue.dfs.core.windows.net')
 
         # List call for non-ascii directory name (/user/Tжейкоб)
-        if sys.version_info[0] > 2:
-          params = {'directory': 'user/T\u0436\u0435\u0438\u0306\u043a\u043e\u0431', 'resource': 'filesystem'}
-        else:
-          params = {u'directory': u'user/T\u0436\u0435\u0438\u0306\u043a\u043e\u0431', u'resource': u'filesystem'}
+        params = {'directory': 'user/T\u0436\u0435\u0438\u0306\u043a\u043e\u0431', 'resource': 'filesystem'}
 
         f = client.execute(
           http_method='GET',
@@ -141,12 +132,60 @@ class TestRazHttpClient():
 
         # List call for directory name having %20 like characters (/user/ab%20cd)
         f = client.execute(http_method='GET', path='/test', params={'directory': 'user/ab%20cd', 'resource': 'filesystem'})
-        url = 'https://gethue.dfs.core.windows.net/test?directory=user%2Fab%2520cd&resource=filesystem'
+        url = 'https://gethue.dfs.core.windows.net/test?directory=user/ab%2520cd&resource=filesystem'
+
+        raz_get_url.assert_called_with(action='GET', path=url, headers=None)
+        raz_http_execute.assert_called_with(
+            http_method='GET',
+            path='/test?directory=user/ab%2520cd&resource=filesystem&sv=2014-02-14&' \
+              'sr=b&sig=pJL%2FWyed41tptiwBM5ymYre4qF8wzrO05tS5MCjkutc%3D&st=2015-01-02T01%3A40%3A51Z&se=2015-01-02T02%3A00%3A51Z&sp=r',
+            data=None,
+            headers=None,
+            allow_redirects=False,
+            urlencode=False,
+            files=None,
+            stream=False,
+            clear_cookies=False,
+            timeout=120
+        )
+
+        # List call for directory name having objects greater than 5000 and having continuation token param
+        f = client.execute(
+          http_method='GET',
+          path='/test',
+          params={
+            'directory': 'user/test-dir',
+            'resource': 'filesystem',
+            'continuation': 'VBbzu86Hto/ksAkYKRgOZmlsZV8xNDQ5OC5jc3YWhK6wsrzcudoDGAAWiOHZ1/ivtdoDOAAAAA=='
+          }
+        )
+        url = 'https://gethue.dfs.core.windows.net/test?directory=user/test-dir&resource=filesystem&' \
+              'continuation=VBbzu86Hto/ksAkYKRgOZmlsZV8xNDQ5OC5jc3YWhK6wsrzcudoDGAAWiOHZ1/ivtdoDOAAAAA%3D%3D'
+
+        raz_get_url.assert_called_with(action='GET', path=url, headers=None)
+        raz_http_execute.assert_called_with(
+            http_method='GET',
+            path='/test?directory=user/test-dir&resource=filesystem&' \
+              'continuation=VBbzu86Hto/ksAkYKRgOZmlsZV8xNDQ5OC5jc3YWhK6wsrzcudoDGAAWiOHZ1/ivtdoDOAAAAA%3D%3D&sv=2014-02-14&sr=b&' \
+              'sig=pJL%2FWyed41tptiwBM5ymYre4qF8wzrO05tS5MCjkutc%3D&st=2015-01-02T01%3A40%3A51Z&se=2015-01-02T02%3A00%3A51Z&sp=r',
+            data=None,
+            headers=None,
+            allow_redirects=False,
+            urlencode=False,
+            files=None,
+            stream=False,
+            clear_cookies=False,
+            timeout=120
+        )
+
+        # List call for testdir~@$&()*!+=; directory name (/user/testdir~@$&()*!+=;)
+        f = client.execute(http_method='GET', path='/test', params={'directory': 'user/testdir~@$&()*!+=;', 'resource': 'filesystem'})
+        url = 'https://gethue.dfs.core.windows.net/test?directory=user/testdir~%40%24%26%28%29%2A%21%2B%3D%3B&resource=filesystem'
 
         raz_get_url.assert_called_with(action='GET', path=url, headers=None)
         raz_http_execute.assert_called_with(
             http_method='GET',
-            path='/test?directory=user%2Fab%2520cd&resource=filesystem&sv=2014-02-14&' \
+            path='/test?directory=user/testdir~%40%24%26%28%29%2A%21%2B%3D%3B&resource=filesystem&sv=2014-02-14&' \
               'sr=b&sig=pJL%2FWyed41tptiwBM5ymYre4qF8wzrO05tS5MCjkutc%3D&st=2015-01-02T01%3A40%3A51Z&se=2015-01-02T02%3A00%3A51Z&sp=r',
             data=None,
             headers=None,